deepseek-ai / deepseek-ai/DeepGEMM
[Bug] SM120 m_grouped_fp8 contiguous kernel: illegal memory access at CUDA graph replay (eager + capture + identical-replay all clean)
- Dominant language
- Cuda
- Stars
- 7.8k
- Forks
- 1.3k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 3
Description
## Description
The SM120 (consumer/workstation Blackwell) `m_grouped_fp8_fp4_gemm_nt_contiguous` kernel runs cleanly in eager mode across our full shape matrix, completes CUDA-graph **capture**, passes a full end-to-end quality probe — and then hits an **illegal memory access at graph replay** (`torch.cuda.CUDAGraph.replay()` → `CUDA_ERROR_ILLEGAL_ADDRESS`).
We believe this is the grouped-GEMM instance of the same family as #414 (`pack_ue8m0_to_int` not capture safe) and #431/#430 (missing tensormap drain): the SM120 1D1D kernels update tensormaps on-device (`ptx::tensor_map_replace_*` + publishing the updated tensormap to a GMEM buffer) inside the stream. Under graph capture the tensormap state from the first invocation is recorded; on replay a new batch has a different per-expert token distribution, so TMA consumes stale descriptors and reads out of bounds.
## Environment
- DeepGEMM: **0.1.5.post3** (wheel, nv_dev line with SM120 grouped support)
- GPU: NVIDIA RTX PRO 5000 72GB Blackwell (capability 12.0), single node TP8
- torch 2.13.0+cu130, CUDA 13.0
- Workload: GLM-5.3-Flash blockwise-FP8 MoE via sglang — G=288 experts, per-rank w13 N=512/K=4096 and w2 N=4096/K=256, ue8m0-packed scales on both operands (weights requantized to 2^k scales, activations quantized with column-major TMA-aligned packed scales)
## What works (eager, everything)
- The official `test_m_grouped_gemm_contiguous` shapes: 32/32 pass on this GPU
- Our exact production shapes (G=288, w13/w2, decode `expected_m_per_group` 1/32/113): all pass eager, outputs clean
- The same official-generator inputs inside `torch.cuda.graph` capture **and a single replay with identical inputs/shape**: passes (stale-descriptor hazard not exercised when nothing changes between capture and replay)
## What fails
1. sglang serve with the deep_gemm MoE runner: decode CUDA-graph **capture completes** (20 batch-size tiers), a 30-request prefill/generate quality probe **passes 30/30** (accept_len 2.99, matching the triton baseline) — then the first real decode step replaying the captured graph dies with `illegal memory access` at `full_cuda_graph_backend.py replay()`. With `CUDA_LAUNCH_BLOCKING=1` the report is stable.
2. Offline: any eager call sequence that changes `m_indices` / per-expert distribution between iterations while the same buffers are reused is the suspected trigger; we reproduced the crash deterministically inside the server and captured the exact kernel arguments via a one-shot dump right before the failure (available in the repro kit below).
## Evidence / repro kit
- Kernel-argument dump → offline replay reproduces the scale-format assertion (fixed on the sglang side, see sgl-project/sglang#39063 comment) and isolates this replay crash
- Gate-pass / replay-crash split: eager path fully functional, graph replay of a *changed* batch fails
- Scripts: argument dumper (monkeypatched dispatch entry), offline replayer, official-generator A/B harness, graph-capture harness — happy to attach or open a PR with a regression test
## Suspicion (for maintainers)
`sm120_fp8_fp4_gemm_1d1d.cuh` publishes updated tensormaps to GMEM per group switch; under graph capture these side effects are frozen. A capture-safe design likely needs either (a) persistent tensormap buffers whose contents are re-derived on replay from the (graph-external) `m_indices`/count tensors, or (b) falling back to static descriptors when captured. Related: #414 (pack kernel capture safety), #431 (tensormap drain port), and sgl-project/sglang#25551 (B300 graph-capture crash in the same integration path).
Happy to test any candidate fix on this hardware on short notice.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with sm120_fp8_fp4_gemm_1d1d.cuh and the graph-capture harness or offline replayer from the repro kit. Reproduce with changed m_indices or per-expert distributions between capture and replay, then trace the tensormap publication path. Done means the grouped kernel replays safely with changed batches without illegal memory access.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, pytorch
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100