From f94d2c566354657a5c3807737418203d9f443264 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Thu, 30 Jul 2026 23:42:06 -0700 Subject: [PATCH] [Fix] Restore online MXFP8 quantization for linear layers (#32953) --- python/sglang/srt/layers/quantization/fp8.py | 4 ---- python/sglang/srt/models/nemotron_h.py | 16 +++++++++------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 8852ed38a..26fc7b94a 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -617,10 +617,6 @@ class Fp8LinearMethod(LinearMethodBase): layer.register_parameter("input_scale", scale) else: layer.register_parameter("input_scale", None) - elif use_mxfp8: - raise ValueError( - "MXFP8 requires fp8-serialized checkpoint for linear layers." - ) def create_weights( self, diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index f678c5c32..71dc6005c 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -261,15 +261,17 @@ class NemotronHMoE(nn.Module): self.fc1_latent_proj = None self.fc2_latent_proj = None - self.use_min_latency_fc1_gemm = ( - self.use_latent_moe - and self.fc1_latent_proj is not None - and _is_cuda - and fused_a_gemm_weight_eligible(self.fc1_latent_proj) - ) + self._use_min_latency_fc1_gemm: bool | None = None def _apply_fc1_latent_proj(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.use_min_latency_fc1_gemm: + if self._use_min_latency_fc1_gemm is None: + self._use_min_latency_fc1_gemm = ( + self.use_latent_moe + and self.fc1_latent_proj is not None + and _is_cuda + and fused_a_gemm_weight_eligible(self.fc1_latent_proj) + ) + if self._use_min_latency_fc1_gemm: return linear_with_fused_a_gemm(self.fc1_latent_proj, hidden_states) return self.fc1_latent_proj(hidden_states)[0]