diff --git a/aiter/mla.py b/aiter/mla.py index a0b8c311df..a294f3e950 100644 --- a/aiter/mla.py +++ b/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" diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index ec9a6374fd..47365bb999 100644 --- a/aiter/ops/attention.py +++ b/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 diff --git a/csrc/kernels/mla/metadata/v1_2_device.cuh b/csrc/kernels/mla/metadata/v1_2_device.cuh index 19fac5f04e..d0c796b676 100644 --- a/csrc/kernels/mla/metadata/v1_2_device.cuh +++ b/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)) diff --git a/csrc/py_itfs_cu/asm_mla.cu b/csrc/py_itfs_cu/asm_mla.cu index f4e5800324..e3c334403a 100644 --- a/csrc/py_itfs_cu/asm_mla.cu +++ b/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; diff --git a/op_tests/test_mla_persistent.py b/op_tests/test_mla_persistent.py index 1448381415..d8a5976783 100644 --- a/op_tests/test_mla_persistent.py +++ b/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 ) diff --git a/op_tests/test_mla_persistent_round_robin.py b/op_tests/test_mla_persistent_round_robin.py index abf440c1c1..c433fef6f8 100644 --- a/op_tests/test_mla_persistent_round_robin.py +++ b/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""", )