[Fix] Restore online MXFP8 quantization for linear layers (#32953)

This commit is contained in:
Brayden Zhong
2026-07-31 06:42:06 +00:00
committed by GitHub
parent 5c6635d8f3
commit f94d2c5663
2 changed files with 9 additions and 11 deletions
@@ -617,10 +617,6 @@ class Fp8LinearMethod(LinearMethodBase):
layer.register_parameter("input_scale", scale) layer.register_parameter("input_scale", scale)
else: else:
layer.register_parameter("input_scale", None) layer.register_parameter("input_scale", None)
elif use_mxfp8:
raise ValueError(
"MXFP8 requires fp8-serialized checkpoint for linear layers."
)
def create_weights( def create_weights(
self, self,
+6 -4
View File
@@ -261,15 +261,17 @@ class NemotronHMoE(nn.Module):
self.fc1_latent_proj = None self.fc1_latent_proj = None
self.fc2_latent_proj = None self.fc2_latent_proj = None
self.use_min_latency_fc1_gemm = ( 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 is None:
self._use_min_latency_fc1_gemm = (
self.use_latent_moe self.use_latent_moe
and self.fc1_latent_proj is not None and self.fc1_latent_proj is not None
and _is_cuda and _is_cuda
and fused_a_gemm_weight_eligible(self.fc1_latent_proj) and fused_a_gemm_weight_eligible(self.fc1_latent_proj)
) )
if self._use_min_latency_fc1_gemm:
def _apply_fc1_latent_proj(self, hidden_states: torch.Tensor) -> torch.Tensor:
if self.use_min_latency_fc1_gemm:
return linear_with_fused_a_gemm(self.fc1_latent_proj, hidden_states) return linear_with_fused_a_gemm(self.fc1_latent_proj, hidden_states)
return self.fc1_latent_proj(hidden_states)[0] return self.fc1_latent_proj(hidden_states)[0]