[ROCm][K3] Extend FP8 asm MLA prefill to non-divisor small head counts (#51040)

Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py

@@ -43,8 +43,20 @@ pytestmark = pytest.mark.skipif(
 )
 
 # DeepSeek-style MLA head dims after kv_b_proj decompression (q/k carry the
-# nope+rope dims, v carries v_head_dim). FP8 prefill requires 16-aligned heads.
-NUM_HEADS = 16
+# nope+rope dims, v carries v_head_dim).
+#
+# 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:
+#   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]
 QK_NOPE_HEAD_DIM = 128
 QK_ROPE_HEAD_DIM = 64
 QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM  # 192
@@ -56,6 +68,10 @@ SCALE = 1.0 / math.sqrt(QK_HEAD_DIM)
 # elements.
 ATOL, RTOL = 1e-1, 5e-2
 
+# The assembly kernel writes bf16 through a raw output pointer. Keep its tensors
+# and workspace reservation on one dtype.
+ATTN_OUT_DTYPE = torch.bfloat16
+
 
 @pytest.fixture(autouse=True)
 def _workspace_manager():
@@ -73,14 +89,14 @@ def _workspace_manager():
     reset_workspace_manager()
 
 
-def _make_impl():
+def _make_impl(num_heads: int):
     """Minimal AiterMLAImpl exposing only what _mla_fp8_prefill_attn reads."""
     from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1
 
     from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl
 
     impl = object.__new__(AiterMLAImpl)
-    impl.num_heads = NUM_HEADS
+    impl.num_heads = num_heads
     impl.v_head_dim = V_HEAD_DIM
     impl.scale = SCALE
     impl._mla_prefill_ps_asm_fwd = mla_prefill_ps_asm_fwd
@@ -88,8 +104,13 @@ def _make_impl():
     return impl
 
 
-def _build_prefill_metadata(seq_lens: list[int], device: torch.device):
-    """Build fp8_prefill_* metadata via the real builder (no hand-rolling)."""
+def _build_prefill_metadata(
+    seq_lens: list[int],
+    device: torch.device,
+    num_heads: int,
+    attn_out_dtype: torch.dtype = ATTN_OUT_DTYPE,
+):
+    """Build metadata with the production builder and output dtype."""
     from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder
 
     qo_indptr_cpu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
@@ -99,12 +120,13 @@ def _build_prefill_metadata(seq_lens: list[int], device: torch.device):
     max_q = max(seq_lens)
 
     builder = object.__new__(AiterMLAMetadataBuilder)
-    builder.num_heads = NUM_HEADS
+    builder.num_heads = num_heads
     builder.mla_dims = SimpleNamespace(v_head_dim=V_HEAD_DIM)
     builder._init_fp8_prefill_ps_buffers(
         max_num_reqs=len(seq_lens),
         max_prefill_qlen=max_q,
         max_num_batched_tokens=total_q,
+        attn_out_dtype=attn_out_dtype,
         device=device,
     )
 
@@ -115,7 +137,7 @@ def _build_prefill_metadata(seq_lens: list[int], device: torch.device):
     metadata = SimpleNamespace(prefill=prefill, num_decodes=0)
     common = SimpleNamespace(query_start_loc_cpu=qo_indptr_cpu)
     builder._build_fp8_prefill_ps_metadata(metadata, common)
-    return metadata, total_q
+    return metadata, total_q, builder
 
 
 def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
@@ -128,34 +150,95 @@ def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tenso
     return out.transpose(0, 1)  # [S, H, Dv]
 
 
+@pytest.mark.parametrize("num_heads", HEAD_COUNTS)
 @pytest.mark.parametrize("seq_len", [128, 512])
 @torch.inference_mode()
-def test_fp8_prefill_matches_reference(seq_len: int) -> None:
+def test_fp8_prefill_matches_reference(seq_len: int, num_heads: int) -> None:
     device = torch.device("cuda")
     torch.manual_seed(0)
 
-    metadata, total_q = _build_prefill_metadata([seq_len], device)
+    metadata, total_q, _ = _build_prefill_metadata([seq_len], device, num_heads)
 
-    q = torch.randn(
-        total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device
-    )
-    k = torch.randn(
-        total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16, device=device
-    )
-    v = torch.randn(total_q, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=device)
+    # Lock the workspace, as GPUModelRunner does after warmup, so the forward
+    # cannot quietly grow it. The builder reserves scratch at num_head_k and
+    # the forward requests it at its own padded width; if a future change makes
+    # the forward's width exceed the builder's, that is an under-reservation
+    # which is invisible on an unlocked workspace (it just grows) and fatal in
+    # a real serve. Locking turns it into a failure here.
+    #
+    # Note this does not catch the reverse -- a builder wider than the forward
+    # merely over-reserves -- so it would not have flagged the pre-PR
+    # max(16, num_heads) sizing, which over-reserves 2.8x at 24 heads.
+    from vllm.v1.worker.workspace import lock_workspace
+
+    lock_workspace()
+
+    qkv_kwargs = dict(dtype=ATTN_OUT_DTYPE, device=device)
+    q = torch.randn(total_q, num_heads, QK_HEAD_DIM, **qkv_kwargs)
+    k = torch.randn(total_q, num_heads, QK_HEAD_DIM, **qkv_kwargs)
+    v = torch.randn(total_q, num_heads, V_HEAD_DIM, **qkv_kwargs)
     out = torch.zeros(
-        total_q, NUM_HEADS * V_HEAD_DIM, dtype=torch.bfloat16, device=device
+        total_q, num_heads * V_HEAD_DIM, dtype=ATTN_OUT_DTYPE, device=device
     )
 
-    impl = _make_impl()
+    impl = _make_impl(num_heads)
     impl._mla_fp8_prefill_attn(q, k, v, metadata, out)
 
     out_ref = _reference(q, k, v)
 
     assert torch.isfinite(out).all()
     torch.testing.assert_close(
-        out.view(total_q, NUM_HEADS, V_HEAD_DIM).float(),
+        out.view(total_q, num_heads, V_HEAD_DIM).float(),
         out_ref.float(),
         atol=ATOL,
         rtol=RTOL,
     )
+
+
+@pytest.mark.parametrize("num_heads", HEAD_COUNTS)
+def test_fp8_prefill_metadata_width_matches_forward(num_heads: int) -> None:
+    """The PS metadata must be built for the head count the kernel is handed.
+
+    ``_init_fp8_prefill_ps_buffers``/``build()`` size the work and reduce maps
+    from ``num_head_k``, while ``_mla_fp8_prefill_attn`` pads q/k/v to its own
+    target. Divergence is not a numerical bug in either direction -- maps built
+    narrower than the kernel width still cover the real heads, since padding
+    only ever appends duplicate heads above them. It is a sizing bug: narrower
+    maps mean a lower head alignment, hence more partial tiles and a larger
+    reservation, while wider ones under-reserve the forward's scratch. Pin the
+    two to one helper rather than to independent expressions.
+    """
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper
+
+    width = AiterMLAHelper.get_fp8_prefill_num_heads(num_heads)
+    assert width % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
+    assert width >= num_heads
+    assert width - num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
+    # Every count the enable gate accepts must have a usable padded width.
+    assert AiterMLAHelper.is_valid_num_heads(num_heads)
+
+    # The builder must actually size its maps at that width. Assert against the
+    # buffer it allocated rather than recomputing the width here, or the test
+    # passes no matter what _init_fp8_prefill_ps_buffers does.
+    #
+    # fp8_ps_reduce_partial_map is sized by the partial-tile count, which is
+    # gcd-driven and so is a sharp probe of the width used: a lower head
+    # alignment yields *more* tiles. At 24 heads, sizing at the unpadded count
+    # gives 4x the tiles (64 vs 16) and ~2.8x the scratch reservation (193.6
+    # vs 68.6 MiB at batch=1/qlen=512) -- not wrong numerically (the 24-wide
+    # maps still cover the real heads) but the reason to pin the widths
+    # together.
+    from aiter import get_ps_metadata_info_v1
+
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla import _FP8_PREFILL_TILE_Q
+
+    def _tiles(nhk: int) -> int:
+        return get_ps_metadata_info_v1(
+            batch_size=1,
+            num_head_k=nhk,
+            max_qlen=512,
+            qlen_granularity=_FP8_PREFILL_TILE_Q,
+        )[5][0]
+
+    _, _, builder = _build_prefill_metadata([512], torch.device("cuda"), num_heads)
+    assert builder.fp8_ps_reduce_partial_map.numel() == _tiles(width)

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

@@ -414,9 +414,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             device=device,
         )
 
-        # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
-        self._fp8_prefill_enabled = (
-            _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
+        # The assembly prefill requires FP8 KV, bf16 output, and 16-aligned
+        # heads. It writes bf16 through a raw output pointer, so fp16 must use
+        # the standard prefill path. Head counts that are not a multiple of 16
+        # are replicate-padded up to one in _mla_fp8_prefill_attn, so the gate
+        # is the same head-count predicate the decode path uses.
+        self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
+            kv_cache_dtype_str == "fp8"
+            and vllm_config.model_config.dtype == torch.bfloat16
+            and AiterMLAHelper.is_valid_num_heads(self.num_heads)
         )
         if self._fp8_prefill_enabled:
             max_prefill_qlen = min(
@@ -427,6 +433,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 max_num_reqs,
                 max_prefill_qlen,
                 vllm_config.scheduler_config.max_num_batched_tokens,
+                vllm_config.model_config.dtype,
                 device,
             )
 
@@ -444,6 +451,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         max_num_reqs: int,
         max_prefill_qlen: int,
         max_num_batched_tokens: int,
+        attn_out_dtype: torch.dtype,
         device: torch.device,
     ) -> None:
         """Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.
@@ -466,13 +474,30 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 is bounded by this budget rather than by a single request's
                 ``max_prefill_qlen`` — concurrent requests can sum to more than
                 ``max_model_len`` when ``max_model_len < max_num_batched_tokens``.
+            attn_out_dtype: Dtype of the attention output buffer, used to size
+                the padded-head output scratch (small head counts only).
             device: Target device for the buffers.
         """
         from aiter import get_ps_metadata_info_v1
 
         # After kv_b_proj decompression, K has num_heads heads (same as Q).
         # So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
-        num_head_k = self.num_heads
+        # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
+        # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
+        # metadata for that same padded count so the work/reduce maps and the
+        # scratch reservations describe the width the kernel is handed.
+        #
+        # This was previously max(16, num_heads), which agrees with the helper
+        # at every head count reachable today (12 and 16 both give 16) and
+        # differs only above 16: at 24 it leaves num_head_k=24 while the
+        # forward pads to 32. That is not a correctness bug -- the 24-wide work
+        # maps still cover head-tiles 0..23, which are the real heads -- but it
+        # is expensive, because a lower head alignment yields more partial
+        # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
+        # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
+        # (measured on gfx950). Sizing both from one helper keeps the widths
+        # equal and takes the cheaper tiling.
+        num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)
         v_head_dim = self.mla_dims.v_head_dim
         # gqa_ratio = 1
         # qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
@@ -516,7 +541,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         from vllm.v1.worker.workspace import current_workspace_manager
 
         max_num_partial_tiles = reduce_partial_map_size
-        current_workspace_manager().get_simultaneous(
+        reservations: list[tuple[tuple[int, ...], torch.dtype]] = [
             (
                 (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
                 torch.float32,
@@ -526,7 +551,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 torch.float32,
             ),
             ((max_num_batched_tokens, num_head_k), torch.float32),
-        )
+        ]
+        if self.num_heads < num_head_k:
+            # Padded head counts also take their kernel output buffer from the
+            # workspace: the caller's [total_q, num_heads * v_head_dim] output
+            # cannot back a num_head_k-head view (see _mla_fp8_prefill_attn).
+            reservations.append(
+                ((max_num_batched_tokens, num_head_k, v_head_dim), attn_out_dtype)
+            )
+        current_workspace_manager().get_simultaneous(*reservations)
 
         logger.info(
             "FP8 MLA prefill PS buffers allocated "
@@ -566,7 +599,22 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         kv_indptr_cpu = qo_indptr_cpu.clone()
         seq_lens_cpu = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(torch.int32)
 
-        num_head_k = self.num_heads
+        # Head counts that are not a multiple of 16 (K3: 12/rank at TP8) are
+        # replicate-padded up to one in _mla_fp8_prefill_attn; build the PS
+        # metadata for that same padded count so the work/reduce maps and the
+        # scratch reservations describe the width the kernel is handed.
+        #
+        # This was previously max(16, num_heads), which agrees with the helper
+        # at every head count reachable today (12 and 16 both give 16) and
+        # differs only above 16: at 24 it leaves num_head_k=24 while the
+        # forward pads to 32. That is not a correctness bug -- the 24-wide work
+        # maps still cover head-tiles 0..23, which are the real heads -- but it
+        # is expensive, because a lower head alignment yields more partial
+        # tiles: gcd-driven, 24 heads -> 64 tiles vs 32 heads -> 16, i.e.
+        # 193.6 MiB of reservations instead of 68.6 MiB at batch=1/qlen=512
+        # (measured on gfx950). Sizing both from one helper keeps the widths
+        # equal and takes the cheaper tiling.
+        num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)
         # gqa_ratio = 1
         # qhead_granularity = max(gqa_ratio, 1)
         # qlen_granularity = _FP8_PREFILL_TILE_Q // qhead_granularity
@@ -927,8 +975,39 @@ class AiterMLAHelper:
         return -(-num_heads // m) * m
 
     @staticmethod
-    def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
-        m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
+    def get_fp8_prefill_num_heads(num_heads: int) -> int:
+        """Head count the FP8 PS asm prefill runs at: the next multiple of 16.
+
+        Deliberately *not* ``get_actual_mla_num_heads``. That one carves out
+        native H24 when ``_aiter_mla_native_h24_supported()``, which probes the
+        asm *decode* reducer and metadata. The prefill is a different kernel
+        pair (``mla_prefill_ps_asm_fwd`` + ``mla_reduce_v1``) with no such
+        probe, so 24 heads pad to 32 here even on a native-H24 build.
+
+        The PS metadata in ``_init_fp8_prefill_ps_buffers``/``build()`` must be
+        sized with this same function, or the work/reduce maps describe a
+        different head count than the kernel is handed.
+
+        This function itself has no upper bound; the ceiling comes from the
+        ``is_valid_num_heads`` gate, which rejects counts above
+        ``_AITER_MAX_PADDED_MLA_HEADS`` (128) unless they are already
+        16-aligned. That bound was established for the asm *decode* padding, so
+        an architecture with, say, 136 heads per rank would be refused the
+        prefill here for a reason that was never measured against this kernel
+        pair. Revisit the constant rather than special-casing prefill.
+        """
+        m = AiterMLAHelper._AITER_MIN_MLA_HEADS
+        return -(-num_heads // m) * m
+
+    @staticmethod
+    def get_mla_padded_q(
+        num_heads: int, q: torch.Tensor, target_heads: int | None = None
+    ) -> torch.Tensor:
+        m = (
+            target_heads
+            if target_heads is not None
+            else AiterMLAHelper.get_actual_mla_num_heads(num_heads)
+        )
         if num_heads == m:
             return q
         if m % num_heads == 0:
@@ -1036,10 +1115,14 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
         self.flash_attn_varlen_func = flash_attn_varlen_func
 
         # FP8 MLA prefill kernel imports (lazy, only when enabled).
-        # Auto-enabled on gfx950 when AITER ships the kernels.
-        # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
-        self._fp8_prefill_enabled = (
-            _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
+        # Auto-enabled on gfx950 when AITER ships the kernels. Only runs when the
+        # KV cache is FP8. Head counts that are not a multiple of 16 are
+        # replicate-padded up to one (see _mla_fp8_prefill_attn).
+        from vllm.utils.torch_utils import is_quantized_kv_cache
+
+        self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
+            is_quantized_kv_cache(kv_cache_dtype)
+            and AiterMLAHelper.is_valid_num_heads(self.num_heads)
         )
         if self._fp8_prefill_enabled:
             from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1
@@ -1082,7 +1165,36 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
 
         fp8_dtype = current_platform.fp8_dtype()
         total_q = q.shape[0]
-        nhead = self.num_heads
+        # PS asm prefill + mla_reduce_v1 require 16-aligned heads, and the PS
+        # metadata is built for get_fp8_prefill_num_heads(num_heads). For head
+        # counts that are not a multiple of 16 (K3 = 12/rank at TP8)
+        # replicate-pad q/k/v up to that count, then slice the output back to
+        # the real head count.
+        #
+        # Counts above 16 that are not multiples of 16 (24, 40, ...) are
+        # handled by the same code but are not reached by any current model
+        # and TP that fits: 96 heads would need TP4, whose weights exceed a
+        # 288 GiB GPU, and 128-head models land on 128/64/32/16/8. The path is
+        # kept general for future architectures and its numerics are covered by
+        # test_fp8_prefill_matches_reference[num_heads=24]. Note the cost is
+        # (padded - real)/real extra FLOPs and q/k/v bytes, which is worst just
+        # above a multiple of 16 (17 heads pad to 32); a future arch landing
+        # there should measure against the flash_attn_varlen_func fallback
+        # rather than assume the asm path wins.
+        #
+        # 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.
+        _real_nhead = self.num_heads
+        nhead = AiterMLAHelper.get_fp8_prefill_num_heads(_real_nhead)
+        _pad = nhead != _real_nhead
+        if _pad:
+            q = AiterMLAHelper.get_mla_padded_q(_real_nhead, q, nhead)
+            k = AiterMLAHelper.get_mla_padded_q(_real_nhead, k, nhead)
+            v = AiterMLAHelper.get_mla_padded_q(_real_nhead, v, nhead)
         v_head_dim = self.v_head_dim
         tile_q = _FP8_PREFILL_TILE_Q
 
@@ -1106,19 +1218,25 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
         num_partial_tiles = attn_metadata.fp8_prefill_num_partial_tiles
         assert num_partial_tiles is not None
 
-        # Reuse the caller's output buffer to skip the per-call alloc + copy.
-        # The ASM and reduce kernels both write to a [total_q, nhead, v_head_dim]
-        # view, which aliases the [total_q, nhead * v_head_dim] storage of out.
-        out_3d = out.view(total_q, nhead, v_head_dim)
-
-        # Per-call scratch (logits, attn_lse, final_lse) is served from the
-        # workspace manager so allocator churn in the prefill hot path is
-        # bounded after warmup, matching the pattern in PR #41002.
-        logits, attn_lse, final_lse = current_workspace_manager().get_simultaneous(
+        # Per-call scratch is served from the workspace manager so allocator
+        # churn in the prefill hot path is bounded after warmup, matching the
+        # pattern in PR #41002.  The builder reserves the maximum shape of every
+        # tensor requested here before the workspace is locked.
+        scratch: list[tuple[tuple[int, ...], torch.dtype]] = [
             ((num_partial_tiles * tile_q, nhead, v_head_dim), torch.float32),
             ((num_partial_tiles * tile_q, nhead), torch.float32),
             ((total_q, nhead), torch.float32),
-        )
+        ]
+        if _pad:
+            # The ASM and reduce kernels write a [total_q, nhead, v_head_dim]
+            # buffer.  With unpadded heads that aliases the caller's
+            # [total_q, nhead * v_head_dim] output, so write straight into it;
+            # padded heads do not fit that storage and need their own buffer.
+            scratch.append(((total_q, nhead, v_head_dim), out.dtype))
+
+        workspace = current_workspace_manager()
+        logits, attn_lse, final_lse, *pad_out = workspace.get_simultaneous(*scratch)
+        out_3d = pad_out[0] if _pad else out.view(total_q, nhead, v_head_dim)
 
         # Phase 1: persistent-scheduling assembly prefill kernel.
         self._mla_prefill_ps_asm_fwd(
@@ -1156,6 +1274,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
             final_lse,
         )
 
+        if _pad:
+            out.view(total_q, _real_nhead, v_head_dim).copy_(out_3d[:, :_real_nhead, :])
+
     def forward_mha(
         self,
         q: torch.Tensor,