diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index cf867a6a1..0cff59ea9 100644 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -19,6 +19,7 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import is_gfx95_supported if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention @@ -30,7 +31,11 @@ try: flash_attn_varlen_func, get_mla_metadata_info_v1, get_mla_metadata_v1, + get_ps_metadata_info_v1, + get_ps_metadata_v1, mha_batch_prefill_func, + mla_prefill_ps_asm_fwd, + mla_reduce_v1, paged_attention_ragged, ) from aiter.mla import mla_decode_fwd, mla_prefill_fwd @@ -49,6 +54,11 @@ logger = logging.getLogger(__name__) # Use aiter mla persist design for fp8-kv cache _use_mla_ps_kernel = get_bool_env_var("SGLANG_AITER_MLA_PERSIST", "True") +# Use fp8 prefill only on gfx95 +_use_fp8_prefill_attn = ( + get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and is_gfx95_supported() +) + # Persist # fast_mode=True if _use_mla_ps_kernel else False # intra_batch_mode=False if _use_mla_ps_kernel else True @@ -308,6 +318,94 @@ class AiterAttnBackend(AttentionBackend): dtype_kv=dtype, ) + def make_mla_prefill_ps_meta_data_buffer( + self, batch_size: int, max_qlen: int, qlen_granularity: int + ): + ( + (work_meta_data_size, work_meta_data_type), + (work_indptr_size, work_indptr_type), + (work_info_size, work_info_type), + (reduce_indptr_size, reduce_indptr_type), + (reduce_final_map_size, reduce_final_map_type), + (reduce_partial_map_size, reduce_partial_map_type), + ) = get_ps_metadata_info_v1( + batch_size=batch_size, + num_head_k=self.num_kv_head, + max_qlen=max_qlen, + qlen_granularity=qlen_granularity, + ) + + device = self.device + work_metadata_ptrs = torch.empty( + work_meta_data_size, dtype=work_meta_data_type, device=device + ) + work_indptr = torch.empty( + work_indptr_size, dtype=work_indptr_type, device=device + ) + work_info = torch.empty(work_info_size, dtype=work_info_type, device=device) + reduce_indptr = torch.empty( + reduce_indptr_size, dtype=reduce_indptr_type, device=device + ) + reduce_final_map = torch.empty( + reduce_final_map_size, dtype=reduce_final_map_type, device=device + ) + reduce_partial_map = torch.empty( + reduce_partial_map_size, dtype=reduce_partial_map_type, device=device + ) + + return ( + work_metadata_ptrs, + work_indptr, + work_info, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + ) + + def make_mla_prefill_ps_meta_data( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + seq_lens: torch.Tensor, + work_metadata: torch.Tensor, + work_indptr: torch.Tensor, + work_info: torch.Tensor, + reduce_indptr: torch.Tensor, + reduce_final_map: torch.Tensor, + 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 + tile_q = 256 + qhead_granularity = gqa_ratio + qlen_granularity = tile_q // qhead_granularity + kvlen_granularity = max(128, self.page_size) + block_size = self.page_size + + qo_indptr_cpu = qo_indptr.to("cpu", dtype=torch.int32) + kv_indptr_cpu = kv_indptr.to("cpu", dtype=torch.int32) + seq_lens_cpu = seq_lens.to("cpu", dtype=torch.int32) + + get_ps_metadata_v1( + qo_indptr_cpu, + kv_indptr_cpu, + seq_lens_cpu, + gqa_ratio, + num_heads_k, + work_metadata, + work_indptr, + work_info, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + qhead_granularity=qhead_granularity, + qlen_granularity=qlen_granularity, + kvlen_granularity=kvlen_granularity, + block_size=block_size, + is_causal=is_causal, + ) + def init_forward_metadata(self, forward_batch: ForwardBatch): """Init auxiliary variables for triton attention backend.""" @@ -587,15 +685,56 @@ class AiterAttnBackend(AttentionBackend): spec_info=None, ) - kv_indices = self.mla_indices_updater_prefill.kv_indices + max_q_len = self.mla_indices_updater_prefill.max_q_len + qo_indptr = self.mla_indices_updater_prefill.qo_indptr + + work_metadata = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + + if _use_fp8_prefill_attn: + tile_q = 256 + qlen_granularity = tile_q // (self.num_head // self.num_kv_head) + ( + work_metadata, + work_indptr, + work_info_set, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + ) = self.make_mla_prefill_ps_meta_data_buffer( + bs, max_q_len, qlen_granularity + ) + + self.make_mla_prefill_ps_meta_data( + qo_indptr, + qo_indptr, + forward_batch.seq_lens, + work_metadata, + work_indptr, + work_info_set, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + is_causal=True, + ) self.forward_metadata = ForwardMetadata( self.mla_indices_updater_prefill.kv_indptr, - kv_indices, - self.mla_indices_updater_prefill.qo_indptr, + self.mla_indices_updater_prefill.kv_indices, + qo_indptr, self.kv_last_page_len[:bs], - self.mla_indices_updater_prefill.max_q_len, + max_q_len, self.mla_indices_updater_prefill.max_kv_len, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, ) else: self.indices_updater_prefill.update( @@ -1047,18 +1186,93 @@ class AiterAttnBackend(AttentionBackend): ): extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) if kv_indices.shape[0] == 0 or extend_no_prefix: - o = flash_attn_varlen_func( - q, - k, - v, - qo_indptr, - qo_indptr, - max_q_len, - max_q_len, - softmax_scale=layer.scaling, - causal=True, - ) - return o + if _use_fp8_prefill_attn: + total_s = q.shape[0] + nhead = layer.tp_q_head_num + v_head_dim = layer.v_head_dim + + if q.dtype != fp8_dtype: + q = q.float().to(fp8_dtype) + if k.dtype != fp8_dtype: + k = k.float().to(fp8_dtype) + if v.dtype != fp8_dtype: + v = v.float().to(fp8_dtype) + one_scale = torch.tensor( + 1.0, dtype=torch.float32, device=q.device + ) + + kv_indptr_asm = qo_indptr + kv_indices_asm = torch.arange( + total_s, device=q.device, dtype=torch.int32 + ) + + tile_q = 256 + reduce_indptr = self.forward_metadata.reduce_indptr + reduce_final_map = self.forward_metadata.reduce_final_map + reduce_partial_map = self.forward_metadata.reduce_partial_map + + logits = torch.empty( + (reduce_partial_map.size(0) * tile_q, nhead, v_head_dim), + dtype=torch.float32, + device=q.device, + ) + attn_lse = torch.empty( + (reduce_partial_map.size(0) * tile_q, nhead), + dtype=torch.float32, + device=q.device, + ) + final_lse = torch.empty( + (total_s, nhead), + dtype=torch.float32, + device=q.device, + ) + output = q.new_empty( + (total_s, nhead, v_head_dim), + dtype=self.input_dtype, + ) + + mla_prefill_ps_asm_fwd( + q, + k, + v, + qo_indptr, + kv_indptr_asm, + kv_indices_asm, + self.forward_metadata.work_indptr, + self.forward_metadata.work_info_set, + max_q_len, + layer.scaling, + True, + logits, + attn_lse, + output, + one_scale, + one_scale, + one_scale, + ) + mla_reduce_v1( + logits, + attn_lse, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + tile_q, + output, + final_lse, + ) + else: + output = flash_attn_varlen_func( + q, + k, + v, + qo_indptr, + qo_indptr, + max_q_len, + max_q_len, + softmax_scale=layer.scaling, + causal=True, + ) + return output elif layer.qk_head_dim != (kv_lora_rank + qk_rope_head_dim): K_Buffer = torch.index_select(K_Buffer, 0, kv_indices) kvc, k_pe = torch.split(