deepseek-ai / deepseek-ai/DeepGEMM
[Bug] `pack_ue8m0_to_int` is not CUDA Graph capture safe
- Dominant language
- Cuda
- Stars
- 7.8k
- Forks
- 1.3k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 3
Description
## Summary
`deep_gemm.utils.math.pack_ue8m0_to_int` performs CUDA-tensor value checks with
Python assertions:
```python
x_int = x.view(torch.int)
assert (x_int >= 0).all() and (x_int & 0x7FFFFF == 0).all()
```
Evaluating these zero-dimensional CUDA tensors as Python booleans forces a
device-to-host synchronization. When the helper runs during CUDA Graph capture,
that synchronization fails with `cudaErrorStreamCaptureUnsupported`.
## Affected integration
FlashInfer MegaMoE block-FP8 input staging calls this helper while SGLang
captures its decode graph. SGLang currently needs a narrowly scoped workaround
that replaces the helper only around the captured MegaMoE forward and restores
the original function immediately afterward:
- https://github.com/wenscarl/sglang/pull/6
## Expected behavior
DeepGEMM should expose a capture-safe UE8M0 packing path so downstream projects
do not need to replace a module-level helper. The dtype and shape checks can
remain host-side; the CUDA value checks could be optional, moved to an eager
validation path, or implemented without a host synchronization during capture.
The packing operation itself is capture safe:
```python
(x.view(torch.int) >> 23).to(torch.uint8).view(torch.int)
```
## Suggested validation
- Compare eager and captured/replayed packed output.
- Cover valid positive, zero-mantissa UE8M0 inputs.
- Preserve eager validation for malformed values if that remains part of the
public contract.
## Version
Observed against DeepGEMM `main` at
`559d79fb6994a58b8a15b4b93bf13ccc16edf247`.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start at deep_gemm.utils.math.pack_ue8m0_to_int and inspect its current CUDA-tensor assertions. Validate eager and captured/replayed outputs for valid positive, zero-mantissa UE8M0 inputs, while checking that malformed-value validation remains preserved.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- hpc, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100