From e534024e490a32cc500555973e2291cb9de9b38f Mon Sep 17 00:00:00 2001 From: lucas-santos-amd Date: Wed, 14 Oct 2026 15:34:28 -0300 Subject: [PATCH] [ROCm][Bugfix][MLA] Fix FP8 MLA prefill unpad for interleave-padded heads (#60907) Signed-off-by: Lucas Santos Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../test_rocm_aiter_mla_fp8_prefill.py | 12 +++++++----- .../v1/attention/backends/mla/rocm_aiter_mla.py | 17 ++++++++++++----- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py index 23dd09128b..de48b195db 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index dfd5dfabb0..e485f5462a 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/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,