[ROCm][DSv4] Make AITER sparse MLA decode cudagraph-clean (follow-up to #40889)

[ChuanLi1101]

Purpose

Cudagraph-clean follow-up to #40889. The persistent-mode AITER sparse MLA decode kernel (mla_a8w8_qh16_qseqlen1_gqaratio16_lse_ps) is the only LSE-returning sparse attention path on MI355X / gfx950, but the original integration allocates fresh per-step indexing tensors (qo_indptr, kv_indptr, kv_indices, kv_last_page_lens, q_scale, kv_scale) and a fresh FP8 query / output buffer on every layer call. That is incompatible with HIP/CUDA-graph capture, so DSv4 currently has to run in eager mode on ROCm.

This PR makes the per-step decode path completely allocation-free and pointer-stable so the decode loop can be wrapped in a HIP graph once the model wires it up.

Changes

• AiterSparseScratch now owns every per-step buffer in addition to the AITER work-plan / reduce buffers:
• qo_indptr, kv_indptr, kv_indices_2d, kv_last_page_lens
• valid_mask, valid_lens, col_arange
• q_fp8 (FP8 query buffer), out_buf (BF16 output)
• q_scale, kv_scale (constant 1.0 tensors)
• _aiter_decode_one_scope rewrites all of those in-place every step (torch.lt(out=), copy_, masked_fill_, torch.cumsum(out=)).
• AiterSparseScratch.refresh_metadata() re-runs aiter.get_mla_metadata_v1 against the current kv_indptr every step and writes the new work plan into the same work_* / reduce_* slots. The persistent ASM kernel encodes per-batch lengths into that plan, so leaving it stale across steps with different kv lengths causes a GPU memory-access fault.
• rebuild() now only allocates buffers and stores static gqa/topk/dtype parameters; it no longer re-runs the metadata builder itself.

Net effect: across decode steps with the same (total_q, nhead, topk, d_qk, d_v, dtype, kvtype) key, every data_ptr() is stable, so a graph captured on step N can be replayed for any step N+k.

Testing

bench_remote/_unit_test_cudagraph.py (run inside chuali_glm51 on MI355X, GPU 5):

• Call 1 with lens=[3, 2] → success, scratch key set.
• Call 2 with the same lens → rebuild() skipped, all 9 tracked data_ptr()s identical to call 1, output bit-identical.
• Call 3 with lens=[4, 1] → all 9 data_ptr()s still identical, no fault, output differs by 2.39 max abs as expected because the valid kv set is different.
• Parity vs. the original ([ROCm] Add AITER-accelerated MLA decode for DeepSeek V4 on MI355X #40889) implementation on the same inputs: max abs diff = 0.000000.

AI-assisted contribution

This PR was prepared with AI assistance (Cursor agent). The submitter reviewed every changed line and ran the unit-test harness above on MI355X.

Stacking

Depends on #40889 (AITER-accelerated MLA decode for DSv4). All four commits in this PR's range will reduce to two cudagraph-only commits once #40889 lands; happy to rebase on request.

[mergify]

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

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

[gemini-code-assist]

Code Review

This pull request implements support for the DeepseekV4 model family, introducing Multi-Head Latent Attention (MLA) and Multi-Token Predictor (MTP) capabilities. The changes include several optimized CUDA and Triton kernels for fused operations such as RMSNorm, RoPE, and quantization, alongside hypercompressed (mHC) blocks implemented via TileLang. Review feedback identifies a critical performance bottleneck where the entire SWA cache is dequantized every decode step, which should be optimized to handle dequantization on-the-fly. Additionally, the implementation of _dequantize_blocked_k_cache is noted to be incompatible with CUDA graphs due to per-call tensor allocations. A potential runtime error due to a shape mismatch in the attention output copy operation was also flagged, with a recommendation to squeeze the sequence dimension.

[gemini-code-assist]

Dequantizing the entire SWA cache every decode step is extremely inefficient and will cause significant performance degradation as the KV cache grows. In vLLM V1, kv_cache typically represents the entire physical block pool. Dequantizing thousands of blocks (e.g., 64MB per layer for 1000 blocks) on every step will likely exceed the latency budget for real-time inference. The AITER kernel should ideally handle dequantization on-the-fly using the provided indices, or at least only dequantize the active sliding window region.

[gemini-code-assist]

Potential shape mismatch during copy_. aiter_sparse_attn_decode is called with a 4D query q.unsqueeze(1) (shape [num_tokens, 1, num_heads, head_dim]), so it likely returns a 4D tensor. However, output is a 3D tensor (shape [num_tokens, num_heads, head_dim]). This will cause a RuntimeError during the copy_ operation. You should explicitly squeeze the sequence dimension from the attention output.

Suggested change

	output.copy_(attn_out.to(output.dtype))
	output.copy_(attn_out.squeeze(1).to(output.dtype))

[gemini-code-assist]

This allocation of the result tensor on every call is not cudagraph-clean and contradicts the PR's objective of making the decode path allocation-free. Since the number of blocks in the KV cache is fixed after initialization in vLLM V1, this buffer should be pre-allocated (e.g., as part of AiterSparseScratch or as a persistent buffer in the layer) to ensure pointer stability and avoid runtime overhead during graph replay.

[Morxi]

Fail on 8 x MI325x

DOCKER_BUILDKIT=1 docker build \
    -f docker/Dockerfile.rocm \
    -t vllm-rocm-custom:latest \
    --build-arg ROCM_VERSION=7.2.2 \
    --build-arg PYTHON_VERSION=3.10 \
    .

docker run -it --rm \
    --name vllm-deepseek-v4-1m \
    --group-add=video \
    --cap-add=SYS_PTRACE \
    --security-opt seccomp=unconfined \
    --device /dev/kfd \
    --device /dev/dri \
    --ipc=host \
    --shm-size 32gb \
    -p 8422:8000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    vllm-rocm-custom:latest \
    --model deepseek-ai/DeepSeek-V4-Pro \
        --host 0.0.0.0 \
        --port 8000 \
        --dtype auto \
        --tensor-parallel-size 8\
        --trust-remote-code

(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/utils.py", line 647, in <ge
nexpr>
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     layer_fn(prefix=f"{prefix}.{idx}") for idx in range(start_layer, end_layer)
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v4.py", line 638, 
in <lambda>
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     lambda prefix: DeepseekV4DecoderLayer(
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]                    ^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v4.py", line 489, 
in __init__
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v4.py", line 235, 
in __init__
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     self.experts = SharedFusedMoE(
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]                    ^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/layer.py", line 5
23, in __init__
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     self.quant_method: FusedMoEMethodBase = _get_quant_method()
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]                                             ^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/layer.py", line 5
15, in _get_quant_method
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     quant_method = self.quant_config.get_quant_method(self, prefix)
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v4.py", line 158, 
in get_quant_method
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     return Mxfp4MoEMethod(layer.moe_config)
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/quantization/mxfp4.py", lin
e 471, in __init__
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py",
 line 503, in select_mxfp4_moe_backend
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879]     raise NotImplementedError(
(Worker_TP7 pid=820) ERROR 04-26 02:28:10 [multiproc_executor.py:879] NotImplementedError: No MXFP4 MoE backend supports the deployment configuration.

[ChuanLi1101]

Moving to draft, blocked on #40889 (and on validating MI355X baseline generation quality). See #40889 for context. Will flip back to ready once #40889 is ready and lm_eval is re-run.