[DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0 by DarkSharpness · Pull Request #30378 · sgl-project/sglang

DarkSharpness · 2026-07-07T09:31:08Z

(generated by claude)

Follow-up to #30274: removes its TEMP allow_topk_v2 gate and fixes the root cause of the GLM 5.2 MTP illegal memory access (failing rerun).

Root cause

Under DP attention, the idle-companion / DP-padded rows of a draft-extend-v2 CUDA graph replay carry the graph's seq_len fill value (1), which is smaller than qo_len (= speculative_num_draft_tokens = 6). The per-row visible kv lengths are expanded as kv - qo + 1 + i, so those rows produced negative lengths ([-4, -3, -2, -1, 0, 1] — confirmed both by cuda-gdb register analysis of the CI coredump and by instrumented local replay). The top-k v2 kernel reads lengths as uint32_t, so -4 reinterprets as ~4.29e9 tokens, sending the row down the cluster path over the wide (max-context) graph score buffer; the selection pass then under-fills the output slots and the page-table transform gathers page_table[garbage >> page_bits] → CUDBG_EXCEPTION_WARP_ILLEGAL_ADDRESS in topk_small_batch_kernel, grid (6,8,1).

This also explains the observed failure pattern:

The verify-side expansion (kv + off + 1) is additive and stays positive → only draft-extend companions crash.
TP-only MTP has no DP idle companions (only occasional bucket-padding rows) → passes / flaky, while DP MTP crashes on the first request whose companions go idle.
The kernel is "robust in isolation" — valid non-negative lengths never hit this.

Changes

Clamp expanded seq_lens to >= 0 in fused_dsa_draft_extend_metadata (triton_ops/dsa_metadata.py) and seqlens_expand_kernel (triton_ops/pad.py). A 0-length row takes the kernel's trivial all--1 output path, matching what the legacy/unfused paths give padded rows.
Remove the TEMP allow_topk_v2 gate from [DSA] Fold page-table into fused top-k v2 (decode): drop page_size=1 expansion #30274 — MTP target-verify / draft-extend dispatch to top-k v2 again (closes the TODO(dsa-topk-v2)).
Document the non-negative seq_lens contract on the v2 Python interface (plan_topk_v2, topk_transform_512_v2, _topk_transform_v2_paged).

Validation (local 4x B200, exact CI DP config)

Check	Before (gate removed, no clamp)	After (this PR)
`--tp 4 --dp 4 --enable-dp-attention` + EAGLE MTP, single request	illegal memory access within ~1s; 3 idle ranks core-dump in `topk_small_batch_kernel<true>` (byte-for-byte the CI signature)	OK
GSM8K 500q / 500 parallel	–	accuracy 0.936 (thres 0.935)
`avg_spec_accept_length`	–	4.09–4.14 on all 4 DP ranks (thres 3.7)
CUDA coredumps	3	0
Unit: triton expand kernels, kv=1 / qo=6 row	`[-4,-3,-2,-1,0,1]`	`[0,0,0,0,0,1]`

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ❌ Run #28863798948
Latest PR Test (Extra): ❌ Run #28863798939

…>= 0 GLM 5.2 MTP (DP attention) crashed with an illegal memory access in topk_small_batch_kernel (#30274 rerun). Root cause: DP-padded / idle-companion draft-extend rows carry the CUDA-graph seq_len fill value (1), which is smaller than qo_len (num_draft_tokens), so the expanded per-row kv lengths went negative ([-4..1]); the top-k v2 kernel reads lengths as uint32_t, so -4 reinterprets as ~4.29e9, sending those rows down the cluster path over garbage scores and off the page table. - Clamp expanded seq_lens to >= 0 in fused_dsa_draft_extend_metadata and seqlens_expand_kernel (0 keeps padded rows on the trivial all-(-1) output path). - Drop the TEMP allow_topk_v2 gate from #30274: MTP target-verify / draft-extend dispatch to top-k v2 again. - Document the non-negative seq_lens contract on the v2 Python interface (plan_topk_v2 / topk_transform_512_v2 / _topk_transform_v2_paged). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

DarkSharpness · 2026-07-07T09:31:22Z

/rerun-test registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py

github-actions · 2026-07-07T09:31:47Z

Results for /rerun-test registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py:

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py
cd test/ && python3 registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py

gemini-code-assist

Code Review

This pull request enables the fused top-k v2 kernel for speculative verification and draft-extend modes by ensuring sequence lengths are clamped to non-negative values (>= 0) in Triton kernels (_fused_dsa_draft_extend_metadata_kernel and seqlens_expand_kernel). This prevents negative lengths from being reinterpreted as large uint32_t values, which previously caused illegal memory access crashes. It also cleans up the obsolete allow_topk_v2 flags and adds detailed documentation regarding this safety requirement. I have no feedback to provide as there are no review comments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

DarkSharpness · 2026-07-07T10:44:03Z

(generated by claude) The previously-failing MTP tests pass on this branch: rerun 28856177008 ✅

TestGLM52NVFP4TPMTP (tp4): 2 tests OK in 339.6s, avg_spec_accept_length=3.95
TestGLM52NVFP4DPMTP (tp4 dp4 + dp-attention): 2 tests OK in 223.4s, avg_spec_accept_length=4.01 — this is the config that previously illegal-memory-accessed during server warmup (failing run)

Both run with the fused top-k v2 path enabled for MTP (TEMP gate removed).

The fused v2 PAGED dispatch (from #30274) asserts that DSAMetadata carries a preprocessed topk_v2_plan instead of silently recomputing it per layer. The equivalence test builds DSAMetadata by hand without the plan, so test_topk_fused_backends_equivalence started failing on trunk for every PR (e.g. #30275's base-b-test-1-gpu-large job). Build the plan the same way init_forward_metadata / _build_forward_metadata_cuda_graph do; the sgl-kernel-vs-flashinfer equivalence check now actually covers the v2 fused path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

DarkSharpness · 2026-07-07T11:48:17Z

(generated by claude) Pushed 34e1ca1 folding in another #30274 regression: test_dsa_indexer.py::test_topk_fused_backends_equivalence fails on trunk for every PR (e.g. #30275's base-b job) with AssertionError: topk_v2_plan must be preprocessed per forward.

The v2 PAGED dispatch asserts the plan is preprocessed with the metadata (by design — no silent per-layer recompute), but the test builds DSAMetadata by hand without it. Fixed by building the plan in the mock metadata the same way init_forward_metadata / _build_forward_metadata_cuda_graph do. Side benefit: the sgl-kernel-vs-flashinfer equivalence check now actually exercises the v2 fused path.

Verified locally: test_dsa_indexer.py — 11/11 OK (was 1 failure / 3 parametrizations).

DarkSharpness · 2026-07-07T11:48:28Z

/rerun-test registered/kernels/test_dsa_indexer.py

github-actions · 2026-07-07T11:48:52Z

Results for /rerun-test registered/kernels/test_dsa_indexer.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/test_dsa_indexer.py

Fridge003 · 2026-07-07T20:08:44Z

/rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_hisparse.py test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py

github-actions · 2026-07-07T20:09:12Z

Results for /rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_hisparse.py test/registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py:

🚀 8-gpu-h200 (3 tests): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_dsa_glm52_tp_mtp.py
cd test/ && python3 registered/models_e2e/test_dsa_glm52_dp_mtp.py
cd test/ && python3 registered/models_e2e/test_dsa_glm52_hisparse.py

🚀 4-gpu-b200 (3 tests): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py
cd test/ && python3 registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py
cd test/ && python3 registered/cuda_graph/piecewise/test_pcg_glm52_fp4.py

…TP: clamp padded-row seq_lens to >= 0 (#30378) (#30427) Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…>= 0 (sgl-project#30378) Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Negative per-row lengths (DP-padded / idle-companion rows, e.g. -4 from GLM 5.2 MTP draft-extend metadata) are reinterpreted as ~4e9-token rows by the unsigned device-side reads in both top-k kernels. Depending on how the unsigned chunk arithmetic wraps, this is either an illegal memory access (v1: any negative; v2: e.g. -1048576 on every dispatch shape) or a silently garbage output row (v2: -1/-4), all reproduced deterministically on B200. sgl-project#30378 clamped two DSA triton producers, but the kernels stayed exposed to every other caller (the DSv4 DP-attention path has no clamp). Per review feedback, handle this with signed comparisons instead of an arithmetic clamp: negative rows take the documented trivial all-(-1) path (a non-negative candidate index never qualifies against a negative length). v1 keeps seq_len signed through the trivial dispatch and naive_transform; v2 dispatches via a signed is_trivial() in both the main and fused small-batch kernels; the plan kernel counts signed and proves rows positive before the (unsigned) pool-routing compare, so a full-uint32 static_cluster_threshold keeps its meaning. Zero cost: paired timing on all six dispatch shapes is unchanged within noise.

…TP: clamp padded-row seq_lens to >= 0 (sgl-project#30378) (sgl-project#30427) Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…>= 0 (sgl-project#30378) Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

DarkSharpness requested review from 1am9trash, BBuf, Fridge003, HaiShaw, HydraQYH, Qiaolin-Yu, YAMY1234, celve, hebiao064, hubertlu-tw, ispobock, kkHuang-amd, merrymercy, rainj-me and yuan-luo as code owners July 7, 2026 09:31

github-actions Bot added the jit-kernel label Jul 7, 2026

gemini-code-assist Bot reviewed Jul 7, 2026

View reviewed changes

DarkSharpness mentioned this pull request Jul 7, 2026

Fix DSV4 prefill large Triton recompilation idle across context lengths #30255

Merged

Fridge003 approved these changes Jul 7, 2026

View reviewed changes

Fridge003 merged commit bbc5370 into main Jul 7, 2026
106 of 120 checks passed

Fridge003 deleted the dsa-topk-v2-mtp-seqlens branch July 7, 2026 20:44

Kangyan-Zhou mentioned this pull request Jul 7, 2026

[Cherry-pick to release/v0.5.15] [DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0 (#30378) #30427

Merged

mmangkad mentioned this pull request Jul 8, 2026

[CI] Fix DSA top-k v2 test metadata #30383

Closed

vincentzed mentioned this pull request Jul 8, 2026

[DSA] Opt-in MXFP4 index-K cache for GLM-5.x / DeepSeek-V3.2 — fused-kernel MXFP4 mode (port+extension of #26209) vincentzed/sglang#10

Open

vincentzed mentioned this pull request Jul 8, 2026

[DSA] Opt-in MXFP4 index-K cache for GLM-5.x / DeepSeek-V3.2 (port of #26209), with fused-kernel MXFP4 mode #30546

Open

amd-bot mentioned this pull request Jul 8, 2026

[CI Monitor] Daily Report - 2026-07-08 bingxche/sglang-ci-bot#129

Open

This was referenced Jul 9, 2026

[CI Monitor] Daily Report - 2026-07-09 bingxche/sglang-ci-bot#130

Open

[CI Monitor] Daily Report - 2026-07-10 bingxche/sglang-ci-bot#131

Open

zkyue mentioned this pull request Jul 14, 2026

[BugFix][DSA] Harden top-k v1/v2 kernels against negative padded seq_lens #31123

Closed

5 tasks

amd-bot mentioned this pull request Jul 15, 2026

[CI Monitor] Daily Report - 2026-07-14 bingxche/sglang-ci-bot#135

Open

weireweire mentioned this pull request Jul 15, 2026

[DSV4] Fix draft-extend CUDA graph padding and WAR ordering #31270

Open

3 tasks

This was referenced Jul 16, 2026

[CI Monitor] Daily Report - 2026-07-16 bingxche/sglang-ci-bot#137

Open

[CI Monitor] Daily Report - 2026-07-18 bingxche/sglang-ci-bot#139

Open

Navigation Menu

[DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0 - #30378

[DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0#30378
Fridge003 merged 2 commits into
mainfrom
dsa-topk-v2-mtp-seqlens

DarkSharpness commented Jul 7, 2026 •

edited by github-actions Bot

Loading

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026 •

edited

Loading

Uh oh!

Fridge003 commented Jul 7, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

DarkSharpness commented Jul 7, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Root cause

Changes

Validation (local 4x B200, exact CI DP config)

CI States

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

DarkSharpness commented Jul 7, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Jul 7, 2026

Uh oh!

github-actions Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

DarkSharpness commented Jul 7, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jul 7, 2026 •

edited

Loading

github-actions Bot commented Jul 7, 2026 •

edited

Loading

github-actions Bot commented Jul 7, 2026 •

edited

Loading