remove special case of NVFP4 static quantizer in general export logic

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

modelopt/torch/export/quant_utils.py

@@ -300,18 +300,15 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") ->
 
     quantization_format = get_quantization_format(module)
 
-    # Handle NVFP4 variants (static or dynamic)
-    is_nvfp4_static = isinstance(weight_quantizer, NVFP4StaticQuantizer)
-    if is_nvfp4_static or quantization_format in [
+    if quantization_format in [
         QUANTIZATION_NVFP4,
         QUANTIZATION_NVFP4_AWQ,
         QUANTIZATION_NVFP4_SVDQUANT,
         QUANTIZATION_W4A8_NVFP4_FP8,
     ]:
-        # Calibrate weight quantizer if amax is not set (only needed for dynamic quantizers)
-        if not is_nvfp4_static:
-            module_name = f"{type(module).__name__}.{weight_name}"
-            _ensure_weight_quantizer_calibrated(weight_quantizer, weight, module_name)
+        # Calibrate weight quantizer if amax is not set
+        module_name = f"{type(module).__name__}.{weight_name}"
+        _ensure_weight_quantizer_calibrated(weight_quantizer, weight, module_name)
 
         if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8:
             # weight_scaling_factor_2 for w4a8 needs to be amax/448, so that the wsf is in range 448/6.
@@ -347,18 +344,16 @@ def get_weight_scaling_factor_2(module: nn.Module, weight_name: str = "weight")
 
     quantization_format = get_quantization_format(module)
 
-    is_nvfp4_static = isinstance(weight_quantizer, NVFP4StaticQuantizer)
-    if is_nvfp4_static or quantization_format in [
+    if quantization_format in [
         QUANTIZATION_NVFP4,
         QUANTIZATION_NVFP4_AWQ,
         QUANTIZATION_NVFP4_SVDQUANT,
         QUANTIZATION_W4A8_NVFP4_FP8,
     ]:
-        # Calibrate weight quantizer if amax is not set (only needed for dynamic quantizers)
-        if not is_nvfp4_static:
-            weight = getattr(module, weight_name)
-            module_name = f"{type(module).__name__}.{weight_name}"
-            _ensure_weight_quantizer_calibrated(weight_quantizer, weight, module_name)
+        # Calibrate weight quantizer if amax is not set
+        weight = getattr(module, weight_name)
+        module_name = f"{type(module).__name__}.{weight_name}"
+        _ensure_weight_quantizer_calibrated(weight_quantizer, weight, module_name)
 
         if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8:
             # weight_scaling_factor_2 for w4a8 needs to be amax/448, so that the wsf is in range 448/6.

modelopt/torch/export/unified_export_hf.py

@@ -495,6 +495,11 @@ def _export_quantized_weight(
         expert_type in type(sub_module).__name__
         for expert_type in ["Llama4TextExperts", "GptOssExperts"]
     )
+    if is_bmm_expert_weight and isinstance(weight_quantizer, NVFP4StaticQuantizer):
+        warnings.warn(
+            "NVFP4StaticQuantizer with BMM-style expert weights (e.g. Llama4TextExperts, "
+            "GptOssExperts) is not yet supported; export may produce incorrect results."
+        )
 
     if quantization_format in [
         QUANTIZATION_NVFP4,
@@ -507,17 +512,11 @@ def _export_quantized_weight(
             weight, is_bmm_expert_weight=is_bmm_expert_weight
         )
 
-        # Check if this is a static NVFP4 quantizer (has pre-computed scales from MSE calibration)
-        # For static NVFP4, weight_scale is already computed from static _amax values in get_weight_scaling_factor
-        is_nvfp4_static = isinstance(weight_quantizer, NVFP4StaticQuantizer)
-
-        if not is_nvfp4_static:
-            # For dynamic NVFP4, compute scales from weights
-            weight_scale = NVFP4QTensor.get_weights_scaling_factor(
-                weight,
-                block_size=block_size,
-                weights_scaling_factor_2=weight_scale_2,
-            )[0]
+        weight_scale = NVFP4QTensor.get_weights_scaling_factor(
+            weight,
+            block_size=block_size,
+            weights_scaling_factor_2=weight_scale_2,
+        )[0]
 
         quantized_weight = to_quantized_weight(
             weight.to(dtype),