From 4a78031a71ddbce7b1e0fc8d2a1e9105eb1e145b Mon Sep 17 00:00:00 2001 From: yctseng0211 Date: Tue, 11 Nov 2025 18:59:10 +0800 Subject: [PATCH] [ROCM] Optimized deepseek-r1 model with rmsnorm + fp8 quant fusion (#12689) should be clean after https://github.com/sgl-project/sglang/pull/13017 landed --- python/sglang/srt/layers/communicator.py | 41 ++++- python/sglang/srt/layers/quantization/fp8.py | 10 ++ .../srt/layers/quantization/fp8_utils.py | 21 ++- python/sglang/srt/models/deepseek_v2.py | 132 ++++++++++++++-- .../quant/test_fused_rms_fp8_group_quant.py | 148 ++++++++++++++++++ test/srt/run_suite.py | 1 + 6 files changed, 331 insertions(+), 22 deletions(-) create mode 100644 test/srt/quant/test_fused_rms_fp8_group_quant.py diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 2f276b319..98d8abb99 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 37f15ae2c..78dfbf711 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ec2c5f05b..ff2d3c2c0 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -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( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c74698ba7..021173194 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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( diff --git a/test/srt/quant/test_fused_rms_fp8_group_quant.py b/test/srt/quant/test_fused_rms_fp8_group_quant.py new file mode 100644 index 000000000..443c590e9 --- /dev/null +++ b/test/srt/quant/test_fused_rms_fp8_group_quant.py @@ -0,0 +1,148 @@ +# test/srt/quant/test_fused_rms_fp8_group_quant.py +import itertools +import unittest + +import torch +import torch.nn.functional as F + +from sglang.test.test_utils import CustomTestCase + + +def _fp8_available() -> bool: + # requirement：1) GPU；2) ROCm；3) torch support float8_e4m3fn + if not torch.cuda.is_available(): + return False + if getattr(torch.version, "hip", None) is None: + return False + return hasattr(torch, "float8_e4m3fn") + + +def _rmsnorm(x, weight, eps=1e-6): + # row-wise RMSNorm + row_norm = (x * x).sum(dim=-1) + norm = torch.rsqrt(row_norm / x.shape[1] + eps) + return x * norm[:, None] * weight[None, :] + + +def _per_token_fp8_group_quant(x, dtype_quant, group_size=128): + """per token、group-size quant, return (quantized, scale)。""" + DTYPE_MAX = torch.finfo(dtype_quant).max + M, N = x.shape + + pad = (group_size - (N % group_size)) % group_size + if pad: + x_reshape = F.pad(x, (0, pad, 0, 0), "constant", 0) + else: + x_reshape = x + + G = (N + group_size - 1) // group_size + x_reshape = x_reshape.view(M, G, group_size).to(torch.float32) + x_max = torch.max(torch.abs(x_reshape), dim=-1, keepdim=True)[0].clamp_min_(1e-10) + x_scale = x_max / DTYPE_MAX + inv = 1.0 / x_scale + + x_q = torch.clamp(x_reshape * inv, -DTYPE_MAX, DTYPE_MAX).to(dtype_quant) + x_q = x_q.view(M, G * group_size) + if pad: + x_q = x_q[:, :N] + x_scale = x_scale.squeeze(-1) # [M, G] + return x_q, x_scale + + +def _upcast_fp8_group(x_q, x_s, out_dtype=torch.float32, group_size=128): + """unqaunt""" + M, N = x_q.shape + G = (N + group_size - 1) // group_size + pad = (group_size - (N % group_size)) % group_size + + if pad: + x_q = F.pad(x_q, (0, pad, 0, 0), "constant", 0) + + x_q = x_q.view(M, G, group_size).to(torch.float32) + x = x_q * x_s.view(M, G, 1) + x = x.view(M, G * group_size)[:, :N] + return x.to(out_dtype) + + +class TestFusedRMSFP8GroupQuant(CustomTestCase): + # + DTYPES = [torch.bfloat16, torch.float16] + # (M, N1, N2) + SHAPES = [(32, 128, 7168), (128, 7168, 7168)] + GROUP_SIZE = [128] + SEEDS = [0] + + @classmethod + def setUpClass(cls): + if not _fp8_available(): + raise unittest.SkipTest("Skip: ROCm/FP8 is not available") + torch.set_default_device("cuda") + + def _run_ref(self, x1, w1, eps1, x2, w2, eps2, res1, dtype_quant, group_size): + s = x1 + (res1 if res1 is not None else 0) + y1 = _rmsnorm(s, w1, eps1) + y2 = _rmsnorm(x2, w2, eps2) if x2 is not None else None + y1_q, y1_s = _per_token_fp8_group_quant(y1, dtype_quant, group_size) + return ( + (y1_q, y1_s), + y1.to(x1.dtype), + (y2.to(x1.dtype) if y2 is not None else None), + (s.to(x1.dtype) if res1 is not None else None), + ) + + def _case(self, M, N1, N2, group_size, dtype, seed): + torch.manual_seed(seed) + fp8 = torch.float8_e4m3fn + device = "cuda" + + x1 = torch.randn(M, N1, dtype=dtype, device=device) / 10 + x2 = torch.randn(M, N2, dtype=dtype, device=device) / 10 + w1 = torch.ones(N1, dtype=torch.float32, device=device) + w2 = torch.ones(N2, dtype=torch.float32, device=device) + res1 = torch.randn(M, N1, dtype=dtype, device=device) / 10 + + # ref + (y1_q_ref, y1_s_ref), y1_ref, y2_ref, s_ref = self._run_ref( + x1, w1, 1e-6, x2, w2, 1e-6, res1, fp8, group_size + ) + + # be tested：aiter fused op + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant + + (y1_q, y1_s), y1, y2, s = fused_rms_fp8_group_quant( + x1, + w1, + 1e-6, + inp2=x2, + inp2_weight=w2, + inp2_epsilon=1e-6, + group_size=group_size, + dtype_quant=fp8, + res1=res1, + output_unquantized_inp1=True, # get unqaunt y1 + ) + + torch.testing.assert_close(y1, y1_ref, atol=0.1, rtol=0.1) + torch.testing.assert_close(y2, y2_ref, atol=0.1, rtol=0.1) + torch.testing.assert_close(s, s_ref, atol=0.1, rtol=0.1) + + # check unquant + y1_up_ref = _upcast_fp8_group( + y1_q_ref, y1_s_ref, out_dtype=torch.float32, group_size=group_size + ) + y1_up = _upcast_fp8_group( + y1_q, y1_s, out_dtype=torch.float32, group_size=group_size + ) + torch.testing.assert_close(y1_up, y1_up_ref, atol=0.1, rtol=0.1) + + def test_fused_rms_fp8_group_quant(self): + for params in itertools.product( + self.SHAPES, self.GROUP_SIZE, self.DTYPES, self.SEEDS + ): + (M, N1, N2), g, dtype, seed = params + with self.subTest(M=M, N1=N1, N2=N2, group_size=g, dtype=dtype, seed=seed): + self._case(M, N1, N2, g, dtype, seed) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 6c8efc5b3..79de8664a 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -421,6 +421,7 @@ suite_amd = { TestFile("openai_server/validation/test_request_length_validation.py", 31), TestFile("quant/test_awq_dequant.py", 2), TestFile("quant/test_block_int8.py", 22), + TestFile("quant/test_fused_rms_fp8_group_quant.py", 10), TestFile("rl/test_update_weights_from_disk.py", 210), TestFile("test_abort.py", 51), TestFile("test_chunked_prefill.py", 410),