[Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) - #8
Closed
JaredforReal wants to merge 2 commits into
Closed
JaredforReal wants to merge 2 commits into
JaredforReal wants to merge 2 commits into
Conversation
FlashKDA (already built as vllm._flashkda_C for Kimi-K3) implements the same bounded-gate KDA recurrence as chunk_kda_with_fused_gate (lower_bound * sigmoid(exp(A_log) * (g + dt_bias)), in-kernel q/k l2norm, raw beta logits). On GB300 it replaces ~15 Triton kernels with 4 and is 1.7-3.8x faster (T=2048x8: 183->49 us, 8192x4: 555->146 us, 16384x1: 1371->787 us per layer). Select it automatically when supported (SM90/SM10x/SM12x, bf16, head_dim 128, bounded gate); additional_config.kda_prefill_backend = triton keeps the old path. The output is written straight into the layer buffer and the pre-sigmoid beta cast is no longer needed on this path. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
There was a problem hiding this comment.
🟡 Changes recommended
Auto-selecting flashkda can crash on CUDA builds where the optional vllm._flashkda_C extension is not present, instead of cleanly falling back to the Triton path.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR wires the optional FlashKDA CUDA extension (vllm._flashkda_C) into the GLM-5.3-Flash KDA chunked prefill path to replace the existing Triton chunked kernels when supported, aiming to reduce prefill latency.
Changes:
- Add backend selection logic (
auto/triton/flashkda) for KDA chunked prefill, based on device capability, dtype, head_dim, and bounded-gate configuration. - Add a FlashKDA prefill implementation that uses the v1 workspace manager for scratch buffers and can write directly into the layer output buffer when not in spec-decode.
File summaries
| File | Description |
|---|---|
| vllm/models/glm5next/nvidia/kda.py | Adds FlashKDA-backed chunked-prefill path and backend selection via additional_config.kda_prefill_backend. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+127
to
+150
| def _resolve_kda_prefill_backend( | ||
| backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None | ||
| ) -> str: | ||
| """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on | ||
| SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton | ||
| ``chunk_kda_with_fused_gate`` path. ``backend`` comes from | ||
| ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" | ||
| if backend not in ("auto", "triton", "flashkda"): | ||
| raise ValueError(f"Unsupported KDA prefill backend: {backend}") | ||
| capability = current_platform.get_device_capability() | ||
| supported = ( | ||
| current_platform.is_cuda() | ||
| and capability is not None | ||
| and capability.major in (9, 10, 12) | ||
| and head_dim == 128 | ||
| and dtype == torch.bfloat16 | ||
| and lower_bound is not None | ||
| ) | ||
| if backend == "flashkda" and not supported: | ||
| raise RuntimeError( | ||
| "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 " | ||
| "and a bounded KDA gate." | ||
| ) | ||
| return "flashkda" if supported and backend != "triton" else "triton" |
This was referenced Sep 7, 2026
Owner
Author
|
Upstream draft: vllm-project#55737 |
…code steps too The FlashKDA path was gated on `not use_spec` so it could write straight into the layer output buffer; a step that also carried spec-decode tokens fell back to the Triton chunk path for its whole prefill segment, which with MTP enabled is almost every step that has a prefill. Follow the Kimi-K3 KDA layer instead: in a spec step FlashKDA writes to a workspace buffer and the non-spec tokens are scattered by non_spec_token_indx, alongside the spec tokens (which keep the recurrent kernel with num_accepted_tokens rollback). The merge now index_copy_s directly into core_attn_out instead of going through a temporary tensor. MTP k=1, TP4 on 4x GB300, same build, Triton -> FlashKDA: 8x2048 TTFT 464 -> 423 ms, 2x8192 390 -> 360 ms, 32k/256 c=16 260 -> 284 tok/s, 1k/512 c=64 2003 -> 2130 tok/s; gsm8k 93.33 vs 93.56 (+-0.7). Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
One of three independent GLM-5.3-Flash perf PRs (#7: decode cleanups; this one: FlashKDA prefill; #9: masked-MHA prefill). Each is benchmarked on its own below.
Summary
Use FlashKDA (
vllm._flashkda_C, already built and used by Kimi-K3) for GLM-5.3-Flash's KDA chunked prefill instead of the ~15-kernel Tritonchunk_kda_with_fused_gatepath.FlashKDA implements the same bounded-gate KDA recurrence GLM-5.3-Flash uses (
lower_bound * sigmoid(exp(A_log) * (g + dt_bias)), in-kernel q/k l2norm, raw beta logits), so it is a drop-in. On GB300 per KDA layer (H=16/rank, D=128,bench/flashkda_vs_chunk.py):Selection: automatic when supported (SM90/SM10x/SM12x, bf16, head_dim 128, bounded gate);
additional_config.kda_prefill_backend = tritonkeeps the old path (flashkdaforces it). The output is written straight into the layer buffer (no merge copy) and the pre-sigmoid beta cast is not needed on this path. Spec-decode tokens themselves keep the Triton recurrent kernel (num_accepted_tokensrollback); in a step that carries both, FlashKDA runs the prefill segment into a workspace buffer and the outputs are scattered back bynon_spec_token_indx/spec_token_indx(same scheme as the Kimi-K3 KDA layer; the merge writes straight intocore_attn_out).Performance of this PR alone
Measured on top of #7 (main + #7) vs the same + this PR, 4x GB300, TP4,
--attention-backend FLASHINFER_MLA_SPARSE --max-model-len 69632 --max-num-seqs 256 --max-num-batched-tokens 16384, prefix caching disabled (hit rate 0.0% checked in the server log),vllm bench serverandom dataset with warmups, back-to-back in the same session. #7 only touches decode, so this is the effect of FlashKDA on an otherwise-main prefill path. Decode tok/s is the steady-state window value; TPOT is the per-request median.Prefill: −12-13% TTFT for 8x2048 / 2x8192 and −8% for 8x32k / 2x64k (per-KDA-layer speedup shrinks with sequence length, see the kernel table above). 32k-context serving (c=16): +6.6% decode tok/s, −6.5% TPOT, −5% TTFT, because each step still prefills the 34 KDA layers of new requests. At c=256 with 1k prompts the same effect gives +4.7% tok/s (new prompts' prefill is interleaved with decode); c=1 / c=64 decode is unchanged (±1%, within the noise floor in the last column of the ablation below).
With speculative decoding (MTP k=1)
Same build, same settings plus
--speculative-config.method=mtp --speculative-config.num_speculative_tokens=1,kda_prefill_backend=tritonvsflashkda, back-to-back. Random-dataset acceptance is ~1.2 tokens/step, so absolute decode numbers are low; the comparison is what matters. Whole-run throughput (the steady-state window accounting is off under spec decode).Accuracy under MTP: gsm8k 93.33 ± 0.69 (Triton) vs 93.56 ± 0.68 (FlashKDA); needle 34/36 on both, within the model's run-to-run spread. A debug hook (removed) confirmed the FlashKDA-in-spec-step path ran in mixed steps with 1, 1024 and 15349 non-spec tokens.
Ablation of the whole series
Cumulative, same session, same settings: main → +#7 → +#8 (this PR) → +#9 → main again.
Accuracy (per build, same session)
Each rung of the ablation was also evaluated on its own: gsm8k (1319 questions, 5-shot, greedy, lm_eval
local-completions), prompt-logprob agreement on real 4k/12k/30k/60k prompts (mean |Δ logprob| per token vs the main build; the "main again" row is the run-to-run noise floor of this FP8 model), and a needle-in-a-haystack retrieval set (12 codes per length at 6k/16k/30k tokens, greedy,reasoning_effort=low), which exercises the 2k-36k prefill range where #9 switches kernels.| build | gsm8k flexible-extract % | gsm8k strict-match % | prompt-logprob mean|Δ| vs main (4k/12k/30k/60k) | next token | needle hits (6k/16k/30k) | needle, 3 more passes (concurrency 6) |
|---|---|---|---|---|---|---|
| main | 93.33 ± 0.69 | 93.10 ± 0.70 | (reference) | | 12/12 / 12/12 / 12/12 | 34/36 / 36/36 / 35/36 |
| +#7 | 93.03 ± 0.70 | 92.95 ± 0.71 | 0.040 / 0.055 / 0.051 / 0.108 | all same | 12/12 / 12/12 / 11/12 | 35/36 / 36/36 / 34/36 |
| +#7 +#8 | 93.25 ± 0.69 | 93.18 ± 0.69 | 0.099 / 0.072 / 0.060 / 0.136 | all same | 12/12 / 11/12 / 11/12 | - |
| +#7 +#8 +#9 | 93.10 ± 0.70 | 93.10 ± 0.70 | 0.099 / 0.072 / 0.060 / 0.127 | all same | 12/12 / 12/12 / 11/12 | 35/36 / 34/36 / 36/36 |
| main again | 92.87 ± 0.71 | 92.87 ± 0.71 | 0.041 / 0.055 / 0.051 / 0.118 | all same | 12/12 / 12/12 / 12/12 | - |
+#7vs main equals the main-vs-main noise floor to three digits, i.e. [Perf][GLM-5.3-Flash] Decode hot-path cleanups: strided KDA recurrent inputs, NoPE MQA query without concat, no duplicate router GEMM #7 is numerically equivalent to main (expected: the strided KDA kernel is bit-identical, the bmm into a transposedoutwas verified identical, and the router logits actually used were always the ones MoERunner recomputes). [Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) #8 and [Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat #9 raise the 4k |Δ| to ~2.4x the noise floor (FlashKDA and FA4 accumulate in a different order than the Triton chunk / per-token sparse kernels); 12k-60k stay at the noise floor and the next token is identical everywhere.Duplicate-work check
gh pr list --repo vllm-project/vllm --state open --search "GLM-5.3-Flash FlashKDA"/"glm5next KDA": no open PR wires FlashKDA for GLM-5.3-Flash (vllm-project#55224 is the KDA conv weight refit fix).Tests
bench/kda_numerics.py: FlashKDA vs the naive recurrent reference (fromtests/models/kimi_k3/test_kda.py) mean|Δ| 2e-5, max 4.8e-4 on outputs, same as the Triton chunk path vs the reference (T=48, 2 sequences, lower_bound -5). Note the Triton chunk path writes its output intov's buffer, so the comparison clonesv.pytest tests/models/kimi_k3/test_kda.py -k flashkda: 2 passed (extension sanity).pre-commit run ruff-check / ruff-format: passed.Notes for review
torch.ops._flashkda_C.fwddirectly (no cross-model import from kimi_k3); the support predicate mirrorskimi_k3.nvidia.kda.is_flashkda_supported..contiguous()copies of the conv-output slices remain (~115 us/layer per 16k chunk); a follow-up could run the short conv per q/k/v into separate buffers.~/notes/glm53flash-perf/REPORT.md, ablation logs in~/notes/glm53flash-perf/logs/e2e/ablate_*.