[ROCm] Build per-sequence indexer decode lengths once per step

[amd-sriram]

Purpose

Alternative to vllm-project#59420 for review comparison. On ROCm the sparse indexer passes the (B, next_n) per-row seq_lens table to the AITER paged MQA-logits kernels, which read one length per sequence, so with native MTP decode they get the wrong lengths. vllm-project#59420 fixes this by copying the table's last column inside rocm_fp8_paged_mqa_logits, which runs once per indexer call (every layer).

Here the indexer metadata builder makes that copy once per step instead, into a persistent buffer, and the AITER calls read it. No AITER change and no per-layer copy.

Details

• DeepSeekV32IndexerDecodeMetadata.per_req_seq_lens: (B,) last column of seq_lens. Set on ROCm only, and only when seq_lens is (B, n > 1) (native MTP); otherwise None.
• Filled after DCP localization and compress-ratio scaling, so it always matches seq_lens. Persistent buffer, so the address is stable for CUDA graphs.
• rocm_fp8_paged_mqa_logits(..., per_req_context_lens=None): used for the AITER Gluon and stage1 calls when given. The C4A Triton and torch paths keep reading the 2D table, which they handle per row.
• Both callers pass it: rocm_aiter_mla_sparse.py and glm5next/amd/sparse_indexer.py.
• CUDA unchanged (DeepGEMM consumes the 2D table).

Test Plan

• pytest tests/v1/attention/test_rocm_indexer_per_req_seq_lens.py (new): the AITER Gluon (gfx950) and stage1 calls get per_req_context_lens when given and seq_lens otherwise.
• MI355X, real tensors through rocm_fp8_paged_mqa_logits against fp8_paged_mqa_logits_torch, B=16, next_n=3, varied lengths, the builder's (B, next_n) table and its last column.

Test Result

• Unit test: 4/4 pass on MI355X.
• Real tensors: 2D table only (main today): 48/48 rows with the wrong -inf mask. With per_req_context_lens: 0/48.
• Not run: the full builder path under a serving run (no unit test covers build() on ROCm with MTP), and e2e perf.

[coderabbitai]

Important

• 🔍 Trigger review

This repository does not receive automatic reviews because it has fewer than 10 stars.

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: d82c9d8a-ae63-470b-a33c-be8fbb2491f0

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}