[AMD] Use fused GEMM with FP8 cast for FP8 prefill (#19422)

python/sglang/srt/layers/attention/aiter_backend.py

@@ -92,6 +92,7 @@ class ForwardMetadata:
     custom_mask: Optional[torch.Tensor] = None
     mask_indptr: Optional[torch.Tensor] = None
     max_extend_len: Optional[int] = None
+    fp8_prefill_kv_indices: Optional[torch.Tensor] = None
 
 
 global_workspace_buffer = None
@@ -757,6 +758,7 @@ class AiterAttnBackend(AttentionBackend):
                 reduce_indptr = None
                 reduce_final_map = None
                 reduce_partial_map = None
+                fp8_prefill_kv_indices = None
 
                 if _use_fp8_prefill_attn:
                     tile_q = 256
@@ -785,6 +787,11 @@ class AiterAttnBackend(AttentionBackend):
                         is_causal=True,
                     )
 
+                    total_s = int(forward_batch.extend_seq_lens.sum())
+                    fp8_prefill_kv_indices = torch.arange(
+                        total_s, device=self.device, dtype=torch.int32
+                    )
+
                 self.forward_metadata = ForwardMetadata(
                     self.mla_indices_updater_prefill.kv_indptr,
                     self.mla_indices_updater_prefill.kv_indices,
@@ -798,6 +805,7 @@ class AiterAttnBackend(AttentionBackend):
                     reduce_indptr=reduce_indptr,
                     reduce_final_map=reduce_final_map,
                     reduce_partial_map=reduce_partial_map,
+                    fp8_prefill_kv_indices=fp8_prefill_kv_indices,
                 )
             else:
                 self.indices_updater_prefill.update(
@@ -1470,20 +1478,21 @@ class AiterAttnBackend(AttentionBackend):
                         nhead = layer.tp_q_head_num
                         v_head_dim = layer.v_head_dim
 
+                        # q is cast here (after RoPE).
+                        # k/v are already FP8 for MXFP4 main model (fused kv_b_proj),
+                        # but need casting for FP8/BF16 weights (e.g. MTP draft model).
                         if q.dtype != fp8_dtype:
-                            q = q.float().to(fp8_dtype)
+                            q = q.to(fp8_dtype)
                         if k.dtype != fp8_dtype:
-                            k = k.float().to(fp8_dtype)
+                            k = k.to(fp8_dtype)
                         if v.dtype != fp8_dtype:
-                            v = v.float().to(fp8_dtype)
+                            v = v.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(
+                        kv_indices_asm = self.forward_metadata.fp8_prefill_kv_indices
-                            total_s, device=q.device, dtype=torch.int32
-                        )
 
                         tile_q = 256
                         reduce_indptr = self.forward_metadata.reduce_indptr

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py

@@ -10,6 +10,9 @@ from sglang.srt.utils import is_hip
 
 _is_hip = is_hip()
 if _is_hip:
+    from aiter.ops.triton.gemm.fused.fused_gemm_afp4wfp4_split_cat import (
+        fused_gemm_afp4wfp4_split_cat,
+    )
     from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
     from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import gemm_afp4wfp4_pre_quant
     from aiter.ops.triton.quant import dynamic_mxfp4_quant
@@ -87,20 +90,29 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
         assert bias is None, "bias is not supported"
 
         three_d = False
+        fused_gemm_split_cat = False
         x_s = None
         y = None
+
         if isinstance(x, tuple):
             assert len(x) in [
                 2,
                 3,
-            ], "For tuple input, only (x, x_s) or (x, x_s, y) formats are accepted"
+                5,
+            ], "For tuple input, only (x, x_s), (x, x_s, y), or (x, y, S1, S2, out_dtype) formats are accepted"
             if len(x) == 2:
                 x, x_s = x
             elif len(x) == 3:
                 x, x_s, y = x
+            elif len(x) == 5:
+                x, y, S1, S2, out_dtype = x
+                fused_gemm_split_cat = True
 
         use_fused_quant_gemm = (
-            x_s is None and y is not None and layer.weight.shape[0] == y.shape[1]
+            not fused_gemm_split_cat
+            and x_s is None
+            and y is not None
+            and layer.weight.shape[0] == y.shape[1]
         )
 
         if x.dim() == 3:
@@ -126,10 +138,23 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
         if use_fused_quant_gemm:
             gemm_afp4wfp4_pre_quant(x_q, layer.weight, layer.weight_scale, y.dtype, y)
             y = y.to(x.dtype)
+        elif fused_gemm_split_cat:
+            k, v = fused_gemm_afp4wfp4_split_cat(
+                x=x_q,
+                w=layer.weight,
+                y=y,
+                x_scale=x_s,
+                w_scale=layer.weight_scale,
+                S1=S1,
+                S2=S2,
+                dtype=out_dtype,
+            )
         else:
             gemm_afp4wfp4(x_q, layer.weight, x_s, layer.weight_scale, self.out_dtype, y)
 
-        if three_d:
+        if fused_gemm_split_cat:
+            return k, v
+        elif three_d:
             return y.view(*output_shape)
-
+        else:
-        return y
+            return y

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -17,7 +17,11 @@ from sglang.srt.models.deepseek_common.utils import (
     _use_aiter_gfx95,
 )
 from sglang.srt.server_args import get_global_server_args
-from sglang.srt.utils import BumpAllocator, next_power_of_2
+from sglang.srt.utils import BumpAllocator, get_bool_env_var, next_power_of_2
+
+_use_fp8_prefill_attn = (
+    get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and _use_aiter_gfx95
+)
 
 if TYPE_CHECKING:
     from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
@@ -28,6 +32,7 @@ if _is_cuda:
 if _use_aiter_gfx95:
     from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
 
+    from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
     from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
 
 # Configs for DeepSeek-V3:
@@ -232,17 +237,31 @@ class DeepseekMHAForwardMixin:
                     q.dtype,
                     forward_batch,
                 )
-        if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
+        if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8:
-            kv = self.kv_b_proj(
+            # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
-                kv_a_quanted,
+            # into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
+            # directly in FP8, avoiding a separate elementwise cast
+            k, v = self.kv_b_proj(
+                (
+                    kv_a,
+                    k_pe.expand(-1, self.num_local_heads, -1),
+                    self.qk_nope_head_dim,
+                    self.v_head_dim,
+                    fp8_dtype,
+                )
             )[0]
         else:
-            kv = self.kv_b_proj(kv_a)[0]
+            if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
-        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
+                kv = self.kv_b_proj(kv_a_quanted)[0]
-        k_nope = kv[..., : self.qk_nope_head_dim]
+            else:
-        v = kv[..., self.qk_nope_head_dim :]
+                kv = self.kv_b_proj(kv_a)[0]
+            kv = kv.view(
+                -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
+            )
+            k_nope = kv[..., : self.qk_nope_head_dim]
+            v = kv[..., self.qk_nope_head_dim :]
 
-        k = self._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
+            k = self._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
         return q, k, v, forward_batch
 
     def forward_normal_core(