[JaredforReal]

Purpose

When every context in a decode batch fits in index_topk (2048 tokens), the GLM-5.3-Flash kpool indexer skips top-k and fills each decode row with the exact causal indices (j for j <= position, else -1). _fill_causal_indices did this with torch ops: arange, cast, broadcast copy, compare and masked fill, i.e. 5 kernels per call, on both NVIDIA and AMD.

This PR replaces them with one Triton kernel on GPU. CPU tensors keep the torch code. The output is unchanged.

The other changes this PR used to carry are gone:

• the --mamba-ssm-cache-dtype fix landed as [Bugfix][GLM-5.3-Flash] Honor --mamba-ssm-cache-dtype for the KDA recurrent state vllm-project/vllm#59147;
• the batched f_b/g_b GEMM is covered by [Perf][GLM-5.3-Flash] Batch paired KDA gate projections (f_b_proj/g_b_proj) into one bmm vllm-project/vllm#57617.

Performance

Microbenchmark on H200, 2048-wide rows. Each call is replayed 78 times in one CUDA graph, as in the decode graph, and the table shows GPU time per call:

decode rows	torch (5 kernels)	this PR (1 kernel)
1	5.39 us	0.88 us
8	7.14 us	1.10 us
64	7.40 us	1.22 us
256	9.36 us	1.57 us

In eager mode both are host-bound at ~33 us per call.

End-to-end serving numbers on current main: TODO.

Test Plan

• tests/v1/attention/test_sparse_indexer_decode_seq_lens.py. The new GPU test compares the Triton fill with the torch fill for widths 8, 1000 and 2176, including positions past the row width.

Test Result

• H200: 10 passed.
• ROCm: not tested yet. The kernel also runs there, because rows.is_cuda is true on ROCm.

Why this is not a duplicate

_fill_causal_indices on upstream main (vllm/models/glm5next/common/sparse_indexer.py) is still the torch version. No open upstream PR touches it.

AI assistance (Claude) was used. The submitter reviewed every changed line and ran the tests above.

---

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

The KDA fused projection currently bypasses LoRA/quantization logic by using raw weights, and the new Triton fast path needs an explicit HAS_TRITON guard to avoid runtime failures when Triton is disabled.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR targets GLM-5.3-Flash MTP decode performance by reducing kernel-launch overhead in two hotspots and fixes a configuration plumbing issue so --mamba-ssm-cache-dtype is honored for the KDA recurrent state.

Changes:

• Plumb cache_config.mamba_ssm_cache_dtype into GLM5Next KDA state dtype selection.
• Reduce decode hot-path launches by (a) introducing a Triton kernel to fill causal top-k rows in one launch and (b) batching adjacent f_b/g_b projections into a single strided-batched GEMM.
• Add a CUDA test ensuring the fused causal-fill matches the torch formulation.

File summaries

File	Description
vllm/models/glm5next/nvidia/kda.py	Pass through SSM cache dtype for KDA state; fuse f_b/g_b projection into one batched GEMM with a cached stacked weight.
vllm/model_executor/layers/sparse_attn_indexer_kpool.py	Add a Triton kernel fast-path to fill causal indices in one launch for CUDA tensors.
tests/v1/attention/test_sparse_indexer_decode_seq_lens.py	Add a CUDA test comparing the fused causal-fill with the torch implementation.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[JaredforReal]

Closing: superseded upstream.

• Commit 1 (honor --mamba-ssm-cache-dtype for KDA) landed as [Bugfix][GLM-5.3-Flash] Honor --mamba-ssm-cache-dtype for the KDA recurrent state vllm-project/vllm#59147.
• f_b/g_b batched GEMM duplicates the open [Perf][GLM-5.3-Flash] Batch paired KDA gate projections (f_b_proj/g_b_proj) into one bmm vllm-project/vllm#57617.
• The fused causal top-k fill is still not upstream (now vllm/models/glm5next/common/sparse_indexer.py after [Refactor][GLM-5.3-Flash] Move sparse_attn_indexer_kpool into the model folder and split AMD/NVIDIA vllm-project/vllm#55358); kept on branch perf/glm53-flash-mtp-decode-cleanups for a possible standalone re-measure.