minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Enable aiter mla asm path through padding attn heads for Kimi K3 (#36356)

Browse Source
This commit is contained in:
billishyahao
2026-08-27 18:30:29 -07:00
committed by GitHub
parent 22f37d414b
commit de2fb50120
2 changed files with 119 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+110 -19
View File
@@ -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(
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py
+9 -1
View File
@@ -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
)