[ROCm][DSv3.2] Adopt new paged-MQA-logits API + cached logits buffer with defensive padding

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -17,6 +17,57 @@
     from vllm import _custom_ops as ops
 
 
+# Defense-in-depth padding for the paged-MQA-logits output buffer.  The
+# aiter gluon `_gluon_deepgemm_fp8_paged_mqa_logits_preshuffle[_varctx]`
+# kernels (as of ROCm/aiter@main, pre-#2866) contain unmasked buffer_store
+# ops on `OutLogits_buffer` that can overshoot the logical row by up to
+# `ChunkKPerStage = 128` float32 elements when
+# `context_length == max_model_len` because `split_context_length` is
+# rounded up to a multiple of `KVBlockSize`.  Over-allocating the row
+# stride by `_PAGED_LOGITS_ROW_PADDING` columns (kept >= ChunkKPerStage
+# for safety) turns those OOB writes into in-bounds writes to a never-
+# read padding region, avoiding HIP memory access faults on gfx950
+# (MI355X) during DeepSeek V3.2 MTP decoding.  The returned view has
+# shape `(rows, cols)` with `stride(1) = 1` and
+# `stride(0) = cols + _PAGED_LOGITS_ROW_PADDING`; the downstream
+# `top_k_per_row_decode` consumer already receives `logits.stride(0)`
+# and `logits.stride(1)` as explicit arguments, so this is transparent.
+# Once aiter#2866 is merged and released this padding can be reduced
+# to 0 safely.
+_PAGED_LOGITS_ROW_PADDING: int = 256
+_cached_paged_logits: torch.Tensor | None = None
+
+
+def _get_paged_logits_buffer(
+    rows: int, cols: int, device: torch.device
+) -> torch.Tensor:
+    """Return a (rows, cols) float32 buffer pre-filled with -inf for the
+    paged MQA-logits kernel to write into.
+
+    The underlying storage is over-allocated by `_PAGED_LOGITS_ROW_PADDING`
+    columns as defense-in-depth against the aiter preshuffle-kernel OOB
+    write (ROCm/aiter#2866).  Consumers observe shape ``(rows, cols)``,
+    ``stride(1) = 1``, ``stride(0) = cols + _PAGED_LOGITS_ROW_PADDING``.
+
+    The buffer is cached across decode steps and reused when the logical
+    shape and device match, saving a ``torch.full(-inf)`` per layer
+    (DSv3.2 has 61 layers per decode step).
+    """
+    global _cached_paged_logits
+    padded_cols = cols + _PAGED_LOGITS_ROW_PADDING
+    if (
+        _cached_paged_logits is not None
+        and _cached_paged_logits.shape[0] == rows
+        and _cached_paged_logits.shape[1] == padded_cols
+        and _cached_paged_logits.device == device
+    ):
+        return _cached_paged_logits[:, :cols]
+    _cached_paged_logits = torch.full(
+        (rows, padded_cols), float("-inf"), device=device, dtype=torch.float32
+    )
+    return _cached_paged_logits[:, :cols]
+
+
 @triton.jit
 def _indexer_k_quant_and_cache_kernel(
     k_ptr,  # [num_tokens, head_dim]
@@ -300,6 +351,7 @@ def rocm_fp8_paged_mqa_logits(
     block_tables: torch.Tensor,
     schedule_metadata: torch.Tensor,
     max_model_len: int,
+    block_size: int = 1,
 ) -> torch.Tensor:
     """Compute FP8 MQA logits using paged KV-cache.
 
@@ -317,6 +369,12 @@ def rocm_fp8_paged_mqa_logits(
         schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
             used to distribute work across SMs.
         max_model_len: Maximum sequence length used to size the logits output.
+        block_size: KV cache block size.  When > 1 and the installed aiter
+            build exposes the newer ``deepgemm_fp8_paged_mqa_logits`` API,
+            the fused preshuffle kernel is used (required for DeepSeek V3.2
+            decode on gfx950).  Defaults to 1, preserving the legacy
+            ``deepgemm_fp8_paged_mqa_logits_stage1`` + ``out_qk.sum(dim=0)``
+            code path.
 
     Returns:
         Logits tensor of shape [B * next_n, max_model_len], dtype
@@ -329,10 +387,36 @@ def rocm_fp8_paged_mqa_logits(
         aiter_paged_mqa_logits_module = paged_mqa_logits_module()
 
     if aiter_paged_mqa_logits_module is not None:
+        batch_size, next_n, heads, _ = q_fp8.shape
+        # Prefer the newer fused `deepgemm_fp8_paged_mqa_logits` API when
+        # the aiter build exposes it AND `block_size > 1` (required by the
+        # MFMA-shape preshuffle kernel).  Fall back to `_stage1` otherwise.
+        _deepgemm_fp8_paged_mqa_logits = getattr(
+            aiter_paged_mqa_logits_module,
+            "deepgemm_fp8_paged_mqa_logits",
+            None,
+        )
+        if _deepgemm_fp8_paged_mqa_logits is not None and block_size > 1:
+            out_logits = _get_paged_logits_buffer(
+                batch_size * next_n, max_model_len, q_fp8.device
+            )
+            _deepgemm_fp8_paged_mqa_logits(
+                q_fp8,
+                kv_cache_fp8,
+                weights,
+                out_logits,
+                context_lens,
+                block_tables,
+                max_model_len,
+                ChunkK=256,
+                Preshuffle=(block_size == 64),
+                KVBlockSize=block_size,
+                WavePerEU=2,
+            )
+            return out_logits
         deepgemm_fp8_paged_mqa_logits_stage1 = (
             aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
         )
-        batch_size, next_n, heads, _ = q_fp8.shape
         out_qk = torch.full(
             (heads, batch_size * next_n, max_model_len),
             float("-inf"),
@@ -625,6 +709,7 @@ def rocm_aiter_sparse_attn_indexer(
             decode_metadata.block_table,
             decode_metadata.schedule_metadata,
             max_model_len=max_model_len,
+            block_size=kv_cache.shape[1],
         )
 
         num_rows = logits.shape[0]