From de2fb501202ba9dd85bd9701a8bf98ab1adf3496 Mon Sep 17 00:00:00 2001 From: billishyahao Date: Fri, 28 Aug 2026 09:30:29 +0800 Subject: [PATCH] [AMD] Enable aiter mla asm path through padding attn heads for Kimi K3 (#36356) --- .../srt/layers/attention/aiter_backend.py | 129 +++++++++++++++--- .../forward_mla_rocm.py | 10 +- 2 files changed, 119 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 244d63cde..89294bb9f 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -97,6 +97,16 @@ _use_fp8_prefill_attn = ( get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and is_gfx95_supported() ) +# (v_head_dim -> query head counts) that aiter's mla_reduce_v1 has an +# instantiation. This map is copied from MLA_REDUCE_ROUTER in +# aiter/csrc/kernels/mla/reduce.cu. See https://github.com/ROCm/aiter/blob/7915f53a4225b3f9cb632a97a23369dbebcf1be0/csrc/kernels/mla/reduce.cu#L923 +_MLA_REDUCE_V1_HEADS = { + 64: frozenset({64}), + 128: frozenset({1, 2, 4, 8, 10, 16, 32, 40, 64, 128}), + 512: frozenset({8, 16, 32, 48, 64, 80, 96, 112, 128}), +} + + # Persist # fast_mode=True if _use_mla_ps_kernel else False # intra_batch_mode=False if _use_mla_ps_kernel else True @@ -195,6 +205,34 @@ class AiterAttnBackend(AttentionBackend): -1 ] + # The asm fp8 prefill reduces through mla_reduce_v1, which only has + # instantiations for the head shapes in _MLA_REDUCE_V1_HEADS; anything + # else is tiled up to one it does carry, or falls back to + # flash_attn_varlen_func when the table has nothing to reach. + self.fp8_prefill_num_head = ( + self.check_fp8_prefill_num_head( + num_head=self.num_head, + num_kv_head=self.num_kv_head, + v_head_dim=self.v_head_dim, + ) + if self.use_mla + else None + ) + self.use_fp8_prefill_attn = ( + _use_fp8_prefill_attn and self.fp8_prefill_num_head is not None + ) + # Padding is only offered at GQA ratio 1, so the kv side takes the same + # delta and the ratio the PS metadata is built for stays put. + self.fp8_prefill_num_kv_head = self.num_kv_head + ( + (self.fp8_prefill_num_head or self.num_head) - self.num_head + ) + if self.use_fp8_prefill_attn and self.fp8_prefill_num_head != self.num_head: + logger.info( + f"aiter asm fp8 MLA prefill pads {self.num_head} query heads to " + f"{self.fp8_prefill_num_head}; mla_reduce_v1 has no " + f"{self.num_head}-head instantiation at head_dim {self.v_head_dim}." + ) + # Parse constants self.max_context_len = model_runner.model_config.context_len self.skip_prefill = skip_prefill @@ -307,14 +345,23 @@ class AiterAttnBackend(AttentionBackend): _valid_heads = self.num_head in (4, 8) or ( self.num_head % 16 == 0 and 16 <= self.num_head <= 128 ) - assert _valid_heads, ( + + may_run_mla_decode = self.may_run_mla_decode_kernel( + decode_attention_backend=model_runner.decode_attention_backend_str, + speculative_algorithm=get_spec().speculative_algorithm, + speculative_attention_mode=get_spec().speculative_attention_mode, + ) + # _mla_decode_fwd_with_head_pad brings any count below 16 up to it, + # by repetition when it divides 16 and by tiling otherwise. + _pad_heads_to_16 = self.num_head < 16 + assert _valid_heads or _pad_heads_to_16 or not may_run_mla_decode, ( f"Aiter MLA supports num_head of 4, 8, or multiples of 16 " f"in [16, 128].\n" f"Provided {self.num_head} number of heads.\n" - "Try adjusting tensor_parallel_size value." + "Try adjusting tensor_parallel_size value, or run decode on " + "another backend (--decode-attention-backend)." ) self.num_head_padded = 16 if self.num_head < 16 else self.num_head - self.head_repeat_factor = 16 // self.num_head if self.num_head < 16 else 1 self.enable_dp_attention = is_dp_attention_enabled() self.qo_indptr_ = torch.zeros( @@ -349,6 +396,41 @@ class AiterAttnBackend(AttentionBackend): self.fix_max_split_per_batch = self.max_split_per_batch + def pad_heads(self, x: torch.Tensor, padded: int) -> torch.Tensor: + num_head = x.shape[1] + reps = -(-padded // num_head) # ceil(padded / num_head) + return x.repeat(1, reps, 1)[:, :padded, :].contiguous() + + def check_fp8_prefill_num_head( + self, *, num_head: int, num_kv_head: int, v_head_dim: int + ) -> Optional[int]: + """Check _MLA_REDUCE_V1_HEADS to get head count to run the asm fp8 + mla prefill, return None if it is invalid so we will fall it back to + aiter fa implementation. + """ + supported = _MLA_REDUCE_V1_HEADS.get(v_head_dim, frozenset()) + if num_head in supported: + return num_head + if num_head != num_kv_head: + return None + larger = [h for h in supported if h > num_head] + return min(larger) if larger else None + + def may_run_mla_decode_kernel( + self, + *, + decode_attention_backend: Optional[str], + speculative_algorithm: Optional[str], + speculative_attention_mode: str, + ) -> bool: + """Decode whether aiter backend will invoke mla_decode_fwd""" + if decode_attention_backend == "aiter": + return True + return ( + speculative_algorithm is not None + and speculative_attention_mode == "prefill" + ) + def _get_aiter_paged_ragged_kv_cache_dtype(self) -> str: """``kv_cache_dtype`` string for ``paged_attention_ragged`` (aiter ``pa/pa_ragged.py``). @@ -483,7 +565,7 @@ class AiterAttnBackend(AttentionBackend): (reduce_partial_map_size, reduce_partial_map_type), ) = get_ps_metadata_info_v1( batch_size=batch_size, - num_head_k=self.num_kv_head, + num_head_k=self.fp8_prefill_num_kv_head, max_qlen=max_qlen, qlen_granularity=qlen_granularity, ) @@ -528,8 +610,8 @@ class AiterAttnBackend(AttentionBackend): reduce_partial_map: torch.Tensor, is_causal: bool = True, ): - gqa_ratio = self.num_head // self.num_kv_head - num_heads_k = self.num_kv_head + gqa_ratio = self.fp8_prefill_num_head // self.fp8_prefill_num_kv_head + num_heads_k = self.fp8_prefill_num_kv_head tile_q = 256 qhead_granularity = gqa_ratio qlen_granularity = tile_q // qhead_granularity @@ -768,20 +850,18 @@ class AiterAttnBackend(AttentionBackend): **kwargs, ): """Wrap mla_decode_fwd with head-dimension padding for num_head < 16. - - When head_repeat_factor > 1 (i.e. num_head is 4 or 8), q is - repeat-interleaved to reach num_head_padded (16) before the kernel - call, and the corresponding output columns are sliced back afterward. + The kernel only accepts 4, 8, or a multiple of 16 heads; any other + count crashes it. So we repeat q's heads to reach 16, real heads first. q / o must already be shaped (..., num_head, head_dim). """ - if self.head_repeat_factor > 1: - q_in = q.repeat_interleave(self.head_repeat_factor, dim=1) + if self.num_head_padded != self.num_head: + q_in = self.pad_heads(q, self.num_head_padded) o = q.new_empty( (q.shape[0], self.num_head_padded, layer.v_head_dim), dtype=self.input_dtype, ) mla_decode_fwd(q_in, k_buffer_flat, o, **kwargs) - return o[:, :: self.head_repeat_factor, :] + return o[:, : self.num_head, :] else: o = q.new_empty( (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), @@ -800,6 +880,15 @@ class AiterAttnBackend(AttentionBackend): total_q = q.shape[0] nhead = layer.tp_q_head_num v_head_dim = layer.v_head_dim + # mla_reduce_v1 dispatches on the head count, so a model it has no + # instantiation for runs on the next one up, with the query heads tiled + # to fill and the extra output columns sliced back off. + head_pad = self.fp8_prefill_num_head - nhead + if head_pad: + q = self.pad_heads(q, self.fp8_prefill_num_head) + k = self.pad_heads(k, self.fp8_prefill_num_head) + v = self.pad_heads(v, self.fp8_prefill_num_head) + nhead = self.fp8_prefill_num_head if q.dtype != fp8_dtype: q = q.to(fp8_dtype) @@ -865,7 +954,7 @@ class AiterAttnBackend(AttentionBackend): output, final_lse, ) - return output + return output[:, : layer.tp_q_head_num, :] if head_pad else output def init_forward_metadata_out_graph( self, @@ -1316,9 +1405,11 @@ class AiterAttnBackend(AttentionBackend): reduce_partial_map = None fp8_prefill_kv_indices = None - if _use_fp8_prefill_attn: + if self.use_fp8_prefill_attn: tile_q = 256 - qlen_granularity = tile_q // (self.num_head // self.num_kv_head) + qlen_granularity = tile_q // ( + self.fp8_prefill_num_head // self.fp8_prefill_num_kv_head + ) ( work_metadata, work_indptr, @@ -2040,7 +2131,7 @@ class AiterAttnBackend(AttentionBackend): ): extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) if kv_indices.shape[0] == 0 or extend_no_prefix: - if _use_fp8_prefill_attn: + if self.use_fp8_prefill_attn: output = self.mla_fp8_prefill_attn( q, k, @@ -2073,7 +2164,7 @@ class AiterAttnBackend(AttentionBackend): k_pe = k_pe.to(dtype) if ( - _use_fp8_prefill_attn + self.use_fp8_prefill_attn and layer.kv_b_proj.weight.dtype == torch.uint8 ): # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat @@ -2113,7 +2204,7 @@ class AiterAttnBackend(AttentionBackend): == forward_batch.extend_seq_lens.shape ) - if _use_fp8_prefill_attn: + if self.use_fp8_prefill_attn: return self.mla_fp8_prefill_attn(q, k, v, layer) else: return flash_attn_varlen_func( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 378f3cbf6..50395b223 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -877,5 +877,13 @@ class DeepseekMLARocmForwardMixin: Skip rope in prepare and let the fused kernel in forward_absorb_rocm_core handle it, when running aiter-backend MLA on gfx95 (i.e., the `else` branch in forward_absorb_rocm_core that calls fused_qk_rope_cat_and_cache_mla). + + A layer without a rotary_emb has nothing to fuse: that branch reads + rotary_emb.cos_cache, so skipping the standalone rope there ends in + AttributeError on None. Kimi-K3 has such layers. """ - return _use_aiter_gfx95 and self.current_attention_backend == "aiter" + return ( + _use_aiter_gfx95 + and self.current_attention_backend == "aiter" + and self.rotary_emb is not None + )