diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index c9861c445..cac206abc 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -87,6 +87,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_FLASHINFER_FP4_GEMM_BACKEND` | Select backend for `mm_fp4` on Blackwell GPUS | `` | | `SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN` | Quantize q_b_proj from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_MOE_NVFP4_DISPATCH` | Use nvfp4 for moe dispatch (on flashinfer_cutlass or flashinfer_cutedsl moe runner backend) | `"false"` | +| `SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE` | Quantize moe of nextn layer from BF16 to FP8 when launching DeepSeek NVFP4 checkpoint | `false` | | `SGLANG_ENABLE_FLASHINFER_FP8_GEMM` (deprecated) | Use flashinfer kernels when running blockwise fp8 GEMM on Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=flashinfer_trtllm` instead. | `false` | | `SGLANG_SUPPORT_CUTLASS_BLOCK_FP8` (deprecated) | Use Cutlass kernels when running blockwise fp8 GEMM on Hopper or Blackwell GPUs. **DEPRECATED**: Please use `--fp8-gemm-backend=cutlass` instead. | `false` | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5bf021bd5..dc12cb37e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -253,6 +253,7 @@ class Envs: SGLANG_MOE_NVFP4_DISPATCH = EnvBool(False) SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN = EnvBool(False) SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2 = EnvBool(False) + SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE = EnvBool(False) # Flashinfer SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index e7f2b2a28..e80e818a0 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -667,6 +667,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): ), requires_grad=False, ) + # w13_weight and w2_weight are always requanted together + w13_weight_scale.format_ue8m0 = False + w2_weight_scale.format_ue8m0 = False layer.register_parameter("w13_weight_scale_inv", w13_weight_scale) layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) assert self.quant_config.activation_scheme == "dynamic" @@ -814,6 +817,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): ), ) and get_moe_runner_backend().is_deep_gemm() + and not layer.w13_weight_scale_inv.format_ue8m0 ): assert isinstance( layer, DeepEPMoE @@ -825,6 +829,8 @@ class Fp8MoEMethod(FusedMoEMethodBase): requant_weight_ue8m0_inplace( layer.w2_weight, layer.w2_weight_scale_inv, weight_block_size ) + layer.w13_weight_scale_inv.format_ue8m0 = True + layer.w2_weight_scale_inv.format_ue8m0 = True return # If checkpoint is fp16 or bfloat16, quantize in place. diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index bbd4af816..5de69a9e3 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -71,12 +71,12 @@ class DeepseekModelNextN(nn.Module): super().__init__() if enable_nextn_moe_bf16_cast_to_fp8(quant_config): # refer to real DeepSeek V3 quant config - moe_quant_config = Fp8Config( + moe_quant_config_override = Fp8Config( is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128], ) else: - moe_quant_config = None + moe_quant_config_override = None if quant_config is not None and quant_config.get_name() == "modelopt_fp4": logger.warning( @@ -115,7 +115,7 @@ class DeepseekModelNextN(nn.Module): config, 0, quant_config=quant_config, - moe_quant_config=moe_quant_config, + moe_quant_config_override=moe_quant_config_override, is_nextn=True, prefix=add_prefix(layer_name, prefix), alt_stream=self.alt_stream, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b0e6ec1af..8c3f15719 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -94,7 +94,7 @@ from sglang.srt.layers.moe import ( get_moe_runner_backend, should_use_flashinfer_cutlass_moe_fp4_allgather, ) -from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.token_dispatcher.base import ( @@ -119,7 +119,6 @@ from sglang.srt.layers.quantization.fp8_utils import ( inverse_transform_scale_ue8m0, normalize_e4m3fn_to_e4m3fnuz, quant_weight_ue8m0, - transform_scale_ue8m0_inplace, ) from sglang.srt.layers.quantization.int8_utils import ( block_dequant as int8_block_dequant, @@ -232,11 +231,16 @@ _is_cublas_ge_129 = is_nvidia_cublas_cu12_version_ge_12_9() logger = logging.getLogger(__name__) +# Optional quantization for DeepSeek nvfp4 checkpoint +NVFP4_CKPT_FP8_ATTN_QUANT_MODULES = ["q_b_proj"] + + def enable_nextn_moe_bf16_cast_to_fp8(quant_config): return ( - quant_config is not None + envs.SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE.get() + and quant_config is not None and quant_config.get_name() == "modelopt_fp4" - and get_moe_a2a_backend().is_deepep() + and get_moe_runner_backend().is_deep_gemm() ) @@ -2742,7 +2746,7 @@ class DeepseekV2DecoderLayer(nn.Module): config: PretrainedConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, - moe_quant_config: Optional[QuantizationConfig] = None, + moe_quant_config_override: Optional[QuantizationConfig] = None, is_nextn: bool = False, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, @@ -2795,7 +2799,7 @@ class DeepseekV2DecoderLayer(nn.Module): if self.is_layer_sparse: self.mlp = DeepseekV2MoE( config=config, - quant_config=moe_quant_config or quant_config, + quant_config=moe_quant_config_override or quant_config, prefix=add_prefix("mlp", prefix), layer_id=self.layer_id, alt_stream=alt_stream, @@ -3621,47 +3625,6 @@ 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 is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config): - self._transform_scale_nextn_moe_ue8m0() - - # TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0) - def _transform_scale_nextn_moe_ue8m0(self): - layer = self.model.decoder - - 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, - ]: - transform_scale_ue8m0_inplace( - module.weight_scale_inv, mn=module.weight.shape[-2] - ) - - experts = layer.mlp.experts - w13_weight_fp8 = ( - experts.w13_weight, - ( - experts.w13_weight_scale_inv - if hasattr(experts, "w13_weight_scale_inv") - else experts.w13_weight_scale - ), - ) - w2_weight_fp8 = ( - experts.w2_weight, - ( - experts.w2_weight_scale_inv - if hasattr(experts, "w2_weight_scale_inv") - else experts.w2_weight_scale - ), - ) - if isinstance(experts, DeepEPMoE): - for w in [ - w13_weight_fp8, - w2_weight_fp8, - ]: - transform_scale_ue8m0_inplace(w[1], mn=w[0].shape[-2]) - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): if is_nextn: @@ -3677,12 +3640,9 @@ class DeepseekV2ForCausalLM(nn.Module): else: raise ValueError("num_nextn_predict_layers is not in the config") - if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get(): - weights = self._quant_attn_to_fp8_ue8m0(weights, is_nextn=is_nextn) - if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config): - weights = self._quant_nextn_moe_to_fp8_ue8m0( - weights, nextn_layer_id=nextn_layer_id - ) + weights = self._maybe_quant_weights_to_fp8_ue8m0( + weights, NVFP4_CKPT_FP8_ATTN_QUANT_MODULES, is_nextn + ) stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -3945,62 +3905,6 @@ class DeepseekV2ForCausalLM(nn.Module): self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) - def _quant_attn_to_fp8_ue8m0(self, weights, is_nextn): - weights_dict = dict(weights) - - # temporarily only support DeepSeek V3/R1 - weight_block_size = [128, 128] - - for layer_id in tqdm.trange( - self.config.num_hidden_layers + int(is_nextn), - desc="quant attn to fp8 ue8m0", - ): - for stem in [ - # may put tensors like `o_proj` here for DeepSeek FP4 ckpt v1 - "q_b_proj", - ]: - partial_name = f"model.layers.{layer_id}.self_attn.{stem}" - original_weight = weights_dict[f"{partial_name}.weight"] - out_w, out_s = quant_weight_ue8m0( - original_weight, weight_block_size=weight_block_size - ) - weights_dict[f"{partial_name}.weight"] = out_w - weights_dict[f"{partial_name}.weight_scale_inv"] = out_s - - return list(weights_dict.items()) - - # TODO avoid code dup - def _quant_nextn_moe_to_fp8_ue8m0(self, weights, nextn_layer_id: int): - weights_dict = dict(weights) - - # temporarily only support DeepSeek V3/R1 - weight_block_size = [128, 128] - - for layer_id in [nextn_layer_id]: - for expert_sub_name in [ - "shared_experts", - *[ - f"experts.{expert_id}" - for expert_id in range(self.config.n_routed_experts) - ], - ]: - for stem in [ - "gate_proj", - "up_proj", - "down_proj", - ]: - partial_name = ( - f"model.layers.{layer_id}.mlp.{expert_sub_name}.{stem}" - ) - original_weight = weights_dict[f"{partial_name}.weight"] - out_w, out_s = quant_weight_ue8m0( - original_weight, weight_block_size=weight_block_size - ) - weights_dict[f"{partial_name}.weight"] = out_w - weights_dict[f"{partial_name}.weight_scale_inv"] = out_s - - return list(weights_dict.items()) - def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight @@ -4034,6 +3938,76 @@ class DeepseekV2ForCausalLM(nn.Module): # of the (i-1)th layer as aux hidden state self.model.layers_to_capture = [val + 1 for val in layer_ids] + # Mark the ue8m0 flag of nextn moe weights as True to avoid requantization + def _mark_nextn_moe_weights_as_ue8m0(self): + experts = self.model.decoder.mlp.experts + w13_scale = ( + experts.w13_weight_scale_inv + if hasattr(experts, "w13_weight_scale_inv") + else experts.w13_weight_scale + ) + w2_scale = ( + experts.w2_weight_scale_inv + if hasattr(experts, "w2_weight_scale_inv") + else experts.w2_weight_scale + ) + w13_scale.format_ue8m0 = True + w2_scale.format_ue8m0 = True + + def _maybe_quant_weights_to_fp8_ue8m0( + self, weights, attn_quant_modules, is_nextn=False + ): + # Quantize some weights to fp8 ue8m0 for DeepSeek nvfp4 checkpoint + partial_names = [] + nextn_layer_id = ( + 0 if self.config.num_hidden_layers == 1 else self.config.num_hidden_layers + ) + weights_dict = dict(weights) + weight_block_size = [128, 128] + + if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get(): + layer_ids = ( + list(range(self.config.num_hidden_layers)) + if not is_nextn + else [nextn_layer_id] + ) + for layer_id in layer_ids: + for stem in attn_quant_modules: + partial_names.append(f"model.layers.{layer_id}.self_attn.{stem}") + + if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config): + for expert_sub_name in [ + "shared_experts", + *[ + f"experts.{expert_id}" + for expert_id in range(self.config.n_routed_experts) + ], + ]: + for stem in [ + "gate_proj", + "up_proj", + "down_proj", + ]: + partial_names.append( + f"model.layers.{nextn_layer_id}.mlp.{expert_sub_name}.{stem}" + ) + + for partial_name in tqdm.tqdm( + partial_names, + desc="quant weights to fp8 ue8m0", + ): + original_weight = weights_dict[f"{partial_name}.weight"] + out_w, out_s = quant_weight_ue8m0( + original_weight, weight_block_size=weight_block_size + ) + weights_dict[f"{partial_name}.weight"] = out_w + weights_dict[f"{partial_name}.weight_scale_inv"] = out_s + + if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config): + self._mark_nextn_moe_weights_as_ue8m0() + + return list(weights_dict.items()) + AttentionBackendRegistry.register("ascend", handle_attention_ascend) AttentionBackendRegistry.register("flashinfer", handle_attention_flashinfer) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index abb3fde8d..d506110ac 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1138,6 +1138,27 @@ class ServerArgs: "Use flashinfer_trtllm as MoE runner backend on sm100 for DeepseekV3ForCausalLM" ) + if ( + self.quantization == "modelopt_fp4" + and self.speculative_algorithm == "EAGLE" + and ( + self.speculative_moe_runner_backend is None + or self.speculative_moe_a2a_backend is None + ) + ): + if envs.SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE.get(): + self.speculative_moe_runner_backend = "deep_gemm" + self.speculative_moe_a2a_backend = "deepep" + logger.info( + "Use deep_gemm moe runner and deepep a2a backend for bf16 nextn layer in deepseek fp4 checkpoint." + ) + else: + self.speculative_moe_runner_backend = "triton" + self.speculative_moe_a2a_backend = "none" + logger.info( + "Use triton fused moe by default for bf16 nextn layer in deepseek fp4 checkpoint." + ) + elif model_arch in ["GptOssForCausalLM"]: # Set attention backend for GPT-OSS if self.is_attention_backend_not_set():