[PR #466] Make `pack_ue8m0_to_int` CUDA graph capture safe #470

New issue

Open

opened 2026-10-26 21:18:15 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:15 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/466
Author: @zhoubiao188
Created: 10/6/2026
Status: 🔄 Open

Base: main ← Head: fix/pack-ue8m0-capture-safe

📝 Commits (1)

063e257 Make pack_ue8m0_to_int CUDA graph capture safe

📊 Changes

2 files changed (+90 additions, -1 deletions)

View changed files

📝 deep_gemm/utils/math.py (+5 -1)
➕ tests/test_pack_ue8m0.py (+85 -0)

📄 Description

Summary

deep_gemm.utils.math.pack_ue8m0_to_int validated its input with Python asserts over CUDA tensors. Evaluating those 0-dim tensors as bools forces a device-to-host sync, which fails with cudaErrorStreamCaptureUnsupported when the helper runs while a CUDA graph is being captured.
This blocks downstream callers that stage inputs during capture, e.g. FlashInfer MegaMoE block-FP8 staging under SGLang, which currently works around it by monkey-patching the helper around captured forwards (see the linked issue and the SGLang PR referenced there).
Skip the value checks while the current stream is capturing a CUDA graph. The packing itself ((x.view(torch.int) >> 23).to(torch.uint8).view(torch.int)) is capture safe, and eager validation is preserved everywhere else.
The check is guarded by x.is_cuda, so CPU-only torch builds never evaluate torch.cuda.is_current_stream_capturing() (which raises on builds without CUDA) and CPU tensors keep full validation unconditionally.

Fixes #414

Validation

Added tests/test_pack_ue8m0.py:
- packing values (hand-computed byte layout) and shape round-trips via unpack_ue8m0_from_int;
- eager validation still rejects negative sign bits, non-zero mantissas, and bad shapes/dtypes;
- on CUDA: warmup on a side stream, torch.cuda.graph capture, replay, and comparison against the eager packing — this capture raises on main before the fix;
- eager CUDA validation preserved outside capture.
CPU-only local checks (run outside this repo's GPU suite): the fixed implementation is bitwise identical to the previous one on valid inputs across shapes; CPU tensors never evaluate the capture check; per_token_cast_to_fp8/per_token_cast_to_fp4 with use_packed_ue8m0=True still round-trip through cast_back_from_*.

_{🔄 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/466 **Author:** [@zhoubiao188](https://github.com/zhoubiao188) **Created:** 10/6/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/pack-ue8m0-capture-safe` --- ### 📝 Commits (1) - [`063e257`](https://github.com/deepseek-ai/DeepGEMM/commit/063e2577798393ce0ae705467278bcc68d889f12) Make `pack_ue8m0_to_int` CUDA graph capture safe ### 📊 Changes **2 files changed** (+90 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/utils/math.py` (+5 -1) ➕ `tests/test_pack_ue8m0.py` (+85 -0) </details> ### 📄 Description ## Summary - `deep_gemm.utils.math.pack_ue8m0_to_int` validated its input with Python asserts over CUDA tensors. Evaluating those 0-dim tensors as bools forces a device-to-host sync, which fails with `cudaErrorStreamCaptureUnsupported` when the helper runs while a CUDA graph is being captured. - This blocks downstream callers that stage inputs during capture, e.g. FlashInfer MegaMoE block-FP8 staging under SGLang, which currently works around it by monkey-patching the helper around captured forwards (see the linked issue and the SGLang PR referenced there). - Skip the value checks while the current stream is capturing a CUDA graph. The packing itself (`(x.view(torch.int) >> 23).to(torch.uint8).view(torch.int)`) is capture safe, and eager validation is preserved everywhere else. - The check is guarded by `x.is_cuda`, so CPU-only torch builds never evaluate `torch.cuda.is_current_stream_capturing()` (which raises on builds without CUDA) and CPU tensors keep full validation unconditionally. Fixes #414 ## Validation - Added `tests/test_pack_ue8m0.py`: - packing values (hand-computed byte layout) and shape round-trips via `unpack_ue8m0_from_int`; - eager validation still rejects negative sign bits, non-zero mantissas, and bad shapes/dtypes; - on CUDA: warmup on a side stream, `torch.cuda.graph` capture, replay, and comparison against the eager packing — this capture raises on `main` before the fix; - eager CUDA validation preserved outside capture. - CPU-only local checks (run outside this repo's GPU suite): the fixed implementation is bitwise identical to the previous one on valid inputs across shapes; CPU tensors never evaluate the capture check; `per_token_cast_to_fp8`/`per_token_cast_to_fp4` with `use_packed_ue8m0=True` still round-trip through `cast_back_from_*`. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:18:15 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:19 +00:00

[PR #470] [CLOSED] Sync/upstream dev e1f418c2 #477