[zhentaocc]

Overview

GLM-5.3-Flash's chunk KDA prefill recomputes FLA's chunk_indices with prepare_chunk_indices. That function reads cu_seqlens back to the host, which is a GPU→CPU sync on every prefill call. GDNAttentionMetadata already builds chunk_indices and chunk_offsets on the host, so this PR passes them through instead.

Related:

• [Perf] Avoid a per-prefill-step GPU->CPU sync in the KDA chunk kernels #51540: the same fix for Kimi Linear. Its non_spec_chunk_indices would also cover the case this PR leaves on the old path.
• [RFC]: DFlash2 speculative decoding performance on ROCm (MI355X) #61761: DFlash2 performance on ROCm.

Claims

• No sync from prepare_chunk_indices in spec batches, or in non-spec batches without decodes.
• Output is bit-identical.
• End-to-end throughput is unchanged on MI355X. The sync was not on the critical path there, so this removes a host sync, not GPU time.

Validation

Test. tests/models/glm5next/test_kda_chunk_metadata.py (new):

• With the builder's metadata, the kernel derives none of its own and its output is bit-identical.
• 3 passed with this PR. On main all 3 fail.

End to end. GLM-5.3-Flash + DFlash2, MI355X TP2, with #63070, max_cudagraph_capture_size 2048, #63071 and #63069. Output tok/s (acceptance length):

	c=1	c=16	c=64	c=64 step
without this PR	388.0 (4.64)	2215.2 (4.48)	4554.1 (4.44)	56.9 ms
with this PR	409.1 (4.95)	2168.9 (4.47)	4589.8 (4.63)	57.2 ms

• Before this PR, py-spy put about 20% of each worker's main-thread samples in prepare_chunk_indices (.tolist()), at c=64.
• The host was already ahead of the GPU at that point, so the wait did not delay the GPU.
• GSM8K with this PR is 0.9704.

Setup and commands

One MI355X node (gfx950). TP2. Image vllm/vllm-openai-rocm:nightly-rocm100-e68efc685a, with the patches named in each row applied on top. VLLM_ROCM_USE_AITER=1 VLLM_USE_V2_MODEL_RUNNER=1.

vllm serve amd/GLM-5.3-Flash-Quark-MXFP4 -tp 2 --max-model-len 16384 --max-num-seqs 256 \
  --gpu-memory-utilization 0.85 --limit-mm-per-prompt '{"image": 0, "video": 0}' \
  --no-enable-prefix-caching \
  [--compilation-config '{"max_cudagraph_capture_size": 2048}'] \
  --speculative-config '{"method": "dflash", "model": "incoai/GLM-5.3-Flash-DFlash2", "num_speculative_tokens": 7}'
vllm bench serve --dataset-name sharegpt --sharegpt-output-len 512 --temperature 0 --seed 0 \
  --num-prompts max(64, 8*C) --max-concurrency C

• Acceptance length comes from the vllm:spec_decode_* counters, read before and after each concurrency.
• GSM8K uses all 1319 test questions: zero-shot chat, greedy, thinking on, up to 6144 output tokens, concurrency 64. The score is the last number after </think>.
• Each row is one run. Acceptance length varies by about ±0.15 between runs, so the step time is given as TPOT × acceptance length, which is one spec step.

Details

• chunk_kda_with_fused_gate takes optional chunk_indices and chunk_offsets, and derives them only when they are not given. They are passed on to solve_tril and chunk_gated_delta_rule_fwd_h.
• The builder's metadata describes non_spec_query_start_loc in every spec batch, and in non-spec batches without decodes.
• A non-spec batch with decodes keeps the old path. The builder builds that batch's metadata prefill-only, because GDN peels decodes off, while KDA chunks the decodes together with the prefills.
• The NVIDIA kernels get the same plumbing, since the model code is shared. I did not run them on NVIDIA.

---

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.

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

🚀

[zhentaocc]

no improvement. closed