[HIP] FIX MLA the nhead fold error for cp round robin

aiter/mla.py

@@ -398,6 +398,20 @@ def _persistent_mla_decode_max_batch():
         return _MLA_DECODE_PERSISTENT_MAX_BATCH_DEFAULT
 
 
+def _fold_seqlen_indptr(indptr, fold_factor):
+    """Repeat each batch's seqlen ``fold_factor`` times (head-folding pseudo-batches)."""
+    lens = indptr[1:] - indptr[:-1]
+    folded_lens = lens.repeat_interleave(fold_factor)
+    out = torch.empty(
+        indptr.shape[0] + (fold_factor - 1) * (indptr.shape[0] - 1),
+        dtype=indptr.dtype,
+        device=indptr.device,
+    )
+    out[0] = 0
+    out[1:] = torch.cumsum(folded_lens, dim=0).to(indptr.dtype)
+    return out
+
+
 def _use_persistent_mla_decode(bs, nhead, max_seqlen_q, q_dtype, kv_dtype):
     """Whether to keep the persistent MLA decode kernel.
 
@@ -813,6 +827,13 @@ def mla_decode_fwd(
                 and q.dtype == dtypes.bf16
                 and kv_buffer.dtype == dtypes.bf16
             )
+            or (
+                get_gfx() == "gfx950"
+                and nhead == 96
+                and q.dtype == dtypes.fp8
+                and kv_buffer.dtype == dtypes.fp8
+                and max_seqlen_q <= 6
+            )
         ):
             # Natively support cases
             pass
@@ -837,6 +858,14 @@ def mla_decode_fwd(
                 o_orig = o
 
             o = o.view(total_s, nhead, -1)
+
+            qo_indptr = _fold_seqlen_indptr(qo_indptr, fold_factor)
+            if g_kv_indptr is not None:
+                g_kv_indptr = _fold_seqlen_indptr(g_kv_indptr, fold_factor)
+                # Each pseudo-batch shares the original batch's local kv begin.
+                kv_indptr = torch.cat(
+                    [kv_indptr[:-1].repeat_interleave(fold_factor), kv_indptr[-1:]]
+                )
             io_transformed = True
         else:
             assert False, f"{nhead=} and {max_seqlen_q=} not supported"

aiter/ops/attention.py

@@ -1224,6 +1224,12 @@ def get_mla_metadata_info_v1(
         and kv_dtype == dtypes.bf16
         and q_dtype == dtypes.bf16
         and num_head_qo != 48
+    ) or (
+        get_gfx() == "gfx950"
+        and q_dtype == dtypes.fp8
+        and kv_dtype == dtypes.fp8
+        and num_head_qo == 96
+        and effective_seqlen_qo <= 6
     ):
         if num_head_qo * 2 > 128:
             max_qo_tiles_per_batch = effective_seqlen_qo
@@ -1634,6 +1640,13 @@ def decode_update_mla_metadata_v1(
             and q_is_fp8
             and kv_is_fp8
         )
+        or (
+            arch_id == "gfx950"
+            and num_heads_per_head_k == 96
+            and q_is_fp8
+            and kv_is_fp8
+            and max_seqlen_qo <= 6
+        )
     )
     cu_num = work_indptr.shape[0] - 1
     tile_reduce_cnt = reduce_indptr.shape[0] - 1

csrc/kernels/mla/metadata/v1_2_device.cuh

@@ -904,6 +904,8 @@ void get_mla_metadata_v1_2_device(const aiter_tensor_t& seqlens_qo_indptr, // [b
         ((arch_id == "gfx942") && (num_heads == 128) && q_is_fp8 && kv_is_fp8) ||
         ((arch_id == "gfx950") && q_is_fp8 && kv_is_fp8 &&
          ((num_heads == 32) || (num_heads == 64) || (num_heads == 128))) ||
+        ((arch_id == "gfx950") && q_is_fp8 && kv_is_fp8 && (num_heads == 96) &&
+         (max_seqlen_qo <= 6)) ||
         hk_mtp_experimental;
 
     if(!natively_supported && (num_heads % 16 == 0))

csrc/py_itfs_cu/asm_mla.cu

@@ -983,7 +983,8 @@ void mla_decode_stage1_asm_fwd(
     } else if (arch_id == "gfx950" && q_type == "fp8" && kv_type == "fp8" && persistent
                && ((gqa_ratio == 32 && max_seqlen_q >= 4)
                    || (gqa_ratio == 64 && max_seqlen_q >= 2)
-                   || (gqa_ratio == 128))){
+                   || (gqa_ratio == 128)
+                   || (gqa_ratio == 96 && max_seqlen_q <= 6))){
         config_max_seqlen_q = 4;
         config_gqa_ratio = 32;
         args.s_MQA = gqa_ratio;

op_tests/test_mla_persistent.py

@@ -649,6 +649,13 @@ def torch_mla_extend_split_kv(
             and max_seqlen_q == 1
         )
         or (get_gfx() == "gfx950" and not is_fp8_q and not is_fp8_kvc)
+        or (
+            get_gfx() == "gfx950"
+            and nheads == 96
+            and is_fp8_q
+            and is_fp8_kvc
+            and max_seqlen_q <= 6
+        )
     ):
         # Natively support cases
         pass
@@ -714,9 +721,21 @@ def torch_mla_extend_split_kv(
         # For a split chunk, we need to account for the position offsets of Q and KV within the batch
         causal_diagonal = None
         if is_causal:
-            q_local_start = qo_start - qo_indptr[batch_idx].item()
+            if q_ratio > 1:
+                # Folded head layout: qo_start/qo_end index the folded token
+                # space, where every (batch, head-group) pair is its own batch of
+                # max_seqlen_q tokens starting at row[0] * max_seqlen_q. qo_indptr
+                # only describes the unfolded batches, so measuring the query
+                # offset against it shifts the diagonal by head_group *
+                # max_seqlen_q for every group past the first.
+                q_local_start = qo_start - row[0].item() * max_seqlen_q
+                total_q_len = max_seqlen_q
+            else:
+                q_local_start = qo_start - qo_indptr[batch_idx].item()
+                total_q_len = (
+                    qo_indptr[batch_idx + 1].item() - qo_indptr[batch_idx].item()
+                )
             kv_local_start = (kv_start - kv_indptr[batch_idx].item()) * page_size
-            total_q_len = qo_indptr[batch_idx + 1].item() - qo_indptr[batch_idx].item()
             causal_diagonal = (
                 q_local_start - kv_local_start + cur_real_kv_seq_len - total_q_len
             )

op_tests/test_mla_persistent_round_robin.py

@@ -736,7 +736,7 @@ def test_mla_cp(
     type=dtypes.str2tuple,
     nargs="*",
     const=None,
-    default=[(16, 4), (32, 3), (64, 1), (64, 2), (128, 2), (16, 8), (64, 17)],
+    default=[(16, 4), (32, 3), (64, 1), (64, 2), (128, 2), (16, 8), (96, 8), (64, 17)],
     help="""Number of heads, decode_qlen pairs.
     e.g.: -n 16,4""",
 )