[Kernel][Perf] Skip GLM-5.3 sparse indexing for short contexts

[zigzagcai]

[Kernel][Perf] Skip GLM-5.3 sparse indexing for short contexts

Purpose

GLM-5.3-Flash uses a 2,048-token sparse-attention index budget. When the full
causal context is at most 2,048 tokens, sparse scoring selects every token and
the paged-MQA logits, top-k, and pool-expansion work is unnecessary.

This change:

• enables the existing exact short-context shortcut on NVIDIA CUDA as well as
  ROCm;
• replaces the NVIDIA shortcut's five eager PyTorch CUDA kernels with one
  stride-aware Triton causal-index fill; and
• keeps the existing PyTorch implementation on CPU, XPU, and ROCm.

This is a three-file performance delta stacked on vllm-project#53906. It targets
ZJY0516/vllm:glm-release, whose head is the exact parent of this PR at
36bb3795b258e1b773cee5c2b725d9b8346b0c8d. It can therefore be reviewed and
tested independently while vllm-project#53906 remains open. After vllm-project#53906 lands, this
commit can be rebased and retargeted to upstream main without carrying the
model-support diff.

Trace and kernel evidence

The baseline GLM sparse-indexer trace contains paged-MQA logits, persistent
top-k, pool/tail expansion, and metadata kernels even when every causal token
fits in the 2,048-token budget. A focused torch profile shows that the legacy
causal fill itself launches five CUDA kernels; the candidate launches one.

H200, PyTorch 2.13.0+cu130, CUPTI 13.3.1, CUDA graph timing, cold L2, seeded
inputs, preallocated outputs, and alternating arm order:

Rows x columns	Legacy	Fused	Speedup
1 x 2,048	7.872 us	1.120 us	7.03x
96 x 2,048	9.056 us	1.185 us	7.64x
512 x 2,048	14.368 us	2.048 us	7.02x
1,536 x 2,048	26.368 us	3.968 us	6.65x

Reproduce with:

PYTHONPATH=. CUDA_VISIBLE_DEVICES=0 <active-env>/bin/python \
  benchmarks/kernels/benchmark_sparse_indexer_causal_fill.py

After rebasing to the current vllm-project#53906 head, the focused suite still reported 10
passed and mypy passed. A fresh H200 run measured 7.59x, 7.00x, 6.85x, and
6.60x speedups for 1, 96, 512, and 1,536 rows respectively.

Matched serving A/B

Configuration: 4x H200, TP4, MTP-5, BF16 KV cache, CUDA graphs, 8,192 max
batched tokens, 512 input / 256 output, concurrency 16, 32 requests,
temperature=0, and ignore_eos. Each arm used one excluded stabilizer and 10
measured runs. All 640 measured requests succeeded without illegal memory
access.

Metric	Baseline mean	Candidate mean	Change	Bootstrap 95% CI	Exact permutation p
Output throughput	1,433.943 tok/s	1,601.974 tok/s	+11.718%	[+5.729%, +17.600%]	0.00126
Mean TTFT	335.063 ms	323.122 ms	-3.564%	[-12.917%, +7.267%]	0.52008
Mean TPOT	8.287 ms	7.776 ms	-6.170%	[-11.981%, -0.049%]	0.08357
MTP acceptance	70.351%	75.023%	+4.672 pp	[+0.933 pp, +8.540 pp]	0.03366

MTP acceptance correlates with throughput in both arms. An OLS ablation that
controls for acceptance rate still estimates a +6.145% throughput effect
(p=0.00550). The adjusted TPOT effect is not significant, so throughput is
the primary claim and no independent TPOT improvement is claimed.

Correctness and model evaluation

The CUDA test covers widths below, equal to, and above one 2,048-element tile,
as well as non-contiguous row and position strides. CPU tests continue to cover
the exact short-context predicate and sparse-context fallback.

The shortcut changes floating-point reduction order by returning ascending
causal indices instead of relevance-ordered indices, so a model evaluation was
run rather than assuming output identity. Using vLLM's GSM8K evaluator with
200 questions, 5-shot prompting, greedy sampling, 256 max output tokens, and
concurrency 16:

Arm	Accuracy	Invalid responses
Baseline	0.925	0.000
Candidate	0.940	0.000

Commands run:

CUDA_VISIBLE_DEVICES=0 <active-env>/bin/python -m pytest \
  tests/v1/attention/test_sparse_indexer_decode_seq_lens.py -q
# 10 passed

pre-commit run ruff-format --files \
  benchmarks/kernels/benchmark_sparse_indexer_causal_fill.py \
  vllm/model_executor/layers/sparse_attn_indexer_kpool.py \
  tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
# passed

pre-commit run ruff-check --files \
  benchmarks/kernels/benchmark_sparse_indexer_causal_fill.py \
  vllm/model_executor/layers/sparse_attn_indexer_kpool.py \
  tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
# passed

pre-commit run mypy-3.12 --files \
  benchmarks/kernels/benchmark_sparse_indexer_causal_fill.py \
  vllm/model_executor/layers/sparse_attn_indexer_kpool.py \
  tests/v1/attention/test_sparse_indexer_decode_seq_lens.py \
  --hook-stage manual
# passed

git diff --check
# passed

Duplicate-work check

Open-PR searches were refreshed on 2026-08-30 in both
vllm-project/vllm and ZJY0516/vllm for GLM-5.3 sparse indexer,
short context sparse indexer, causal indices Triton, and
NVIDIA causal index fill. No exact duplicate was found.

There is no linked issue, so an issue-number lookup and the corresponding
<issue_number> in:body search are not applicable.

• [Model] add GLM-5.3-Flash support vllm-project/vllm#53906 is the required GLM-5.3 model-support base; it contains the ROCm
  shortcut but does not enable or fuse the NVIDIA path.
• [Perf][Attention][DSA] Shard prefill indexer rows across TP ranks vllm-project/vllm#54394 shards long-context CUDA prefill indexer rows across TP ranks. This PR
  removes short-context decode scoring and fuses causal-index materialization;
  the execution regimes and implementations are distinct.
• [Bugfix] Support NoPE models on FLASHINFER_MLA_SPARSE_SM120 and validate effective topk buffer width vllm-project/vllm#53969 fixes the effective top-k buffer width for an SM120 NoPE backend.
• [Attention] Overlap sparse MLA indexer on a native CUDA side stream vllm-project/vllm#47355 overlaps sparse indexer computation on a side stream.
• [Minimax-M3] Enable fp8 index cache on the Triton indexer (non-SM100) vllm-project/vllm#47665 enables an FP8 index cache for a different indexer path.
• [rocm] perf: drop redundant -inf prefill of decode paged MQA-logits buffer vllm-project/vllm#50008 and [ROCm][Perf][MiniMax-M3] Speed up the lightning indexer and add an opt-in fp8 index cache vllm-project/vllm#53448 optimize ROCm-specific sparse-indexer work.
• [Kernel][Perf] Fuse clamped MoE activation and UE8M0 FP8 block quantization vllm-project/vllm#53875 is a separate clamped MoE activation/quantization fusion.
• ZJY0516/vllm [Bugfix] Make GLM-5.3 kpool metadata graph-safe without prefix caching #7 fixes kpool graph-metadata lifetime, and [Bugfix] Reject invalid GLM kpool indices #8 rejects invalid
  kpool indices. [Bugfix] Reject invalid GLM kpool indices #8 touches the same implementation file but does not alter the
  short-context predicate or causal fill; both fixes are orthogonal.

AI assistance and human accountability

AI assistance was used for profiling, implementation, benchmarking, tests,
statistical analysis, and drafting this description.

• Before submission, the human submitter reviewed every changed line,
  understands the implementation and evidence, and can defend the change end
  to end.
• Before submission, the human submitter reran or directly reviewed the
  reported tests, serving A/B, and model evaluation artifacts.

The PR should be rebased and retargeted to upstream main after vllm-project#53906 lands.

[ZJY0516]

short context optimization is not very useful

[zigzagcai]

short context optimization is not very useful

Hi @ZJY0516 Thank you for your reply! I understand that short context scenarios are not very useful, but it just fills in the gaps in the benefits of short context scenarios on NV (I see there already exist short context path for AMD) and has some benefit in 2k context length.

[zigzagcai]

Update: I close the PR #11, since this optmization is limited to short context length, which is not very useful in real-world scenarios.