From 5172c378456fa9735ba6429a2ca7688d1152d256 Mon Sep 17 00:00:00 2001 From: Thomas Wang Date: Fri, 27 Feb 2026 07:14:52 +0800 Subject: [PATCH] [AMD] Use fused GEMM with FP8 cast for FP8 prefill (#19422) --- .../srt/layers/attention/aiter_backend.py | 21 ++++++++--- .../quark/schemes/quark_w4a4_mxfp4.py | 35 +++++++++++++++--- .../attention_forward_methods/forward_mha.py | 37 ++++++++++++++----- 3 files changed, 73 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 0062a5573..2d7692449 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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( - total_s, device=q.device, dtype=torch.int32 - ) + kv_indices_asm = self.forward_metadata.fp8_prefill_kv_indices tile_q = 256 reduce_indptr = self.forward_metadata.reduce_indptr diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index 984e15696..04f7cd246 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/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) - - return y + else: + return y diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index ceb18afc7..7c83904f8 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/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: - kv = self.kv_b_proj( - kv_a_quanted, + if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8: + # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat + # 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] - 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 :] + if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn: + kv = self.kv_b_proj(kv_a_quanted)[0] + else: + 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(