Conversation
…gits On SM100 the indexer metadata builder schedules DeepGEMM's paged MQA logits in varlen mode (per-token rows tagged with request indices), but the kpool indexer launched the kernel without `indices`, so the non-varlen kernel ran on a varlen schedule. The varlen scheduler pairs adjacent rows of a request into one atom charged with the later row's context length; the non-varlen walk treats every row as its own request and stops only on an exact (q_atom, kv_split) end match. Once a pair straddles a 256-state split boundary the walk steps past an end marker it can never hit, the TMA warp stops prefetching Q and the math warps wait forever on the Q-stage barrier: 100% GPU utilization with idle host threads until the engine times out. Single-token decode has one row per request, so both schedules coincide and the bug never shows without speculative decoding. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. WalkthroughThe decode path now passes ChangesDecode attention index propagation
Estimated code review effort: 1 (Trivial) | ~2 minutes Suggested reviewers: Merge Risk: ⚪ Minimal · up to The decode path now supplies varlen indices to paged MQA logits, preventing the affected speculative-decoding hang. The change is narrowly scoped and has no remaining merge-readiness risk. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
JaredforReal
left a comment
JaredforReal
left a comment
There was a problem hiding this comment.
while we understand we should align with ds's pattern passing indices=decode_metadata.indices down to kernel, the mtp hanging mechanism need a deepdive and repro to comfirm
|
/ci run |
|
✅ Triggered Buildkite CI #87214 for commit |
ivanium
changed the title
|
Update: the hang issue is caused by a stale deepgemm 2.5.0 version. Updating to 2.6.1 resolved it. This PR is still a little bit helpful to get metadata aligned, but it did not fix anything (there is nothing to fix) |
|
One additional data point: this mismatch is not only relevant to the hang behavior seen with older DeepGEMM revisions; it also has a noticeable performance impact with the current varlen-capable path. On GLM-5.3-Flash with context length 120K, batch size 40, and speculative decoding steps = 5, passing decode_metadata.indices changes sm100_paged_mqa_logits from the non-varlen instantiation (is_varlen=false, next_n=1) to the varlen instantiation (is_varlen=true). In an nsys capture, kernel latency dropped from ~214.3 us to ~51.8 us, about a 4.14x speedup for this kernel. Reference: #56811 |
4 participants
Note
Update: the original hang issue is caused by a stale deepgemm 2.5.0 version. Updating to 2.6.1 resolved it. This PR is still a little bit helpful to get metadata aligned, but it did not fix anything (there is nothing to fix)
Purpose
Make the kpool indexer's DeepGEMM paged-MQA launch agree with the schedule metadata it was built with. On SM100
DeepseekV32IndexerMetadataBuilderschedules the paged MQA logits in varlen mode (per-token rows tagged with request indices,get_paged_mqa_logits_metadata(..., indices=...)), andsparse_attn_indexer.pypassesindices=decode_metadata.indicestofp8_fp4_paged_mqa_logits. The kpool indexer did not, so DeepGEMM selects the non-varlennext_n=1kernel and runs it on a varlen schedule. Under MTP a request's verify rows share one request index, so the two decompositions differ whenever a request's rows straddle a 256-state split boundary (1024 tokens atindex_kpool=4).What that mismatch does depends on the DeepGEMM revision, verified with a standalone kernel script (same fp8 cache, block table and weights; schedule built with
indices; launched with and withoutindices; compared against a pure non-varlen reference):indiceson a varlen schedulemainruns fine (0.9704 / 0.9731).So with the pinned DeepGEMM this is a latent contract violation rather than a live bug. The one-line fix matches the non-kpool indexer, removes the dependence on the scheduler tolerating a mismatched launch, and keeps the kpool path working on DeepGEMM revisions with the strict scheduler.
Duplicate check: #55201 touches the kpool top-k sentinel handling, not this call; no open PR passes
indiceshere.Test Plan
--speculative-config '{"method":"mtp","num_speculative_tokens":3}', GSM8K 5-shot x2 (2638 requests, concurrency 32), single node TP4 and TP4-to-TP4 NIXL PD.Test Result
mainwithin 4 min every time; with the fix, single node 2638/2638, 0 request errors, 0.9723 / 0.9738 (0.9727 / 0.9742 without MTP); PD 2638/2638, 0 errors, 0.9708 / 0.9719.mainwithout the fix: 0.9704 / 0.9731, 0 errors.No unit test: the observable failure is a device deadlock on one DeepGEMM revision, and the op needs a live GPU, DeepGEMM and a populated kpool cache to reach the call.
AI assistance (Claude Code) was used for the capture, analysis and fix; the submitter reviewed the change.