[ROCm][DSv4] Tighten KV-workspace zero comment + FP8 encoding notes (review #42893)

Replace the long rationale around ``kv.zero_()`` (in both prefill paths) with a brief TODO that names the proper fix: mask invalid rows in the indexer (score = -inf) or in the sparse-attention kernel (skip indices >= valid_len). The current zero is the minimal interim workaround; the underlying bug is arch-independent (uninitialized workspace + indexer that scores the entire M dim) so the call stays unchanged on every platform until the indexer/kernel fix lands. No behavior change.

Also condense the duplicate FNUZ-vs-OCP comments at the dequant call sites and in ``_sparse_attn_decode_ragged_kernel``: the wrapper docstring already explains the asymmetry, so per-call-site repetition was just noise.

Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: maeehart

vllm/models/deepseek_v4/amd/rocm.py

@@ -790,17 +790,10 @@ def _forward_prefill(
         kv = workspace_manager.get_simultaneous(
             ((cls.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
         )[0]
-        # The workspace allocator returns uninitialized memory and is shared
-        # across requests + other layers. dequantize_and_gather_k_cache only
-        # writes the compressed-K prefix (rows [0, seq_len/compress_ratio))
-        # and the SWA window (rows [N, N+gather_lens)) for each chunk row,
-        # leaving holes in the M dimension. rocm_sparse_attn_prefill then
-        # reads ``kv.view(-1, 1, head_dim)`` via ragged indices which can
-        # reach those holes for very short sequences, causing data-dependent
-        # non-determinism across otherwise-identical temperature=0 requests.
-        # Zero once per call so the holes are deterministic (zero attention
-        # contribution). The cost is one bf16 fill of the workspace tile,
-        # which is dwarfed by the FP8 dequant + sparse attention themselves.
+        # TODO: workspace is torch.empty() and only the compressed-K prefix +
+        # SWA window are written per chunk row; the indexer's topK can land in
+        # the unwritten holes for short sequences. Proper fix is to mask invalid
+        # rows in the indexer (score = -inf) or in rocm_sparse_attn_prefill.
         kv.zero_()
         for chunk_idx in range(num_chunks):
             chunk_start = chunk_idx * cls.PREFILL_CHUNK_SIZE
@@ -810,12 +803,7 @@ def _forward_prefill(
                 assert attn_metadata is not None
                 assert compressed_k_cache is not None
                 block_table = attn_metadata.block_table[num_decodes:]
-                # The compressed-K encoder (Triton _fused_kv_compress_norm_...
-                # in fused_compress_quant_cache.py) writes bytes via
-                # tl.float8e4nv with FP8_MAX=448.0 regardless of platform.
-                # The SWA-side C++ encoder, by contrast, switches to FNUZ on
-                # gfx942 (PR #42893), so the two caches need different
-                # use_fnuz settings even on the same MI300X.
+                # compressed_k_cache is OCP on every platform (Triton encoder).
                 dequantize_and_gather_k_cache(
                     kv[:chunk_size],
                     compressed_k_cache,

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -565,16 +565,11 @@ def rocm_fp8_mqa_logits(
 
     k_fp8, scale = kv
 
-    # gfx942 (MI300X): the AITER ``fp8_mqa_logits`` wrapper bundled in the
-    # currently-pinned aiter wheel launches its Triton kernel with
-    # ``(BLOCK_KV=128, num_stages=2)``, which requests ~96 KiB of LDS for
-    # the DSv4 sparse indexer shape. MI300X CUs have 64 KiB of LDS, so
-    # the launch JIT-aborts with ``OutOfResources: shared memory`` on the
-    # first inference. Route gfx942 callers to a vLLM-vendored copy of
-    # the same kernel that selects ``(BLOCK_KV=64, num_stages=1)`` when
-    # the default tile doesn't fit (~33 KiB), matching the fix in
-    # ROCm/aiter#3257. This entire branch can be removed once vLLM bumps
-    # to an AITER version that includes that PR.
+    # gfx942: AITER's bundled fp8_mqa_logits launches with BLOCK_KV=128 +
+    # num_stages=2 (~96 KiB LDS), exceeding MI300X's 64 KiB LDS so it aborts
+    # with OutOfResources. Route gfx942 to a vendored copy that drops to
+    # BLOCK_KV=64 + num_stages=1 (~33 KiB) per ROCm/aiter#3257. Remove this
+    # branch once vLLM bumps AITER to a version that includes that PR.
     if _ON_GFX942 and rocm_aiter_ops.is_enabled():
         from vllm.v1.attention.ops.triton_fp8_mqa_logits import (
             fp8_mqa_logits_gfx942,
@@ -1193,10 +1188,8 @@ def _sparse_attn_decode_ragged_kernel(
     NOPE_DIM: tl.constexpr,
     NOPE_BLOCK: tl.constexpr,
     ROPE_DIM: tl.constexpr,
-    # `main_cache` is the SWA K-cache (written by the C++ encoder, FNUZ on
-    # gfx942 / OCP on gfx950). `extra_cache` is the compressed K-cache
-    # (Triton encoder, OCP on every platform). Reading both with the same
-    # `IS_FNUZ` would mis-decode one of them by the FNUZ/OCP scale ratio.
+    # SWA K-cache (main): C++ encoder writes FNUZ on gfx942, OCP on gfx950.
+    # Compressed K-cache (extra): Triton encoder writes OCP everywhere.
     IS_FNUZ_MAIN: tl.constexpr,
     IS_FNUZ_EXTRA: tl.constexpr,
     BLOCK_H: tl.constexpr,
@@ -1601,10 +1594,6 @@ def _rocm_sparse_attn_decode_ragged_triton(
         NOPE_DIM=nope_head_dim,
         NOPE_BLOCK=triton.next_power_of_2(nope_head_dim),
         ROPE_DIM=rope_head_dim,
-        # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950).
-        # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere).
-        # Reading both with a single IS_FNUZ would mis-decode one of them by
-        # the FNUZ/OCP scale ratio (~1.87×).
         IS_FNUZ_MAIN=current_platform.is_fp8_fnuz(),
         IS_FNUZ_EXTRA=False,
         BLOCK_H=block_h,