[Bugfix][ROCm] Expand indexer block tables when kernel blocks span several storage blocks

[amd-dlimpus]

Overview

Fixes the sparse-attention indexer reading the wrong cache slots on GLM-5.3-Flash/ROCm once a sequence exceeds 2048 tokens (the dense-shortcut threshold). The indexer block table is in kernel-block units (1152 tokens), but the compressed indexer cache is paged in smaller storage blocks (128 tokens). The builder only converted the table when the kernel block was smaller than the storage block.

Claims

• Fixes corrupted sparse top-k selection beyond 2048 tokens for GLM-5.3-Flash on ROCm (MI355X), in both prefill and decode.
• Restores GPQA-Diamond accuracy from 80.8–85.4% (with degenerate, non-terminating generations) to 89.9–91.4%. This matches the pre-regression baseline.
• No behavior change when the kernel block size equals the storage block size or divides it (the existing path). That covers NVIDIA, DeepSeek-V3.2 and every model that doesn't set storage_block_size.

Validation

Unit tests (tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py):

• test_to_storage_block_table: CPU-only checks of the conversion helper. It covers no kernel block size, equal sizes, kernel smaller than storage (the existing collapse) and kernel larger than storage (the new expand).
• test_indexer_builder_expands_kernel_blocks_to_storage_blocks: builds the indexer metadata with a 1152-token kernel block over a 128-token / tokens_per_state=4 storage spec. It checks every prefill compressed slot and every decode block-table entry against slots computed independently.

python -m pytest -q tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
# with fix:    33 passed
# without fix: 5 failed (the new tests), 28 passed

End to end, GLM-5.3-Flash on MI355X, TP=2:

Decode-vs-prefill consistency. For each sequence, I generate greedily, then re-score the same tokens with a single prefill. "Gap" is the mean |logprob(decode) − logprob(prefill)| over the first 64 generated tokens. "Flips" is the share of positions where the argmax differs.

Build	Mean gap	Argmax flips
Older base (pre-regression)	0.074	6.1%
Current main, unfixed	0.53–0.69	23–25%
Current main + fix, bf16	0.075	5.3%
Current main + fix, MXFP4 MLA	0.069	4.3%

GPQA-Diamond (198 questions):

Build	Accuracy
Current main, unfixed (8 runs, bf16 + MXFP4)	80.81–85.35%, degenerate/runaway generations
Current main + fix, MXFP4, seeds 42/43/44	90.40 / 90.40 / 91.41%
Current main + fix, bf16, seed 42	89.90%
Older base, bf16 / MXFP4	88.89–91.41% / 90.40–91.92%

With the fix, wall time per GPQA run also drops from ~3000–3700 s to ~2000–2400 s, because generations no longer run away to the max token limit.

Details

Root cause. On ROCm, Glm5NextIndexerCache.get_kv_cache_spec sets storage_block_size = page_size * kpool = 128 tokens. A 1152-token hybrid-aligned block holds 288 kpool states, which isn't a multiple of 64, so it is stored as nine 32-state pages. The ROCm indexer backend accepts MultipleOf(16) kernel blocks, so set_kernel_block_size(1152) leaves the block table in 1152-token units. DeepseekV32IndexerMetadataBuilder only handled storage % kernel == 0 (the [:, ::f] // f collapse). Otherwise it passed the 1152-unit table through as if it held 128-token ids. Slots were computed as T[pos // 128] * 32 + (pos // 4) % 32, so pos // 128 indexed past the real table entries. Beyond the first few pages, states read zero padding and aliased onto null block 0, which every request shares. Below 2048 tokens the indexer takes a dense causal shortcut, so the bug only appears on long sequences.

Worked example, for a request whose block table is [5, 7, 2] (each physical block is 9 indexer pages):

token 1400
  fixed: 1400 // 1152 = entry 1 -> block 7 -> page 7*9 + (1400-1152)//128 = 63 + 1 = 64
  old:   1400 // 128  = 10 -> table[10] is past the 3 entries -> padding 0 -> page 0 (null block)

early tokens (pos < 384)
  old:   table[pos // 128] = 5, 7 or 2, used directly as a page id -> pages 5/7/2, also inside null block 0

So every request read and wrote its compressed keys in the same few null-block pages.

Fix. Add _to_storage_block_table, which re-expresses the kernel-block table in storage blocks. It keeps the existing collapse when the kernel block divides the storage block. When the storage block divides the kernel block, it expands kernel block k into k*f … k*f+f-1. Both call sites (the compressed slot mapping / prefill table, and the decode paged-MQA block table) use it. If neither size divides the other, it returns None and the table passes through unchanged, as before.

Scope. Only specs with storage_block_size set and smaller than the kernel block take the new path. Today that is GLM-5.3 on ROCm. On NVIDIA the indexer kernel block is 64 tokens, so the existing collapse branch applies. DeepSeek V3.2 uses compress_ratio = 1 without a storage split.

---

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.

[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-dlimpus]

Closing as a duplicate of #59412, which fixes the same root cause (#58858) by selecting page-aligned kernel blocks. I posted the end-to-end accuracy evidence from this PR there. #56381 and the KV-layout refactor in #59297 also overlap.