fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414)

[XFDG]

Problem

pack_ue8m0_to_int called .all() on device tensors:

assert (x.view(torch.int) >> 23 >= 0).all()  # device→host sync!
assert (x.view(torch.int) & 0x7FFFFF == 0).all()  # device→host sync!

Inside a CUDA graph capture region, .all() triggers a device→host sync which raises cudaErrorStreamCaptureUnsupported, breaking SGLang's decode graph via the MegaMoE FP8 staging path (issue #414).

Fix

Use torch.cuda.is_current_stream_capturing() to gate validation:

• Outside capture: bit-level assertions run, catching non-negative sign bit and zero mantissa — direct callers get a loud error on bad input
• Inside capture: assertions skipped — no device→host sync, capture-safe

Bit-level check avoids float rounding edge cases (e.g. subnormals like 2^-126 that x.round() would incorrectly round to 0).

if not torch.cuda.is_current_stream_capturing():
    x_bits = x.view(torch.int)
    assert ((x_bits >> 31) == 0).all(), "scale values must be non-negative"
    assert ((x_bits & 0x7FFFFF) == 0).all(), "scale values must have zero mantissa"

Test (H200, CUDA 12.8)

Case	Result
Valid ue8m0 from ceil_to_ue8m0	✅ PASS
Non-ue8m0 input (mantissa set)	✅ assert fires
Negative input	✅ assert fires
Padding value 1.0 (valid ue8m0)	✅ PASS
CUDA graph capture + invalid input	✅ no assert, no sync
CUDA graph capture + valid input	✅ replay works

Addresses ds-review-bot v6 concern: validation is preserved outside capture paths.

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

该改动解决了 CUDA Graph 捕获同步问题，但不必要地在所有执行路径移除了既有输入校验，导致非法输入被静默错误编码。

v5

该 MR 移除了 deep_gemm/utils/math.py 中 pack_ue8m0_to_int 内对设备张量的两个 .all() 运行时断言，使该函数在 CUDA Graph 捕获区域内可安全调用（.all() 会触发 device->host 同步，捕获时抛出 cudaErrorStreamCaptureUnsupported）。审查确认：1) 函数体内已无任何 .all() 调用，仅保留主机侧结构性断言（dtype、size(-1) % 4 == 0），捕获安全；2) 被移除的断言属于冗余不变量，由唯一的 scale 生产者 ceil_to_ue8m0 保证（exp.clamp(1,254) << 23 恒产生正指数与零尾数），NOTE 注释已明确责任归属并说明畸形输入仍会在内核层报错；3) 核心打包逻辑 (x_int >> 23).to(torch.uint8).view(torch.int) 未变，功能等价；4) 现有调用方（per_token_cast_to_fp8 与 FP4 路径）在 scale 填充时使用 1.0（零尾数的 2^0），完全满足不变量，行为兼容。无 API 变更、无新依赖，变更小且聚焦，评审通过，建议合并。

v4p

本 MR 移除 pack_ue8m0_to_int 中对设备张量的两个 .all() 校验，以消除 CUDA Graph capture 期间的 device→host 同步，使该 helper 可被捕获。整体实现符合目标：仓库内调用路径（per_token_cast_to_fp8/per_token_cast_to_fp4）的 scale 均由 ceil_to_ue8m0 或 1.0 padding 产生，打包逻辑本身保持不变；仅有一处注释对错误检测路径的描述不准确。

Files reviewed: 1
仓库不存在或无权限: #414
Issues found: 🟡 1 warning | 🔵 2 suggestion
Inline comments posted: 3

[XFDG]

Thanks for the detailed review! Here's what was addressed:

v6 warning (silent corruption on invalid input outside capture):
Added conditional bit-level validation using torch.cuda.is_current_stream_capturing():

• Outside capture: asserts fire on bad input (sign bit ≠ 0 or mantissa bits ≠ 0)
• Inside capture: assertions skipped to avoid cudaErrorStreamCaptureUnsupported
• Used bit-level checks (x_bits >> 31, x_bits & 0x7FFFFF) instead of float-based comparison to avoid false failures on subnormals like 2^-126

v4p suggestion (inaccurate comment "kernel itself traps on malformed scales"):
Removed — downstream kernels (sm100_mqa_logits, sm100_fp8_fp4_gemm_1d1d) only validate shape/dtype, not numerical content.

v5 suggestion (document preconditions):
Added a docstring explicitly stating the UE8M0 precondition, why violations aren't caught downstream, and the capture-safe rationale.

Tested on H200 (CUDA 12.8):

• Valid ceil_to_ue8m0 output: ✅ passes
• Non-UE8M0 input (mantissa set): ✅ asserts outside capture
• Negative input: ✅ asserts outside capture
• CUDA graph capture + invalid input: ✅ no assert, no device→host sync
• CUDA graph capture + valid input: ✅ capture and replay work

[200lz]

I independently reproduced #414 while investigating the SGLang MegaMoE workaround. Your current capture-aware guard fixes the failure for valid inputs — verified on an RTX 3060 with PyTorch 2.6.0 + CUDA 11.8 (eager warm-up, then torch.cuda.graph capture/replay).

I found two small follow-ups that may be useful:

1. Guard the capture query with x.is_cuda, so CPU inputs keep eager value validation without touching CUDA capture state (this also lets the existing x_int bit view be reused).
2. Regression coverage for CUDA graph capture/replay, including mutating the static input buffer before replay to verify that packing is actually recorded and recomputed by the graph. I also kept an eager malformed-input test (nonzero mantissa, negative sign) to ensure the existing validation behavior outside capture is preserved.

I prepared the changes as a single commit on top of your current head (ba14b91):
200lz@763b73d

Local results:

valid eager == capture/replay (bit-exact): PASS
mutated static input replay: PASS
eager malformed-input rejection: PASS
CPU short-circuit (CUDA_VISIBLE_DEVICES=""): PASS
git diff --check: PASS

Happy for you to cherry-pick the commit, or I can open a PR against your branch if that is easier.

[XFDG]

Thanks a lot @200lz — cherry-picked your commit 763b73df as 2a5989a (authorship preserved). Both changes are clear wins:

• The x.is_cuda and ... short-circuit means CPU inputs always run the value validation instead of needlessly querying capture state, so malformed CPU input still errors loudly.
• The new tests/test_pack_ue8m0.py is exactly the capture/replay coverage the review asked for. The mutate-static-buffer-then-replay check is a nice touch — it proves the pack is recorded and recomputed by the graph rather than constant-folded at capture time.

Verified on an H200 (torch 2.10 / CUDA): all four checks pass against the real math.py, and a negative control confirms the pre-fix (unguarded) body still fails capture with cudaErrorStreamCaptureUnsupported — so the test genuinely guards the fix.

Testing CUDA graph capture/replay:
 > Capture/replay matches eager (bit-exact)
 > Replay after static-input mutation matches eager (bit-exact)
Testing eager malformed-input rejection:
 > AssertionError raised as expected (nonzero mantissa)
 > AssertionError raised as expected (negative sign)

Appreciate the independent repro and the follow-up.

[XFDG]

Closed the three remaining review threads after rechecking the current head 2a5989a against each item: the stale warning is fixed by eager bit-level validation outside capture, the precondition is documented, the inaccurate downstream-kernel comment is removed, and the CUDA graph capture/replay plus malformed-input regressions are present. No additional code change was needed. Please take another look when convenient.