[Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) - #55737
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>
|
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. 📝 SummarySummary by CodeRabbit
WalkthroughThe KDA layer validates backend constraints and prepares FlashKDA resources. FlashKDA prefill now supports workspace-backed output for mixed speculative steps and returns both output and final state. Results are written directly into ChangesGLM-5.3-Flash KDA prefill
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to This change enables FlashKDA prefill on supported hardware and routes mixed speculative outputs directly into layer buffers. No concrete merge-blocking risk remains. Sequence Diagram(s)sequenceDiagram
participant GLM5NextKDA
participant FlashKDA
participant WorkspaceManager
participant core_attn_out
GLM5NextKDA->>FlashKDA: execute prefill with optional output
FlashKDA->>WorkspaceManager: obtain output and final-state buffers
WorkspaceManager-->>FlashKDA: return workspace buffers
FlashKDA-->>GLM5NextKDA: return output and final state
GLM5NextKDA->>core_attn_out: write ordinary or scattered speculative results
Suggested reviewers: 🚥 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, CI is now available for this PR.
|
| safe_gate=safe_gate, | ||
| lower_bound=lower_bound, | ||
| ) | ||
| if self.kda_prefill_backend == "flashkda" and not use_spec: |
There was a problem hiding this comment.
I remember flashkda also supports spec decoding?
There was a problem hiding this comment.
I remember not,need more profiling and optimization on SpecDecode path tho..
There was a problem hiding this comment.
you can check this: vllm/models/kimi_k3/nvidia/kda.py, no such guard
There was a problem hiding this comment.
thanks, will work on it
There was a problem hiding this comment.
🔵 Needs a closer look
It changes a performance-critical prefill kernel path to call into a fused CUDA extension with new runtime selection and workspace usage, which warrants final human review for correctness/compatibility across execution modes.
Pull request overview
This PR wires the existing FlashKDA CUDA extension (vllm._flashkda_C) into the GLM-5.3-Flash KDA chunked prefill path, replacing the current Triton chunk_kda_with_fused_gate route when the platform/config is compatible (CUDA SM90/SM10x/SM12x, bf16, head_dim=128, bounded gate). The goal is to reduce prefill latency by using a more fused implementation and avoiding the pre-sigmoid beta casting and the merge copy on the non-spec prefill path.
Changes:
- Add backend resolution (
auto/triton/flashkda) for KDA prefill and store the chosen backend on the layer. - Add a FlashKDA prefill wrapper that uses the v1 workspace manager to allocate per-call scratch/state buffers and invokes
torch.ops._flashkda_C.fwd. - Route non-spec prefill through FlashKDA when enabled; otherwise keep the existing Triton chunk path (and always keep Triton for spec-decode steps).
File summaries
| File | Description |
|---|---|
vllm/models/glm5next/nvidia/kda.py |
Adds FlashKDA backend selection and a FlashKDA-based chunked prefill execution path, falling back to the existing Triton chunk implementation when unsupported or disallowed. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
/ci run |
|
✅ Triggered Buildkite CI #87558 for commit |
…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>
|
/ci run |
|
✅ Triggered Buildkite CI #87906 for commit |
|
Field report from a different stack, in case it helps the review: three DGX Spark GB10 nodes (sm_121, arm64, CUDA 13.0) serving GLM-5.3-Flash as an EXL3 4.05 bpw checkpoint at TP=3 + expert parallelism, vLLM pinned at 487ecf1 with this PR's
Full write-up, harness and the port: https://github.com/NNNtrance/GLM-5.3-Flash-EXL3-DGX-Spark/blob/main/results/gates/flashkda-ab-10sep.md. It is in our production as of today. Thanks for the port. |
|
/ci run |
|
✅ Triggered Buildkite CI #88089 for commit |
|
@JaredforReal I noticed you are using |
|
@gaby nope, now glm-5.3-flash in vllm has only one attention backend, which is FLASHINFER_MLA_SPARSE |
|
/ci run |
|
✅ Triggered Buildkite CI #88550 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #88754 for commit |
Purpose
One of three independent GLM-5.3-Flash perf PRs
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).Test Plan
lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --gen_kwargs temperature=0(1319 questions), prompt-logprob agreement on real 4k/12k/30k/60k prompts vs main, and a 36-needle retrieval set at 6k/16k/30k (greedy,reasoning_effort=low), repeated 3x on main / +[Perf][GLM-5.3-Flash] Decode hot-path cleanups: strided KDA recurrent inputs, NoPE MQA query without concat, no duplicate router GEMM #55736 / full stack.Test Result
Performance of this PR alone
Measured on top of #55736 (main + #55736) 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. #55736 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.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 (per build, same session)
Each rung of the ablation was also evaluated on its own: gsm8k (1319 questions, 5-shot, greedy, lm_eval
local-completionsDuplicate-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 (#55224 is the KDA conv weight refit fix).vllm/models/kimi_k3/...) and adds a--kda-prefill-backendengine arg; this PR touches onlyvllm/models/glm5next/.... Once [Perf] Integrate FlashInfer KDA kernels #55364 lands, theadditional_config.kda_prefill_backendknob here should move to that arg.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
.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.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.