[GH-ISSUE #414] [Bug] `pack_ue8m0_to_int` is not CUDA Graph capture safe #189

New issue

Open

opened 2026-10-26 21:14:27 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:14:27 +00:00

Owner

Originally created by @foraxe on GitHub (Aug 24, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/414

Summary

deep_gemm.utils.math.pack_ue8m0_to_int performs CUDA-tensor value checks with
Python assertions:

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:

(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.

Originally created by @foraxe on GitHub (Aug 24, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/414 ## 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`.

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepGEMM#189

No description provided.

Rows
Columns

[GH-ISSUE #414] [Bug] pack_ue8m0_to_int is not CUDA Graph capture safe #189