diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 6903104ac..bf27def43 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -534,11 +534,16 @@ class Fp8LinearMethod(LinearMethodBase): should_deepgemm_weight_requant_ue8m0, ) + # Only requantize to UE8M0 if DeepGEMM can actually run + # this layer. If the dtype or shape is unsupported, the GEMM + # falls back to triton at runtime, which needs float32 scales. if ( should_deepgemm_weight_requant_ue8m0( weight_block_size=getattr( self.quant_config, "weight_block_size", None ), + output_dtype=getattr(layer, "orig_dtype", None), + weight_shape=layer.weight.shape, ) and ( self.w8a8_block_fp8_linear diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 11b289475..51dc933d1 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -496,8 +496,11 @@ def flashinfer_gemm_w8a8_block_fp8_linear_with_fallback( input_2d = input.view(-1, input.shape[-1]) backend = _get_flashinfer_groupwise_backend() - # TRTLLM backend requires K dimension >= 256. - if backend == "trtllm" and input_2d.shape[1] < 256: + # TRTLLM backend requires K >= 256 and weight scales in UE8M0/R128c4 + # packed format. Fall back to triton when scales are plain float32. + if backend == "trtllm" and ( + input_2d.shape[1] < 256 or not getattr(weight_scale, "format_ue8m0", False) + ): return triton_w8a8_block_fp8_linear( input, weight, block_size, weight_scale, input_scale, bias ) diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index 3eb70cedf..1201e7553 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -249,13 +249,41 @@ def get_architecture_class_name(model_config: ModelConfig) -> str: return get_model_architecture(model_config)[1] -def should_deepgemm_weight_requant_ue8m0(weight_block_size): - """Should we requant fp8 weights into UE8M0 format when loading the model""" - return ( +def post_load_weights(model: nn.Module, model_config: ModelConfig): + # Model weight loading consists of two stages: + # 1. Initial weight loading. + # 2. Post-processing of weights, including assigning specific member variables. + # For `dummy_init`, only the second stage is required. + if hasattr(model, "post_load_weights"): + if model_config.hf_config.architectures[0] == "DeepseekV3ForCausalLMNextN": + model.post_load_weights(is_nextn=True) + else: + model.post_load_weights() + + +def should_deepgemm_weight_requant_ue8m0( + weight_block_size, output_dtype=None, weight_shape=None +): + """Should we requant fp8 weights into UE8M0 format when loading the model. + + When output_dtype or weight_shape are provided, also checks that DeepGEMM + can actually run this layer at runtime (bf16 output, N%64==0, K%128==0). + Without these checks, scales would be converted to UE8M0 but the GEMM would + fall back to triton which expects float32 scales, causing wrong results. + """ + if not ( deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 and weight_block_size is not None - ) + ): + return False + if output_dtype is not None and output_dtype != torch.bfloat16: + return False + if weight_shape is not None and ( + weight_shape[0] % 64 != 0 or weight_shape[1] % 128 != 0 + ): + return False + return True def should_async_load(weight: torch.Tensor) -> bool: