[Bugfix][ROCm] Bound sparse indexer decode workspace at layer build time

[Fangzhou-Ai]

Overview

ROCm-only alternative to #58014 for #55132. The ROCm sparse-MLA indexer works out its decode-row bound when the layer is built and passes it to the ROCm ops, instead of reading the vLLM config during the profiling run. gpu_worker.py is not changed.

Claims

• The decode-logits workspace reserved during memory profiling is max_num_seqs * (1 + num_speculative_tokens) rows again, not max_num_batched_tokens rows. This covers the AITER FP8 indexer (rocm_aiter_sparse_attn_indexer) and the MXFP4 indexer (rocm_mxfp4_sparse_attn_indexer / rocm_mxfp4_sparse_mqa_indexer).
• DeepSeek-V4.1-Flash on 8x MI355X at 1M context gets about 32 GiB more KV cache per GPU.
• The CUDA path and shared worker / model-runner code are unchanged.

Validation

End-to-end on 8x MI355X (gfx950) with dummy weights:

VLLM_ROCM_USE_AITER=1 vllm serve deepseek-ai/DeepSeek-V4.1-Flash --load-format dummy -tp 8 \
  --max-model-len 1048576 --max-num-seqs 16 --max-num-batched-tokens 8192 \
  --attention-config '{"indexer_kv_dtype":"mxfp4"}'

	Available KV cache memory (per GPU)	GPU KV cache size
main (89439db727)	174.95 GiB	90,340,997 tokens
this PR	206.92 GiB	106,889,551 tokens

The 31.97 GiB difference matches the expected saving: (8192 − 96) rows × 1M columns × 4 B ≈ 31.6 GiB. With the fix, 16 concurrent completions (prompts of 301–4801 tokens, 32 output tokens each) all finished with no workspace errors. That confirms the smaller locked workspace still covers real decode steps.

Tests:

• pytest tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py -k profiling_bounds passes (new test). It builds RocmSparseAttnIndexer while the config is set, then runs the profiling branch with no current config and checks that the reserved rows are 16 * (1 + 5).
• pytest tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py tests/models/glm5next/test_sparse_indexer_topk_dispatch.py tests/v1/attention/test_rocm_glm5next_sparse.py tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py: 49 passed, 4 skipped, and 18 failed. The same 18 tests fail on main on this machine. The cause is the environment: AITER/FlyDSL can't write its JIT cache under a read-only site-packages. This PR changes nothing in those tests.
• pre-commit run passes.

This change only affects how much memory is reserved, not model outputs, so no accuracy eval was run.

Details

Root cause: Worker.determine_available_memory calls profile_run() without set_current_vllm_config. _max_decode_logits_rows() called get_current_vllm_config(), which raised. The except branch then quietly returned max_num_batched_tokens. The workspace is locked after profiling, so the oversized reservation stayed for the life of the engine and counted as non-KV memory.

#58014 fixes this by setting the config context around profile_run() in gpu_worker.py. That is a change to the core worker for a problem that only shows up in ROCm code. This PR keeps the fix on the ROCm side instead:

• max_decode_logits_rows() (in rocm_aiter_mla_sparse.py) reads the bound once, when the layer is constructed. The model is built inside set_current_vllm_config, so the config is available then. This follows how SparseAttnIndexer already handles its DCP scalars.
• The ROCm ops take a new max_decode_rows: int = 0 argument, and the profiling branch reserves min(num_tokens, max_decode_rows) rows. A value of 0 (bound unknown) keeps the old conservative behavior.
• SparseAttnIndexer (the shared layer) gets one new attribute, and only its forward_hip reads it.

Not in this PR: the GLM5-Next kpool profiling path (sparse_attn_indexer_kpool in models/glm5next/amd/sparse_indexer.py) has the same root cause with the opposite effect. With no config it reserves 0 decode-logits rows, so it under-reserves. That can be fixed the same way in a follow-up.

---

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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Fangzhou-Ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork