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

[ROCM] Optimized deepseek-r1 model with rmsnorm + fp8 quant fusion (#12689)

Browse Source 
should be clean after https://github.com/sgl-project/sglang/pull/13017 landed
This commit is contained in:
yctseng0211
2025-11-11 18:59:10 +08:00
committed by GitHub
parent ea10a9d165
commit 4a78031a71
6 changed files with 331 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
python/sglang/srt/layers/communicator.py
View File

@@ -66,6 +66,8 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
_is_gfx95_supported = is_gfx95_supported()
if _use_aiter and _is_gfx95_supported:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048
@@ -248,7 +250,7 @@ class LayerCommunicator:
hidden_states: torch.Tensor,
residual: torch.Tensor,
forward_batch: ForwardBatch,
qaunt_format: str = "",
quant_format: str = "",
):
if hidden_states.shape[0] == 0:
residual = hidden_states
@@ -267,7 +269,7 @@ class LayerCommunicator:
if residual is None:
residual = hidden_states
if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format):
if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
hidden_states, *_, _ = fused_rms_mxfp4_quant(
hidden_states,
self.input_layernorm.weight,
@@ -277,10 +279,26 @@ class LayerCommunicator:
None,
None,
)
elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format):
hidden_states, _, _, _res = fused_rms_fp8_group_quant(
hidden_states,
self.input_layernorm.weight,
self.input_layernorm.variance_epsilon,
inp2=None,
inp2_weight=None,
inp2_epsilon=None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
)
else:
hidden_states = self.input_layernorm(hidden_states)
else:
if _use_aiter and _is_gfx95_supported and ("mxfp4" in qaunt_format):
if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
hidden_states, *_, residual = fused_rms_mxfp4_quant(
hidden_states,
self.input_layernorm.weight,
@@ -290,6 +308,23 @@ class LayerCommunicator:
None,
residual,
)
elif _use_aiter and _is_gfx95_supported and ("fp8" in quant_format):
# RMSNorm + FP8 per-group quant
# return hidden_states：
# out_fp8 : FP8 activation → a8w8 GEMM
# out_bs : block-scale → gemm_a8w8_blockscale.x_scale
hidden_states, _, _, residual = fused_rms_fp8_group_quant(
hidden_states,
self.input_layernorm.weight,
self.input_layernorm.variance_epsilon,
inp2=None,
inp2_weight=None,
inp2_epsilon=None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=residual,
output_unquantized_inp1=False,
)
else:
hidden_states, residual = self.input_layernorm(
hidden_states, residual

10
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -488,6 +488,16 @@ class Fp8LinearMethod(LinearMethodBase):
True, # is_vnni
)
if isinstance(x, tuple):
return self.w8a8_block_fp8_linear(
input=x[0],
weight=layer.weight,
block_size=self.quant_config.weight_block_size,
weight_scale=layer.weight_scale_inv,
input_scale=x[1],
bias=bias,
)
return self.w8a8_block_fp8_linear(
input=x,
weight=layer.weight,

21
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -263,19 +263,32 @@ def aiter_w8a8_block_fp8_linear(
input_scale: Optional[torch.Tensor] = None,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
assert input_scale is None
# assert input_scale is None
input_2d = input.view(-1, input.shape[-1])
output_shape = [*input.shape[:-1], weight.shape[0]]
q_input, x_scale = aiter_per1x128_quant(input_2d, quant_dtype=aiter.dtypes.fp8)
# if input_scale not None, input is quanted
if input_scale is not None:
q_input = input_2d
x_scale = input_scale
else:
q_input, x_scale = aiter_per1x128_quant(input_2d, quant_dtype=aiter.dtypes.fp8)
output = gemm_a8w8_blockscale(
q_input, weight, x_scale, weight_scale, dtype=input.dtype
q_input,
weight,
x_scale,
weight_scale,
dtype=torch.bfloat16 if input_scale is not None else input.dtype,
)
if bias is not None:
output += bias
return output.to(dtype=input_2d.dtype).view(*output_shape)
return output.to(
dtype=torch.bfloat16 if input_scale is not None else input_2d.dtype
).view(*output_shape)
def triton_w8a8_block_fp8_linear(

132
python/sglang/srt/models/deepseek_v2.py
View File

@@ -153,6 +153,8 @@ _is_gfx95_supported = is_gfx95_supported()
_use_aiter_gfx95 = _use_aiter and _is_gfx95_supported
if _use_aiter_gfx95:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
from sglang.srt.layers.quantization.quark.utils import quark_post_load_weights
from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
batched_gemm_afp4wfp4_pre_quant,
@@ -1507,13 +1509,14 @@ class DeepseekV2AttentionMLA(nn.Module):
q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split(
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
)
q_lora = self.q_a_layernorm(q)
q = self.q_b_proj(q_lora)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
# NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
if self.use_nsa:
q_lora = self.q_a_layernorm(q)
q = self.q_b_proj(q_lora)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
_ = self.indexer(
x=hidden_states,
q_lora=q_lora,
@@ -1522,6 +1525,26 @@ class DeepseekV2AttentionMLA(nn.Module):
layer_id=self.layer_id,
return_indices=False,
)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
q, _, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
else:
q = self.q_a_layernorm(q)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
else:
q = self.q_proj(hidden_states)[0].view(
-1, self.num_local_heads, self.qk_head_dim
@@ -1531,7 +1554,31 @@ class DeepseekV2AttentionMLA(nn.Module):
_, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
latent_cache = latent_cache.unsqueeze(1)
kv_a = self.kv_a_layernorm(kv_a)
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
kv_a,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
None,
None,
None,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=True, # return unqaunt kv_a
)
kv = self.kv_b_proj(
kv_a_quanted,
)[0]
else:
kv_a = self.kv_a_layernorm(kv_a)
kv = self.kv_b_proj(kv_a)[0]
# kv_a = self.kv_a_layernorm(kv_a)
k_pe = latent_cache[:, :, self.kv_lora_rank :]
if self.rotary_emb is not None:
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
@@ -1617,8 +1664,27 @@ class DeepseekV2AttentionMLA(nn.Module):
self.kv_a_layernorm.variance_epsilon,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
q, _, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
# q_lora needed by indexer
if self.use_nsa:
@@ -1946,8 +2012,27 @@ class DeepseekV2AttentionMLA(nn.Module):
self.kv_a_layernorm.variance_epsilon,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
q, _, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
q_lora = q.clone() # required for topk_indices
k_nope = k_nope.unsqueeze(1)
@@ -2673,12 +2758,29 @@ class DeepseekV2DecoderLayer(nn.Module):
) -> torch.Tensor:
quant_format = (
"mxfp4"
if _is_gfx95_supported
and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None) is not None
and getattr(self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None)
is not None
and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype == torch.uint8
else ""
if (
_is_gfx95_supported
and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
is not None
and getattr(self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None)
is not None
and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype == torch.uint8
)
else (
"fp8"
if (
_is_gfx95_supported
and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
is not None
and getattr(
self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None
)
is not None
and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype
== getattr(torch, "float8_e4m3fn", None)
)
else ""
)
)
hidden_states, residual = self.layer_communicator.prepare_attn(