[PR #420] [CLOSED] fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414) #432

New issue

Closed

opened 2026-10-26 21:17:42 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:42 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/420
Author: @XFDG
Created: 8/26/2026
Status: ❌ Closed

Base: main ← Head: fix-issue-414-capture-safe

📝 Commits (1)

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

📊 Changes

1 file changed (+11 additions, -1 deletions)

View changed files

📝 deep_gemm/utils/math.py (+11 -1)

📄 Description

Fix DeepGEMM #414: pack_ue8m0_to_int CUDA Graph capture safe

PR URL

https://github.com/XFDG/DeepGEMM/pull/new/fix-issue-414-capture-safe

Branch

fix-issue-414-capture-safe on XFDG/DeepGEMM (fork of deepseek-ai/DeepGEMM)

PR Title

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

PR Body

Summary

pack_ue8m0_to_int calls .all() on device tensors to validate scale-factor
invariants at runtime. Each .all() forces a device→host synchronisation,
which is illegal inside a CUDA Graph capture region and raises
cudaErrorStreamCaptureUnsupported.

Downstream projects (e.g. SGLang) already hit this when capturing their decode
graphs that call the MegaMoE FP8 staging path, and currently work around it by
monkey-patching the helper. This fix removes the need for that workaround.

Root cause

# deep_gemm/utils/math.py  (before)
def pack_ue8m0_to_int(x: torch.Tensor):
    assert x.dtype == torch.float and x.size(-1) % 4 == 0
    x_int = x.view(torch.int)
    assert (x_int >= 0).all() and (x_int & 0x7fffff == 0).all()  # ← device sync
    return (x_int >> 23).to(torch.uint8).view(torch.int)

The two .all() calls are redundant invariants guaranteed by
ceil_to_ue8m0, the only function that produces these scale factors:

def ceil_to_ue8m0(x: torch.Tensor):
    bits = x.abs().float().view(torch.int)
    exp = ((bits >> 23) & 0xFF) + (bits & 0x7FFFFF).bool().int()
    return (exp.clamp(1, 254) << 23).view(torch.float)
    #             ↑ always zero mantissa  ↑ always positive exponent [1,254]

Change

Remove the two runtime .all() assertions; document the invariant in a
NOTE comment so callers know it is the producer's responsibility.

# deep_gemm/utils/math.py  (after)
def pack_ue8m0_to_int(x: torch.Tensor):
    # NOTE: keep this helper CUDA-graph-capture safe. ...
    assert x.dtype == torch.float and x.size(-1) % 4 == 0
    x_int = x.view(torch.int)
    # The value invariants are guaranteed by ceil_to_ue8m0; do NOT
    # re-check them here on device.
    return (x_int >> 23).to(torch.uint8).view(torch.int)

Testing

CPU host-only: confirmed ceil_to_ue8m0 output always satisfies both
removed invariants (nonneg=True, zero-mantissa=True) across random inputs.
AST verification: no .all() call remains in pack_ue8m0_to_int.
py_compile syntax check passes.
Functional equivalence: pack body unchanged; output is still the exponent byte.

Compatibility

No API change.
No new dependencies.
For valid callers (all current per_token_cast_to_fp8 callers) behaviour is
identical.
Malformed inputs are still caught by the kernel itself at runtime.

deepseek-ai/DeepGEMM#414

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/420 **Author:** [@XFDG](https://github.com/XFDG) **Created:** 8/26/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix-issue-414-capture-safe` --- ### 📝 Commits (1) - [`d1945e0`](https://github.com/deepseek-ai/DeepGEMM/commit/d1945e07faa427c50eb16348338c9bf5763ba898) fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414) ### 📊 Changes **1 file changed** (+11 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/utils/math.py` (+11 -1) </details> ### 📄 Description # Fix DeepGEMM #414: pack_ue8m0_to_int CUDA Graph capture safe ## PR URL https://github.com/XFDG/DeepGEMM/pull/new/fix-issue-414-capture-safe ## Branch `fix-issue-414-capture-safe` on `XFDG/DeepGEMM` (fork of `deepseek-ai/DeepGEMM`) --- ## PR Title `fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414)` ## PR Body ### Summary `pack_ue8m0_to_int` calls `.all()` on device tensors to validate scale-factor invariants at runtime. Each `.all()` forces a device→host synchronisation, which is illegal inside a CUDA Graph capture region and raises `cudaErrorStreamCaptureUnsupported`. Downstream projects (e.g. SGLang) already hit this when capturing their decode graphs that call the MegaMoE FP8 staging path, and currently work around it by monkey-patching the helper. This fix removes the need for that workaround. ### Root cause ```python # deep_gemm/utils/math.py (before) def pack_ue8m0_to_int(x: torch.Tensor): assert x.dtype == torch.float and x.size(-1) % 4 == 0 x_int = x.view(torch.int) assert (x_int >= 0).all() and (x_int & 0x7fffff == 0).all() # ← device sync return (x_int >> 23).to(torch.uint8).view(torch.int) ``` The two `.all()` calls are **redundant invariants** guaranteed by `ceil_to_ue8m0`, the only function that produces these scale factors: ```python def ceil_to_ue8m0(x: torch.Tensor): bits = x.abs().float().view(torch.int) exp = ((bits >> 23) & 0xFF) + (bits & 0x7FFFFF).bool().int() return (exp.clamp(1, 254) << 23).view(torch.float) # ↑ always zero mantissa ↑ always positive exponent [1,254] ``` ### Change Remove the two runtime `.all()` assertions; document the invariant in a `NOTE` comment so callers know it is the producer's responsibility. ```python # deep_gemm/utils/math.py (after) def pack_ue8m0_to_int(x: torch.Tensor): # NOTE: keep this helper CUDA-graph-capture safe. ... assert x.dtype == torch.float and x.size(-1) % 4 == 0 x_int = x.view(torch.int) # The value invariants are guaranteed by ceil_to_ue8m0; do NOT # re-check them here on device. return (x_int >> 23).to(torch.uint8).view(torch.int) ``` ### Testing - **CPU host-only**: confirmed `ceil_to_ue8m0` output always satisfies both removed invariants (`nonneg=True`, `zero-mantissa=True`) across random inputs. - AST verification: no `.all()` call remains in `pack_ue8m0_to_int`. - `py_compile` syntax check passes. - Functional equivalence: pack body unchanged; output is still the exponent byte. ### Compatibility - No API change. - No new dependencies. - For valid callers (all current `per_token_cast_to_fp8` callers) behaviour is identical. - Malformed inputs are still caught by the kernel itself at runtime. --- ## Related issue deepseek-ai/DeepGEMM#414 --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:17:42 +00:00