diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 5db8967c4..aa21283fe 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -44,6 +44,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( dispatch_w8a8_block_fp8_linear, input_to_float8, normalize_e4m3fn_to_e4m3fnuz, + requant_weight_ue8m0_inplace, ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod from sglang.srt.layers.quantization.marlin_utils_fp8 import ( @@ -261,6 +262,7 @@ class Fp8LinearMethod(LinearMethodBase): layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition layer.orig_dtype = params_dtype + layer.executed_weight_requant_ue8m0 = False # WEIGHT weight_dtype = ( @@ -347,7 +349,34 @@ class Fp8LinearMethod(LinearMethodBase): ) return else: + # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0 + from sglang.srt.layers.quantization.fp8_utils import ( + deepgemm_w8a8_block_fp8_linear_with_fallback, + ) + from sglang.srt.model_loader.utils import ( + should_deepgemm_weight_requant_ue8m0, + ) + + if ( + should_deepgemm_weight_requant_ue8m0( + weight_block_size=getattr( + self.quant_config, "weight_block_size", None + ), + ) + and ( + self.w8a8_block_fp8_linear + is deepgemm_w8a8_block_fp8_linear_with_fallback + ) + and (not layer.executed_weight_requant_ue8m0) + ): + requant_weight_ue8m0_inplace( + layer.weight, + layer.weight_scale_inv, + self.quant_config.weight_block_size, + ) + layer.executed_weight_requant_ue8m0 = True weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data + layer.weight.data = weight.data layer.weight_scale_inv.data = weight_scale.data else: diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index ea704e5ec..e1b341cef 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -220,7 +220,6 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ) self.logits_processor = LogitsProcessor(config) - self._executed_weight_requant_ue8m0 = False @torch.no_grad() def forward( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 65e5436bd..8f4badbdd 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -108,7 +108,6 @@ from sglang.srt.layers.quantization.fp8_kernel import ( per_token_group_quant_mla_deep_gemm_masked_fp8, ) from sglang.srt.layers.quantization.fp8_utils import ( - ENABLE_FLASHINFER_FP8_GEMM, block_quant_dequant, block_quant_to_tensor_quant, channel_quant_to_tensor_quant, @@ -3419,7 +3418,6 @@ class DeepseekV2ForCausalLM(nn.Module): } ) self.capture_aux_hidden_states = False - self._executed_weight_requant_ue8m0 = False self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() if self.nsa_enable_prefill_cp: @@ -3594,13 +3592,14 @@ class DeepseekV2ForCausalLM(nn.Module): weight = w weight_scale = self_attn.kv_b_proj.weight_scale_inv + # In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading. if ( should_deepgemm_weight_requant_ue8m0( weight_block_size=getattr( self.quant_config, "weight_block_size", None ) ) - and self._executed_weight_requant_ue8m0 + and self_attn.kv_b_proj.executed_weight_requant_ue8m0 ): weight_scale = inverse_transform_scale_ue8m0( weight_scale, mn=weight.shape[-2] @@ -3716,27 +3715,15 @@ class DeepseekV2ForCausalLM(nn.Module): self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous()) self_attn.use_deep_gemm_bmm = True - if ( - not ENABLE_FLASHINFER_FP8_GEMM - and should_deepgemm_weight_requant_ue8m0( - weight_block_size=getattr(self.quant_config, "weight_block_size", None) - ) - and not self._executed_weight_requant_ue8m0 - ): - self._executed_weight_requant_ue8m0 = True - self._weight_requant_ue8m0(is_nextn) - - # TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading - if ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 - and get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN") - ): - self._transform_scale_ue8m0(is_nextn) + # Requant the weights and scales of MoE layers + if get_moe_runner_backend().is_deep_gemm(): + self._maybe_moe_weight_requant_ue8m0(is_nextn) if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config): self._transform_scale_nextn_moe_ue8m0() - def _weight_requant_ue8m0(self, is_nextn=False): + def _maybe_moe_weight_requant_ue8m0(self, is_nextn=False): + # Dense fp8 layers will be processed in Fp8LinearMethod.process_weights_after_loading + # So we only need to process sparse MoE layers here weight_block_size = self.quant_config.weight_block_size moe_layers = list( @@ -3755,70 +3742,15 @@ class DeepseekV2ForCausalLM(nn.Module): else: layer = self.model.layers[layer_id] - module_list = [ - layer.self_attn.kv_b_proj, - layer.self_attn.o_proj, - ] - - if self.config.q_lora_rank is not None: - module_list.append(layer.self_attn.fused_qkv_a_proj_with_mqa) - module_list.append(layer.self_attn.q_b_proj) - else: - module_list.append(layer.self_attn.kv_a_proj_with_mqa) - module_list.append(layer.self_attn.q_proj) - - for module in module_list: - requant_weight_ue8m0_inplace( - module.weight, module.weight_scale_inv, weight_block_size - ) - if layer_id in moe_layers or is_nextn: - shared_experts = getattr(layer.mlp, "shared_experts", None) - if shared_experts is not None: - for module in [ - shared_experts.gate_up_proj, - shared_experts.down_proj, - ]: - requant_weight_ue8m0_inplace( - module.weight, module.weight_scale_inv, weight_block_size - ) - experts = layer.mlp.experts + # TODO: move this logic to Fp8MoEMethod.process_weights_after_loading if isinstance(experts, DeepEPMoE): for w in [ (experts.w13_weight, experts.w13_weight_scale_inv), (experts.w2_weight, experts.w2_weight_scale_inv), ]: requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size) - else: - mlp = layer.mlp - assert isinstance(mlp, DeepseekV2MLP) - for module in [ - mlp.gate_up_proj, - mlp.down_proj, - ]: - requant_weight_ue8m0_inplace( - module.weight, module.weight_scale_inv, weight_block_size - ) - - # TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading - def _transform_scale_ue8m0(self, is_nextn=False): - num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers - - for layer_id in range(num_hidden_layers): - if is_nextn: - layer = self.model.decoder - else: - layer = self.model.layers[layer_id] - - module_list = [] - if self.config.q_lora_rank is not None: - module_list.append(layer.self_attn.q_b_proj) - - for module in module_list: - transform_scale_ue8m0_inplace( - module.weight_scale_inv, mn=module.weight.shape[-2] - ) # TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0) def _transform_scale_nextn_moe_ue8m0(self): diff --git a/python/sglang/test/test_block_fp8_deep_gemm_blackwell.py b/python/sglang/test/test_block_fp8_deep_gemm_blackwell.py index ac7239ea0..833c23e7c 100644 --- a/python/sglang/test/test_block_fp8_deep_gemm_blackwell.py +++ b/python/sglang/test/test_block_fp8_deep_gemm_blackwell.py @@ -222,7 +222,7 @@ class TestDeepGemmBlackwell(CustomTestCase): with torch.inference_mode(): ref_out = native_w8a8_block_fp8_matmul( - A_q, B_q, A_s, B_s, block_size, out_dtype + A_qu[0], B_qu[0], A_qu[1], B_qu[1], block_size, out_dtype ) out = torch.empty_like(ref_out) fp8_gemm_nt(A_qu, B_qu, out)