[PR #421] fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414) #433

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/421
Author: @XFDG
Created: 8/26/2026
Status: 🔄 Open

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

📝 Commits (5)

e33de5a fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414)
7798589 fix: add capture-safe value validation to pack_ue8m0_to_int
32611f0 fix: use bit-level checks for ue8m0 validation (avoid float rounding edge case)
6f5df09 fix: add docstring and fix inaccurate comment per ds-review-bot suggestions
8e207fc fix: tighten UE8M0 capture guard and tests

📊 Changes

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

View changed files

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

📄 Description

Problem

pack_ue8m0_to_int called .all() on device tensors:

assert (x.view(torch.int) >> 23 >= 0).all()  # device→host sync!
assert (x.view(torch.int) & 0x7FFFFF == 0).all()  # device→host sync!

Inside a CUDA graph capture region, .all() triggers a device→host sync which raises cudaErrorStreamCaptureUnsupported, breaking SGLang's decode graph via the MegaMoE FP8 staging path (issue #414).

Fix

Use torch.cuda.is_current_stream_capturing() to gate validation:

Outside capture: bit-level assertions run, catching non-negative sign bit and zero mantissa — direct callers get a loud error on bad input
Inside capture: assertions skipped — no device→host sync, capture-safe

Bit-level check avoids float rounding edge cases (e.g. subnormals like 2^-126 that x.round() would incorrectly round to 0).

if not torch.cuda.is_current_stream_capturing():
    x_bits = x.view(torch.int)
    assert ((x_bits >> 31) == 0).all(), "scale values must be non-negative"
    assert ((x_bits & 0x7FFFFF) == 0).all(), "scale values must have zero mantissa"

Test (H200, CUDA 12.8)

Case	Result
Valid ue8m0 from `ceil_to_ue8m0`	✅ PASS
Non-ue8m0 input (mantissa set)	✅ assert fires
Negative input	✅ assert fires
Padding value 1.0 (valid ue8m0)	✅ PASS
CUDA graph capture + invalid input	✅ no assert, no sync
CUDA graph capture + valid input	✅ replay works

Addresses ds-review-bot v6 concern: validation is preserved outside capture paths.

_{🔄 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/421 **Author:** [@XFDG](https://github.com/XFDG) **Created:** 8/26/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-issue-414-capture-safe` --- ### 📝 Commits (5) - [`e33de5a`](https://github.com/deepseek-ai/DeepGEMM/commit/e33de5adf7dec84b2278ceb048f637e0aa3a6287) fix: make pack_ue8m0_to_int CUDA-graph capture safe (#414) - [`7798589`](https://github.com/deepseek-ai/DeepGEMM/commit/7798589fa339fca29ba08deb2c183a8e4a8ff11d) fix: add capture-safe value validation to pack_ue8m0_to_int - [`32611f0`](https://github.com/deepseek-ai/DeepGEMM/commit/32611f08212630af4148f357c79580163c05be06) fix: use bit-level checks for ue8m0 validation (avoid float rounding edge case) - [`6f5df09`](https://github.com/deepseek-ai/DeepGEMM/commit/6f5df0941b1542c7315b9124a1ada74efe88b27b) fix: add docstring and fix inaccurate comment per ds-review-bot suggestions - [`8e207fc`](https://github.com/deepseek-ai/DeepGEMM/commit/8e207fca8ce58c3aa5897d703d193502a8c96ac1) fix: tighten UE8M0 capture guard and tests ### 📊 Changes **2 files changed** (+75 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/utils/math.py` (+18 -1) ➕ `tests/test_pack_ue8m0.py` (+57 -0) </details> ### 📄 Description ## Problem `pack_ue8m0_to_int` called `.all()` on device tensors: ```python assert (x.view(torch.int) >> 23 >= 0).all() # device→host sync! assert (x.view(torch.int) & 0x7FFFFF == 0).all() # device→host sync! ``` Inside a CUDA graph capture region, `.all()` triggers a device→host sync which raises `cudaErrorStreamCaptureUnsupported`, breaking SGLang's decode graph via the MegaMoE FP8 staging path (issue #414). ## Fix Use `torch.cuda.is_current_stream_capturing()` to gate validation: - **Outside capture**: bit-level assertions run, catching non-negative sign bit and zero mantissa — direct callers get a loud error on bad input - **Inside capture**: assertions skipped — no device→host sync, capture-safe Bit-level check avoids float rounding edge cases (e.g. subnormals like 2^-126 that `x.round()` would incorrectly round to 0). ```python if not torch.cuda.is_current_stream_capturing(): x_bits = x.view(torch.int) assert ((x_bits >> 31) == 0).all(), "scale values must be non-negative" assert ((x_bits & 0x7FFFFF) == 0).all(), "scale values must have zero mantissa" ``` ## Test (H200, CUDA 12.8) | Case | Result | |------|--------| | Valid ue8m0 from `ceil_to_ue8m0` | ✅ PASS | | Non-ue8m0 input (mantissa set) | ✅ assert fires | | Negative input | ✅ assert fires | | Padding value 1.0 (valid ue8m0) | ✅ PASS | | CUDA graph capture + invalid input | ✅ no assert, no sync | | CUDA graph capture + valid input | ✅ replay works | Addresses ds-review-bot v6 concern: validation is preserved outside capture paths. --- <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:17:44 +00:00

bsakaguchi referenced this issue

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

[PR #434] fix: derive the SM120 fp8 mqa_logits swizzle mode from head_dim #445