[Spec decode] Support variable-length decode for Kimi-K3 adaptive ver - #52988
Conversation
| max_query_len = ( | ||
| int(query_lens_cpu.max().item()) if query_lens_cpu.numel() else 1 | ||
| ) |
There was a problem hiding this comment.
we should pull max_query_len from CommonAttentionMetadata; with adaptive spec-decode we uniformly distribute the budget on the CPU side but this may not be the case on the GPU side. So CommonAttentionMetadata.max_query_len represents the real max ( query_lens_cpu.max().item() may be less)
There was a problem hiding this comment.
thanks @LucasWilkinson now _build_decode now takes max_query_len from CommonAttentionMetadata instead of measuring query_lens_cpu.max(). I also made the FlashInfer MLA flatten map rows via searchsorted on the device offsets and routed every active KDA decode through the spec path, would appreciate you sanity checking those two.
There was a problem hiding this comment.
@LucasWilkinson could you take a look? thanks!
|
This pull request has merge conflicts that must be resolved before it can be |
mergify
Bot
added
needs-rebase
rocm
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Hi @qiching, would you mind rebasing? |
f94ee8e to
9e69f07
Compare
TheEpicDolphin
left a comment
TheEpicDolphin
left a comment
There was a problem hiding this comment.
Thx for the PR! Left some feedback
| # Ragged decode (flashinfer #3238); unset for single-token decode. | ||
| cum_seq_lens_q: torch.Tensor | None = None |
There was a problem hiding this comment.
This property seems redundant because it's just query_start_loc_device when max_query_len > 1. Can we remove it?
Also, can we rename query_start_loc_device => query_start_loc, which conforms better to the general metadata naming?
| assert decode_backend is None, ( | ||
| "FlashInferMLA ragged decode requires trtllm-gen, but num_heads=" | ||
| f"{runtime_num_heads} forces the cute-dsl backend, which does not " | ||
| "support cum_seq_lens_q." | ||
| ) |
There was a problem hiding this comment.
Is this assert needed? codex is flagging that cutedsl in the currently pinned flashinfer version supports cum_seq_lens_q.
| class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]): | ||
| # trtllm-gen tiles ragged queries from cum_seq_lens_q (flashinfer #3238), so one | ||
| # k+1 graph replays any 1..k+1 mix (full varlen decode, not piecewise). | ||
| _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
There was a problem hiding this comment.
ALWAYS is too lenient here as well, for the same reason as above for kda.py. This would allow replaying the captured graph for prefill/mixed batches during cudagraph_mode=FULL, but the capture run does not record the prefill kernels because it doesn't enter this branch: https://github.com/vllm-project/vllm/blob/main/vllm/models/kimi_k3/nvidia/mla.py#L685-L696
Should be downgraded to UNIFORM_BATCH
There was a problem hiding this comment.
Kept, adaptive_verification.py rejects anything below ALWAYS, so UNIFORM_BATCH makes the engine raise at initialize_kv_cache. Same combination already exists upstream: DeepseekV4FlashMLAMetadataBuilder declares ALWAYS while branching on num_prefill_tokens > 0.
| # ALWAYS (overrides GDN's UNIFORM_BATCH): KDA reads per-request offsets off | ||
| # device within a fixed k+1 window, so one k+1 graph replays any 1..k+1 mix. | ||
| _cudagraph_support = AttentionCGSupport.ALWAYS |
There was a problem hiding this comment.
I think ALWAYS is too lenient. When cudagraph_mode=FULL, prefills and mixed batches will be admitted into full cudagraphs, but that's unsafe for KDA because the captured batch is always decode-shaped, but KimiK3DeltaAttention._forward branches on composition:
vllm/vllm/models/kimi_k3/nvidia/kda.py
Line 1098 in c1b69aa
You can downgrade this to UNIFORM_BATCH and adaptive verification should still work. query_len_support = VARLEN is the property that matters for ragged decode.
There was a problem hiding this comment.
Same reason above. VARLEN is set as you noted, but alone it's not enough, UNIFORM_BATCH fails at startup with "must report AttentionCGSupport.ALWAYS, but KimiK3KDAAttentionBackend reports UNIFORM_BATCH". Verified on B300 8×GPU TP=8, 144 points across K=1..7 and adaptive, ±RecoverSSM.
|
/ci run |
|
✅ Triggered Buildkite CI #90244 for commit |
|
@qiching the test failures seem related. Plz take a look when you get the chance |
The merge picked up tests added on main that predate this PR's signature change: three _build_decode call sites in the aiter MTP split tests, and the FlashInfer MLA DCP tests whose MagicMock metadata never stubbed decode.max_query_len. The mamba hybrid SimpleNamespace gets the field InputBatch already carries. Signed-off-by: Albert Cheng <albecheng@nvidia.com>
Head branch was pushed to by a user without write access
|
/ci run |
|
✅ Triggered Buildkite CI #90469 for commit |
The test drives forward_mqa with a MagicMock self, so on the non-causal multi-token path self._flattened_decode_metadata resolved to a mock whose return value unpacked to nothing. Signed-off-by: Albert Cheng <albecheng@nvidia.com>
|
/ci run |
|
❌ This PR is 8 commits behind upstream |
Signed-off-by: Albert Cheng <albecheng@nvidia.com>
|
/ci run |
|
✅ Triggered Buildkite CI #90507 for commit |
The related test failures are fixed. The remaining ones are unrelated: the kernels MoE IPC test is broken on main by #57312 (fix in #58107), and the H200 MIG failures came from a network outage that broke HF model downloads (Network is unreachable / DNS errors). |
|
/ci retry |
|
✅ Queued 7 failed job(s) for retry in Buildkite CI #90507. |
…id models The scheduler pads a one-token prompt tail over prior state (e.g. a P/D decode-node arrival) with K placeholder drafts to keep the K+1 spec-decode shape (vllm-project#45237). For recurrent-state layers these rows must run the spec-decode kernels, which keep the running state and roll the rejected placeholders back; the prefill kernels store only the state after all K+1 tokens (vllm-project#55178). vllm-project#52988 changed the hybrid model state to classify spec-decode rows by request state (not prefilling) instead of num_scheduled_tokens, which adaptive verification rewrites. That sent every padded tail's draft count to -1, so Mamba's padded-tail handling no longer fires and GDN/KDA build the row as a prefill, folding the placeholder tokens into the request's recurrent state. Also count prefilling rows with exactly one remaining prompt token and prior state as decodes, still using request state only. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
|
@qiching @TheEpicDolphin @benchislett I found that this introduced a regression, please see #58434. |
…id models The scheduler pads a one-token prompt tail over prior state (e.g. a P/D decode-node arrival) with K placeholder drafts to keep the K+1 spec-decode shape (vllm-project#45237). For recurrent-state layers these rows must run the spec-decode kernels, which keep the running state and roll the rejected placeholders back; the prefill kernels store only the state after all K+1 tokens (vllm-project#55178). vllm-project#52988 changed the hybrid model state to classify spec-decode rows by request state (not prefilling) instead of num_scheduled_tokens, which adaptive verification rewrites. That sent every padded tail's draft count to -1, so Mamba's padded-tail handling no longer fires and GDN/KDA build the row as a prefill, folding the placeholder tokens into the request's recurrent state. Also count prefilling rows with exactly one remaining prompt token and prior state as decodes, still using request state only. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Rename get_varlen_decode_cudagraph_max_query_len() to get_varlen_cudagraph_max_query_len() and return None both for builders reporting ALWAYS, which replay any batch, and for builders that cannot replay variable-length batches. Only a bounded builder returns a length: FlashInfer returns 1 + num_speculative_tokens while its TRTLLM-GEN varlen decode path is active. Adaptive verification accepts a target builder that reports ALWAYS or a bound of at least the verification width, which drops the max_num_batched_tokens stand-in for "no limit". Keep batches with a prefill out of varlen decode graphs. Dispatch matched them on max(num_scheduled_tokens) alone, so decodes plus a prefill chunk of at most 1 + num_speculative_tokens tokens replayed a decode graph. FlashInfer and FlashInfer MLA run such chunks on their decode kernels, but the Kimi-K3 KDA builder (ALWAYS since vllm-project#52988) classifies rows by request state and does not restage its graph buffers when a prefill row is present. The runner now passes max_query_len=None for batches with a prefill; dispatch already never matches None against a bounded graph. Validation on GB300: pre-commit and mypy pass, and 3,201 tests pass across the adaptive verification, CUDA graph, attention, MLA, sparse MLA, config, Kimi-K3 and Gemma4 suites. The DCP distributed tests and two gated Hugging Face tests also fail on the previous revision, and one fp8 MLA case passed 3/3 on rerun. In a run mixing verification with 2-7 token prompts, all 32 steps that pair them now run piecewise; the previous revision replayed the varlen decode graph for all 31. Gemma4 + DSpark K=7 captures the same FULL graphs as before (fixed 2 + 32, adaptive 2 + 35); GSM8K is 25/50 fixed and 24/50 adaptive, against 24/50 and 24/50. Signed-off-by: zixi-qi <zixi@inferact.ai> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Thanks @njhill, my oversight on the padded prompt tail case. Approved #58434. |
6 participants
Summary
Make Kimi-K3's MLA and KDA decode paths capturable as FULL varlen CUDA graphs so they work under adaptive DSpark verification (#47808), which schedules ragged per-request draft budgets and requires device-sourced query lengths. Before this change K3's
FLASHINFER_MLAreportedUNIFORM_BATCH/UNIFORM, and the decode builders derived per-request query length from a measured/averaged count, so a graph captured on a uniform dummy did not match ragged replay.Changes
FlashInfer MLA (
flashinfer_mla.py)_cudagraph_supporttoAttentionCGSupport.ALWAYSandquery_len_supporttoQueryLenSupport.VARLEN.cum_seq_lens_q+max_q_lentotrtllm_batch_decode_with_kv_cache_mla(flashinfer Is it always better to have more GPUs during parallel inference? Why is it better for me to have two graphics cards in parallel than four graphics cards #3238); assert this path runs on trtllm-gen and does not return LSE (mutually exclusive with DCP).searchsortedon the device query offsets instead ofrepeat_interleave(uniform_len).query_start_loc_device,cum_seq_lens_q, andmax_query_lenon the decode metadata; setcum_seq_lens_qonly whenmax_query_len > 1.MLA common (
mla_attention.py+ subclass builders)max_query_lenthroughMLACommonMetadataBuilder._build_decode(clamped toreorder_batch_threshold) and theflashmla,flashattn_mla,rocm_aiter_mla, anddots3_notebuilders.KDA (
kda_metadata.py)KimiK3KDAMetadataBuildertoAttentionCGSupport.ALWAYS.is_prefillingrather than draft count.Truefromsupports_device_cpu_query_lens_mismatch.KDA RecoverSSM (
kda_metadata.py, opt-in via--use-replayssm, off by default)num_spec+1slots to 1 per request (kernels already in [K3] support recoverssm for K3 #51855):self.layer_nameson the builder so RecoverSSM's commit context can resolve the per-layer forward context (the GDN base does not retain it; inert when RecoverSSM is off).query_lenexceedsnum_spec+1out of the RecoverSSM path instead of raising. Adaptive's cost-table profiler probes token counts past the cudagraph limit by even-splitting them over the capped request slots, producing eager dummy rows longer than RecoverSSM'snum_spec+1verify buffers; real captured/serving decode rows are always<= num_spec+1, so they are unaffected.Mamba-hybrid model state (
mamba_hybrid.py)max_query_len, and classify spec-decode rows offis_prefillingrather thannum_scheduled_tokens == draft_count + 1.K3 DSpark draft (
dspark_mla.py)compute_confidence.Related Issue: #51867
Test plan
_build_decodesignature (
test_mla_backends.py,test_rocm_aiter_mla_mtp_split.py).Validation
Kimi-K3 (target) +
Inferact/Kimi-K3-DSpark(draft), TP8 on 8×B300, FP8 KV cache,attention_backend=FLASHINFER_MLA,num_speculative_tokens=7,enable_adaptive_verification=true. (Numbers below are RecoverSSM off; it is opt-in.)Accuracy — lossless vs. no-spec (GSM8K, in-tree runner, 1319 questions, 5-shot, temperature 0)
Coherence (MTBench, 80 prompts, temperature 1)
Adaptive: 80/80 successful requests, 0 failed, ~19.5K generated tokens; no repetition or degeneration in the saved completions.
Speculative-decoding health (speed_bench, temperature 1, c16)
Acceptance rate 28.95%, acceptance length 3.03, per-position acceptance decaying 74.5% (pos 0) → 6.9% (pos 6) — non-zero at every position, confirming the draft + confidence head + varlen verification path is exercised end to end.
RecoverSSM (opt-in) health
Adaptive + RecoverSSM serves through full-graph capture with healthy acceptance length (2.4–2.9 across c1–c256), confirming the RecoverSSM verify/commit path is exercised end to end; matches base-adaptive accuracy/coherence.
Benchmarking
Pareto sweep of the varlen decode path: 16 configurations × 9 concurrency levels = 144 serving runs. Configurations are fixed draft length K=1..7 plus adaptive, each with and without RecoverSSM.
Kimi-K3 (target) +
Inferact/Kimi-K3-DSpark(draft), TP8 on 8×B300, FP8 KV cache,attention_backend=FLASHINFER_MLA,draft_sample_method=probabilistic,num_speculative_tokens=7for adaptive.speed_benchqualitative,--speed-bench-output-len 2048,--max-model-len 16384, prompt count scaled with concurrency. Axes: per-user token rate (1000 / median TPOT) vs aggregate output throughput.Overlaying both families lets a non-RecoverSSM point appear to overtake a RecoverSSM one it never competes with, so each family is also plotted with its own frontier.
RecoverSSM, adaptive off vs on