[ROCm][Bugfix][MLA] Fix FP8 MLA prefill unpad for interleave-padded heads (#60907)

Signed-off-by: Lucas Santos <Lucas.Santos@amd.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py

@@ -47,16 +47,18 @@ pytestmark = pytest.mark.skipif(
 #
 # The kernel requires 16-aligned heads, so _mla_fp8_prefill_attn replicate-pads
 # q/k/v up to get_fp8_prefill_num_heads(num_heads) and slices the output back.
-# Cover all three regimes:
+# Cover all four regimes:
+#    8 -> padded to 16 (divisor, repeat_interleave), so the real heads sit at
+#         stride 2 in the padded output. Live case: Kimi-K2.x at TP8.
 #   12 -> padded to 16 (non-divisor, tile+slice). Live case: a K3 rank at TP8.
 #   16 -> no padding, output aliases the caller's buffer.
 #   24 -> padded to 32. No current model and TP reaches this band (96 heads
 #         would need TP4, which does not fit in 288 GiB; 128-head models give
 #         128/64/32/16/8), so this case exists to keep the general path honest
-#         for future architectures: it is the only one of the three where
-#         padding above 16 actually happens, so it is what proves the
-#         replicate-pad/slice-back argument holds for a target other than 16.
-HEAD_COUNTS = [12, 16, 24]
+#         for future architectures: it is the only one where padding above
+#         16 actually happens, so it is what proves the replicate-pad/slice-back
+#         argument holds for a target other than 16.
+HEAD_COUNTS = [8, 12, 16, 24]
 QK_NOPE_HEAD_DIM = 128
 QK_ROPE_HEAD_DIM = 64
 QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM  # 192

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -2317,10 +2317,12 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
         #
         # Exact, not approximate: after kv_b_proj gqa_ratio is 1, so q, k and v
         # all carry num_heads heads and attention is independent per head.
-        # Padding all three identically makes padded head j a duplicate of real
-        # head j % num_heads, so the real heads [0:num_heads] are bit-identical
-        # to the unpadded result. Same argument as the decode path; only the
-        # target width differs.
+        # Padding all three identically makes every padded head a duplicate of
+        # a real head, so the real heads are bit-identical to the unpadded
+        # result. get_mla_padded_q repeat_interleaves when the target is a
+        # multiple of num_heads (real heads at stride target // num_heads) and
+        # tiles otherwise (real heads [0:num_heads]). Same argument as the
+        # decode path; only the target width differs.
         _real_nhead = self.num_heads
         nhead = AiterMLAHelper.get_fp8_prefill_num_heads(_real_nhead)
         _pad = nhead != _real_nhead
@@ -2408,7 +2410,12 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
         )
 
         if _pad:
-            out.view(total_q, _real_nhead, v_head_dim).copy_(out_3d[:, :_real_nhead, :])
+            real_out = (
+                out_3d[:, :: nhead // _real_nhead, :]
+                if nhead % _real_nhead == 0
+                else out_3d[:, :_real_nhead, :]
+            )
+            out.view(total_q, _real_nhead, v_head_dim).copy_(real_out)
 
     def forward_mha(
         self,