Fp8 prefill attn kernel integration (#18528)

Co-authored-by: kkHuang-amd <wunhuang@amd.com>

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.kv_indices,
-                    self.mla_indices_updater_prefill.qo_indptr,
+                    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(
+                    if _use_fp8_prefill_attn:
-                        q,
+                        total_s = q.shape[0]
-                        k,
+                        nhead = layer.tp_q_head_num
-                        v,
+                        v_head_dim = layer.v_head_dim
-                        qo_indptr,
+
-                        qo_indptr,
+                        if q.dtype != fp8_dtype:
-                        max_q_len,
+                            q = q.float().to(fp8_dtype)
-                        max_q_len,
+                        if k.dtype != fp8_dtype:
-                        softmax_scale=layer.scaling,
+                            k = k.float().to(fp8_dtype)
-                        causal=True,
+                        if v.dtype != fp8_dtype:
-                    )
+                            v = v.float().to(fp8_dtype)
-                    return o
+                        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(