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

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

Browse Source
This commit is contained in:
Thomas Wang
2026-02-27 07:14:52 +08:00
committed by GitHub
parent b79fa9ccd7
commit 5172c37845
3 changed files with 73 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/sglang/srt/layers/attention/aiter_backend.py
View File

@@ -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

35
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py
View File

@@ -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

37
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
View File

@@ -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(