[suiyoubi]

Overview

This PR moves GLM-5.3-Flash's KDA (Kimi Delta Attention) layers onto FlashInfer kernels:

• the fused single-token decode (fused_kda_decode);
• the packed multi-token decode for draft-verify rows (packed_fused_kda_decode);
• optionally, the CuTe DSL chunked prefill (recurrent_kda(backend="cute-dsl")).

Kimi K3 already uses these kernels through the same vllm.utils.flashinfer wrappers and the same
kda_decode_backend / kda_spec_decode_backend / kda_prefill_backend options. The model has 34 KDA layers
(64 heads of 128). Today each decode step runs three Triton kernels per layer, and mixed steps send their decode
rows through the chunked prefill kernel.

• Decode: default auto; activates where FlashInfer supports the layer.
• Prefill: opt-in --kda-prefill-backend flashinfer.

Important

Requires FlashInfer 0.7.3. This PR assumes the FlashInfer pin bump to 0.7.3 (a separate PR) has landed.
0.7.3 contains flashinfer-ai/flashinfer#7132 (16/64 heads in fused_kda_decode / packed_fused_kda_decode, and
the DS conv-state layout) and #6201 (SM107 prefill). GLM-5.3-Flash's 64-head (TP1 / DP attention) and 16-head
(TP4) ranks, and the DS layout that NIXL P/D needs, rely on #7132. With the current pin (0.7.2.post1) the static
support checks accept configurations that 0.7.2 rejects at its first call, so this PR should merge after the bump.

Claims

• Decode kernels (VR200, 34 layers per step, CUDA graph, L2 flushed):
  • 1.9–2.1x faster at 1–8 rows per GPU;
  • 1.3–1.8x faster at 32–512 rows (H64 fp32 at 512 rows: 19.5 → 12.4 ms per step; H64 bf16: 16.2 → 9.2 ms);
  • pure decode steps run in a FULL CUDA graph (FULL_AND_PIECEWISE, KDA captured), so the GPU times above are
    what serving pays; in eager (PIECEWISE, mixed/prefill) steps the host enqueue also drops, 3.0 → 0.95 ms per
    34 layers.
• Mixed prefill + decode steps (one layer, H64): 46–61% faster (32K prefill tokens + 512 decode rows:
  6.96 → 2.72 ms). Most of it is the decode-row peel.
• FlashInfer prefill (opt-in): 8–12% per layer at 4+ long sequences; 3–8% slower in eager steps with 1 sequence
  at H64 or a few at H16, mostly from the call's host time (about 95 µs vs FlashKDA's 18 µs on VR200, where
  FlashInfer's kernel is faster). This is why it stays opt-in.
• Serving on GB300 and GB200, native srt-slurm sa-bench 8K/1K, Dynamo, base vs this PR (+ the FlashInfer change):
  • agg DEP2 (64 KDA heads per rank): +13.7% / +24.9% TPS/GPU on GB300 at C512 / C1024 (2847 → 3237, 2402 → 2999);
    +13.5% / +14.8% on GB200, which is KV-capacity bound;
  • disagg with the DS conv layout via NIXL, C512: +5.1% (GB300) and +6.0% (GB200) TPS/GPU;
  • agg TP4 (16 heads per rank): +1.4–3.0% at C1–C32.
• Where it comes from: decode steps run one fused FlashInfer kernel per layer instead of three Triton kernels,
  and mixed steps send their decode rows to that kernel instead of the chunked prefill kernel (the decode-row peel).
• FlashInfer prefill (opt-in): faster for several long sequences per step; slower for one or a few sequences at
  low concurrency, mostly from the call's host time (Follow-ups). This is why it stays opt-in.
• Also validated on VR200 (SM107): unit tests, GSM8K accuracy (with and without MTP) and end-to-end serving,
  including disaggregated P/D; performance numbers for VR200 are not included here.
• Accuracy: 40-shot GSM8K (7–8K-token prompts), agg DEP2 on VR200, 6 runs each: base 91.85 ± 0.59%, PR (decode + prefill on FlashInfer) 91.88 ± 0.76%.
• Numerics vs an fp32 reference:
  • FlashInfer decode error ≤ the Triton path's;
  • conv states bit-identical to the Triton path's;
  • in mixed steps, the decode rows equal a decode-only step bit for bit, and the prefill rows (outputs and
    states) are bit-identical to the Triton path.

Validation

Unit tests

tests/models/glm5next/test_kda_flashinfer_decode.py (new):

• backend resolution and fallback;
• decode steps for 8/16/32/64 heads × fp32/bf16 state × SD/DS layout, against the Triton path and an fp32
  reference, with null and padding slots;
• mixed steps (peeled decode rows; prefill rows with the Triton or FlashKDA prefill);
• spec-verify rows, pure and mixed with prefills, at 16/32/64 heads;
• CUDA-graph capture and replay of a decode step.

tests/models/glm5next/test_kda_flashinfer_prefill.py (new):

• FlashInfer prefill vs the Triton chunk path and an fp32 reference, with fresh and continued sequences,
  1-token and >16-token sequences;
• untouched slots; one-token steps;
• FlashInfer decode + prefill in one step;
• prefill rows alongside spec-verify rows;
• a step without draft rows in an MTP deployment, whose slot indices are a strided block-table column.

Existing: tests/models/glm5next/ (one test updated for the _forward signature) and
tests/models/kimi_k3/test_kda.py.

Where	FlashInfer	Result
VR200 (SM107), nightly c50aca1 + this branch	0.7.2 + final #7132 files + #6201 prefill	180 passed, 18 skipped (ROCm-only and Kimi K3-gated), 0 failed; all 50 FlashInfer KDA tests ran
GB300 (SM103), nightly c50aca1 + this branch	0.7.2 + final #7132 files	183 passed, 15 skipped, 0 failed; all 50 FlashInfer KDA tests ran

pytest -v tests/models/glm5next tests/models/kimi_k3/test_kda.py

The strided-slot-index test fails without the prefill fix (VR200, the same test file run against the prefill commit
without .contiguous()) and passes with it:

$ pytest tests/models/glm5next/test_kda_flashinfer_prefill.py -k block_table_column
test_block_table_column_state_indices[prefill] FAILED
test_block_table_column_state_indices[mixed] FAILED
E   ValueError: backend='cute-dsl' does not support this recurrent_kda prefill contract
# with this branch
test_block_table_column_state_indices[prefill] PASSED
test_block_table_column_state_indices[mixed] PASSED

Accuracy (GSM8K 40-shot, 1319 questions, greedy, 6 runs per arm, agg DEP2, mnbt 32768)

Arm	Runs	Mean ± sd
base (nightly 8e2beaa)	92.42, 91.51, 91.51, 92.04, 91.05, 92.57	91.85 ± 0.59
this PR, decode + prefill on FlashInfer	90.67, 91.36, 92.49, 92.72, 92.12, 91.89	91.88 ± 0.76

Greedy runs differ by about ±1 point from batch-dependent numerics, so 6 runs per arm were taken; the difference is
+0.03 points (t = 0.06). Reference from earlier VR200 work: 91.40%.

MTP (speculative decoding, 2 draft tokens)

GSM8K 40-shot as above with --speculative-config '{"method":"mtp","num_speculative_tokens":2}', agg DEP2, VR200,
nightly c50aca1. Covers the draft-verify path (packed_fused_kda_decode) and prefill/decode with spec-decoding
metadata.

Arm	Runs	Mean	Draft acceptance (accepted / drafted tokens, both runs)
base (nightly c50aca1)	92.27, 91.89	92.08	88.6% (162475 / 183370)
this PR, decode + prefill on FlashInfer	91.21, 91.89	91.55	88.7% (160379 / 180818)

Accuracy is within the ±1-point spread of greedy runs seen above, and draft acceptance is unchanged. The PR arm's
logs show the FlashInfer prefill and the packed decode (3 tokens per row) in use.

Serving with MTP (sa-bench 8K/1K random, agg DEP2, C128 and C256, VR200) also completed every request in both arms,
with the same draft acceptance.

This check found a bug in an earlier revision: under speculative decoding the builder's non-spec slot indices are a
strided column of the [N, 1 + num_spec] block table, which FlashInfer's prefill rejects. Warmup failed with
ValueError: backend='cute-dsl' does not support this recurrent_kda prefill contract. The prefill now passes
contiguous indices, and a unit test covers it.

Serving (GB300 and GB200)

Native srt-slurm sa-bench, ISL 8192 / OSL 1024, random data, Dynamo frontend, on Lyris (GB300 and GB200, 4 GPUs per
node). Settings:

• max-num-batched-tokens 32768;
• fp8 KV;
• FlashInfer TRT-LLM NVFP4 MoE;
• SD conv layout for agg, DS for disagg.

Base: nightly 8e2beaa. PR: the same image + this branch's kda.py + the final flashinfer-ai/flashinfer#7132
files, with --kda-prefill-backend flashinfer. Single runs; all points completed every request. The same recipes
were also run on VR200.

GPU	Config	C	Base TPS/GPU	PR TPS/GPU	Δ	Base tok/s/user	PR tok/s/user	Base TPOT ms	PR TPOT ms
GB300	agg DEP2 (2 GPU)	512	2847	3237	+13.7%	11.13	12.66	82.75	71.48
GB300	agg DEP2 (2 GPU)	1024	2402	2999	+24.9%	4.76	5.93	158.91	133.83
GB300	agg TP4 (4 GPU)	1	54.7	55.4	+1.4%	218.66	221.66	4.44	4.36
GB300	agg TP4 (4 GPU)	8	299.4	307.3	+2.6%	149.73	153.68	6.10	5.88
GB300	agg TP4 (4 GPU)	32	678.4	692.0	+2.0%	84.84	86.53	10.36	10.10
GB300	4P(TP1)+1D(DEP4) (8 GPU), DS layout	512	2657	2793	+5.1%	43.02	45.27	20.74	19.18
GB200	agg DEP2 (2 GPU)	512	2090	2372	+13.5%	8.41	9.52	61.91	64.22
GB200	agg DEP2 (2 GPU)	1024	2100	2410	+14.8%	4.26	4.87	62.30	64.05
GB200	agg TP4 (4 GPU)	1	53.1	54.7	+3.0%	212.47	218.90	4.57	4.41
GB200	agg TP4 (4 GPU)	8	294.7	302.8	+2.8%	147.38	151.45	6.16	5.94
GB200	agg TP4 (4 GPU)	32	671.6	683.3	+1.7%	83.98	85.44	10.45	10.24
GB200	2P(DEP2)+1D(DEP4) (8 GPU), DS layout	512	2199	2330	+6.0%	35.73	38.03	18.58	17.63

• DEP2 at C1024 with fp32 KDA state is capacity-bound: the state takes about 142 MB per sequence, so requests
  queue and every step is a mixed step. That is where the decode-row peel pays most.
• TP4 at low concurrency is decode-bound; KDA is only part of each step, so the per-step gain is a few percent.
• GB200 agg DEP2 is KV-capacity bound (about 186 GB per GPU): TPOT stays at 62–64 ms from C512 to C1024 and
  requests queue (mean TTFT 58 s for base at C512). The PR's gain there is admission: faster mixed steps cut mean TTFT
  58 → 42 s (C512) and 176 → 145 s (C1024), while per-token decode time is slightly higher with the larger running
  batch. Hence +14.8% at C1024 rather than about +25%.
• GB200 disagg shape: one GB200 cannot hold the 191 GB checkpoint at TP1, so each prefill worker is DEP2 (TP1 per
  rank, experts over 2 GPUs). A TP2 prefill does not work with a TP1-per-rank decode for this model: the hybrid KDA
  block size follows the per-rank KDA state, and NIXL rejects the mismatched blocks (8576 vs 4352).
• TTFT at TP4, C1–C8 is 20–47 ms higher in the PR arm (GB300 C1 143 → 163 ms, C8 603 → 650 ms; GB200 C1
  147 → 169 ms, C8 652 → 681 ms). These steps prefill one or a few sequences at 16 heads. On GB300 the FlashInfer
  prefill call costs about 190 µs of host time vs FlashKDA's 50–56 µs (34 layers per step), and at 1–4 sequences its
  kernel is also 7–13% slower; this is why the FlashInfer prefill stays opt-in. Decode alone was not run on GB300/GB200.

Details

Decode.

• Plain decode steps call flashinfer_fused_kda_decode once per layer: conv + SiLU with the conv-state update,
  the recurrence, and the gated RMSNorm, writing the layer output buffer. Slot ≤ 0 rows are null rows.
• Draft-verify steps (MTP) call flashinfer_packed_fused_kda_decode with the existing
  spec_state_indices_tensor / num_accepted_tokens. Mixed steps with verify rows keep the Triton spec path.
• Mixed steps without spec rows run the leading decode rows (vLLM orders decodes first) through the fused
  kernel. The conv then skips those rows as null slots, so the builder's existing causal-conv metadata stays
  valid. The prefill kernel uses the builder's prefill tail (prefill_query_start_loc / prefill_state_indices
  / prefill_has_initial_state), as the GDN layer already does. FlashKDA's prepare kernel thus no longer scans
  hundreds of one-token sequences, and gather / scatter no longer move their states. With prefix-cache
  checkpoints the step keeps today's path.
• Selection: --kda-decode-backend (additional_config.kda_decode_backend) and
  additional_config.kda_spec_decode_backend take auto / flashinfer / triton.
  • As for Kimi K3, a static check decides (_is_flashinfer_kda_decode_supported): CUDA SM10x, D=128, conv width
    4, bf16 activations and conv state, fp32 or bf16 recurrent state, a bounded gate, and the head counts and
    layouts of FlashInfer with [Frontend] Support embeddings in the run_batch API #7132 (single-token rows: 8/12/16/24/32/48/64/96 heads, SD or DS; draft-verify rows:
    12–96 heads, SD only, at most 7 draft tokens).
  • auto falls back to Triton when the check fails; flashinfer raises.
  • The prefill opt-in uses the same pattern (_is_flashinfer_kda_prefill_supported: SM100 / SM103 / SM107,
    D=128, bf16 activations, fp32 or bf16 state, bounded gate) and raises at startup when it fails.
• Norm placement: the fused kernel applies o_norm, so with the FlashInfer backend _forward returns normed
  output for every step type. In PIECEWISE (mixed/prefill) steps, breakable CUDA graphs run _forward eagerly but the code after it is
  replayed as captured, so whether to norm must not depend on the step type.

Prefill (opt-in).

• Call: --kda-prefill-backend flashinfer (already a CLI choice, previously rejected for GLM-5.3) calls
  recurrent_kda(backend="cute-dsl") with the raw gate and beta logits.
• In place: it reads and writes the recurrent state pool through ssm_state_indices, so there is no gather,
  scatter or final-state buffer.
• Fresh sequences: FlashInfer has no has-initial-state flag, so their slots are zeroed first with one
  index_fill_; rows with state zero the unused null slot.
• Edge cases:
  • a step whose prefill rows total one token uses the recurrent kernel;
  • steps with verify rows write their own output buffer and need no FlashKDA workspace.
• Opt-in (as for Kimi K3), and it produces no prefix-cache checkpoints
  (kda_prefill_checkpoint_alignment is None).

Numerics. FlashInfer's kernels are not bit-identical to Triton/FlashKDA (different summation orders). Against an
fp32/fp64 reference they are as close or closer:

• decode error equals the Triton path's;
• FlashInfer prefill: 4.2e-3 vs FlashKDA 5.1e-3 (max abs error, H64, 8 x 8192 tokens).

GSM8K above is the end-to-end check.

Limitations / follow-ups.

• Requires feat(kda): fused KDA decode for 16/64 heads and the DS conv-state layout flashinfer-ai/flashinfer#7132 (and [Bug]: vLLM 0.5.1 tensor parallel 2 stuck with Mixtral-8x7B-Instruct-v0.1 #6201 for SM107 prefill); bump the FlashInfer pin to the first
  release that has it before this PR leaves draft.
• Host cost (eager PIECEWISE steps only; FULL-graph decode steps do not pay it): FlashInfer's validating decode
  wrapper costs more per call than its cached kernel launch; a cheaper FlashInfer launch path would help low
  concurrency.
• Not in this PR: a faster prefill conv (planar / TMA output; overlaps [Perf][GLM-5.3-Flash] Update FlashKDA pin, consume row-strided QKV directly #61614), 64-row norm tiles,
  and a pipelined bf16-state decode kernel.
• CI: tests/models/glm5next runs on H200 (SM90), where these tests skip or exercise FlashInfer's kernels only
  where FlashInfer supports them. SM10x coverage needs a Blackwell lane.
• Follow-up: few-sequence prefill steps. With --kda-prefill-backend flashinfer, TTFT at low concurrency is higher
  (GB300 TP4: C1 143 → 163 ms, C8 603 → 650 ms). Most of the cost is the prefill call's host time, paid 34 times per
  eager prefill step: on GB300 about 190 µs on FlashInfer's two-kernel route (2 x sequences x heads <= SMs) and
  110 µs on its one-kernel route, vs FlashKDA's 50–56 µs; at 1–4 sequences the kernel is also 7–13% slower than
  FlashKDA's. A prototype that routes such steps to FlashKDA brings C1 TTFT back to base, with unchanged throughput,
  but reserves FlashKDA's workspace (less KV cache) and ties the cutover to FlashInfer's routing rule. It is left out
  of this PR; the host overhead is reported to FlashInfer, whose fix would make it unnecessary.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR (Claude Code); the submitter reviewed every line.

• Design Fit

• Testing and Validation

• Code Quality and Style

• Pull Request Contents

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @suiyoubi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @suiyoubi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork