deepseek-ai / deepseek-ai/DeepGEMM

[Bug] `pack_ue8m0_to_int` is not CUDA Graph capture safe

Open
#414 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.