[ormandj]

Motivation

Fixes #39173.

With SGLANG_RAGGED_VERIFY_MODE=compact, DSpark target verify can pack unequal numbers of tokens per request. A compact CUDA graph tier can also hold more requests than full verify windows: for example, 42 tokens over 8 slots captures as [6, 6, 5, 5, 5, 5, 5, 5]. Two DeepSeek-V4.1 components still assume draft_token_num tokens per request:

• EngramHasher.forward asserts num_tokens == bs * draft_token_num. Capture fails with engram target-verify expects one equal block per request, got 42 tokens for 8 requests of 6.
• token_req_indices calls repeat_interleave(req, draft_token_num, output_size=num_tokens), which fails, or maps tokens to the wrong request, once the layout is not uniform.

As a result, DeepSeek-V4.1 compact verify fails on these layouts.

#40567 by @lidekang160122 identified both call sites and fixes them as part of a larger PD change. This PR fixes only the crash. It derives the token-to-request mapping in one shared helper and reuses Engram's existing extend kernel path, so it needs no kernel changes.

Depends on the shared-layout fix from #42709, now in main: both components read the verify layout inside the CUDA graph, and without it the DeepSeek-V4.1 graph variants captured first for a token tier would read a layout that replay does not refresh.

Modifications

• ragged_verify.py: add ragged_verify_token_rows(layout, num_tokens). It is a fixed-size searchsorted over the layout's qo_indptr, so it is safe under CUDA graph capture. Empty rows own no tokens, and the graph-padding tail maps to row bs.
• EngramHasher.forward: with a ragged layout, target verify uses the existing MODE_EXTEND path with row = ragged_verify_token_rows(layout, num_tokens).clamp_max(bs - 1) and starts = layout.qo_indptr_device[:-1]. As before, verify commits no history. Graph-padding tokens are hashed with the row the staged layout assigns them (a padding row, or the last row for the tail) and then discarded.
• token_req_indices: with a ragged layout, each token gets its request's pool index. The tail past the last request gets pool slot 0, like the padded rows of a dense verify graph; padding the staged layout assigns to a row gets that row's slot.
• Without a layout (static verify), both functions take the existing dense path unchanged.
• No change to capture slot sizing. The current min(num_tokens, max_bs) slots are valid once these components accept unequal rows.

Accuracy Tests

• test/registered/kernels/ops/attention/dsv4/test_dsv41_compact_verify.py (CPU tests and a CUDA graph replay test; all pass on an RTX PRO 6000 Blackwell):
  • Packed Engram hashes equal the hashes of each request verified alone on the dense path. Cases: the 42-token / 8-slot capture geometry, a batch with empty rows and a graph-padding tail, and one request in a 3-token graph. All three cases run with maximum n-gram sizes 4 and 8 and include a short prefix; the first two also include an image token. Engram history is unchanged.
  • A uniform layout gives the same Engram hashes and request indices as the dense path.
  • token_req_indices maps empty rows and the padding tail correctly.
  • A CUDA graph captured once is replayed with three different layouts, and Engram hashes on real tokens and request indices on all tokens match eager.
• TestCapturedVariantsReadStagedLayout (CPU, same file) calls DecodeCudaGraphRunner's layout-capture helper twice for a 12-token tier, as the two DeepSeek-V4.1 attention variants do, stages live lengths [6, 1] for request slots 7 and 9, and checks that token_req_indices and EngramHasher read the live boundaries from either variant's layout. Without the shared layout from [Spec] Fix wrong tokens from compact-verify CUDA graphs when a token tier is captured more than once #42709, the first variant maps the real tokens to [7, 7, 9, 9, 0, 0, 0]. All six tests in the file, including the CUDA graph replay test, pass on an RTX PRO 6000 Blackwell.
• test/registered/spec/dspark/test_ragged_verify.py: unit test of ragged_verify_token_rows.
• An ad-hoc comparison with [Feat] dspark: Enable DeepSeek-V4.1 dynamic verify on PD decode #40567's implementation on an RTX PRO 6000 Blackwell (SM120), not part of this PR, gave bit-identical Engram hashes on real tokens and bit-identical request indices on all tokens. Coverage was 146 eager layouts (73 at each of n-gram sizes 4 and 8) and 22 CUDA graph replays, including the 42-token geometry, single-request graphs of 1 to 5 tokens, empty rows, and padding tails.
• End to end, I served DeepSeek-V4.1-Flash with [Feat] dspark: Enable DeepSeek-V4.1 dynamic verify on PD decode #40567's implementation and [Spec] Fix wrong tokens from compact-verify CUDA graphs when a token tier is captured more than once #42709 (not with this exact diff). Setup: 4x RTX PRO 6000 Blackwell (SM120, PCIe), TP4/EP2, DSpark gamma=5, compact verify with a profiled SPS table. I compared target log-probs of generated tokens with a plain prefill of the same text. The maximum per-token difference was 0.79-0.86 nats, the same as static verify. A temperature-1 run of 32 streams had no repetition loops.

Speed Tests and Profiling

The change itself adds one searchsorted per Engram call and per low-ratio request lookup in compact verify. Static verify is unchanged. Not benchmarked separately.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #41943891256
Latest PR Test (Extra): ❌ Run #41943891193
Latest PR Test (AMD ROCm 10): ✅ Run #41943891412

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[Qiaolin-Yu]

/rerun-test registered/kernels/ops/attention/dsv4/test_dsv41_compact_verify.py registered/spec/dspark/test_ragged_verify_backend_capability.py registered/unit/speculative/test_dspark_engram_commit_in_graph.py registered/core/test_basic_sanity_dspark.py registered/e2e/models/test_deepseek_v41_flash_gb300_tp2ep2.py registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py

[github-actions]

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

cd test/ && python3 registered/kernels/ops/attention/dsv4/test_dsv41_compact_verify.py
cd test/ && python3 registered/core/test_basic_sanity_dspark.py

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/dspark/test_ragged_verify_backend_capability.py
cd test/ && python3 registered/unit/speculative/test_dspark_engram_commit_in_graph.py

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

cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_tp2ep2.py
cd test/ && python3 registered/e2e/models/test_deepseek_v41_flash_gb300_dp2ep2.py