Address some comments

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>

Author: jingyu-ml

CHANGELOG.rst

@@ -15,6 +15,7 @@ Changelog
 - Enable PTQ workflow for the Step3.5-Flash MoE model with NVFP4 W4A4 + FP8 KV cache quantization. See `modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml <https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml>`_ for more details.
 - Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/vllm_serve#load-qatptq-model-and-serve-in-vllm-wip>`_ for more details.
 - [Early Testing] Add Claude Code PTQ skill (``.claude/skills/ptq/``) for agent-assisted post-training quantization. The skill guides the agent through environment detection, model support checking, format selection, and execution via the launcher or manual SLURM/Docker/bare GPU paths. Includes handling for unlisted models with custom module patching. This feature is in early testing — use with caution.
+- Add implicit GEMM CUDA kernel for Conv3D with fused NVFP4 fake quantization (``modelopt.torch.kernels.conv``). When NVFP4 quantization is applied to an ``nn.Conv3d`` layer via ModelOpt PTQ, the implicit GEMM path is used automatically instead of cuDNN. Uses BF16 WMMA tensor cores (SM80+) with FP32 accumulation and in-kernel FP4 (E2M1) activation quantization. Grouped convolution (``groups > 1``) falls back to the default cuDNN path. Inference only — training mode falls back to cuDNN with a warning.
 
 **Backward Breaking Changes**
 

examples/diffusers/quantization/quantize.py

@@ -116,41 +116,48 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any:
 
         if self.config.format == QuantFormat.INT8:
             if self.config.algo == QuantAlgo.SMOOTHQUANT:
-                quant_config = mtq.INT8_SMOOTHQUANT_CFG
+                base_cfg = mtq.INT8_SMOOTHQUANT_CFG
             else:
-                quant_config = INT8_DEFAULT_CONFIG
+                base_cfg = INT8_DEFAULT_CONFIG
             if self.config.collect_method != CollectMethod.DEFAULT:
                 reset_set_int8_config(
-                    quant_config,
+                    base_cfg,
                     self.config.percentile,
                     n_steps,
                     collect_method=self.config.collect_method.value,
                     backbone=backbone,
                 )
         elif self.config.format == QuantFormat.FP8:
-            quant_config = FP8_DEFAULT_CONFIG
+            base_cfg = FP8_DEFAULT_CONFIG
         elif self.config.format == QuantFormat.FP4:
             if self.model_config.model_type.value.startswith("flux"):
-                quant_config = NVFP4_FP8_MHA_CONFIG
+                base_cfg = NVFP4_FP8_MHA_CONFIG
             else:
-                quant_config = NVFP4_DEFAULT_CONFIG
-            # Override block size if non-default
-            if self.config.block_size != 16:
-                import copy
-
-                quant_config = copy.deepcopy(quant_config)
-                for entry in quant_config["quant_cfg"]:
-                    if isinstance(entry, dict) and "block_sizes" in entry.get("cfg", {}):
-                        entry["cfg"]["block_sizes"][-1] = self.config.block_size
+                base_cfg = NVFP4_DEFAULT_CONFIG
         else:
             raise NotImplementedError(f"Unknown format {self.config.format}")
+
+        # Build a fresh config dict so we never mutate the global constants.
+        quant_cfg_list = list(base_cfg["quant_cfg"])
+
+        if self.config.format == QuantFormat.FP4:
+            for i, entry in enumerate(quant_cfg_list):
+                if isinstance(entry, dict) and "block_sizes" in entry.get("cfg", {}):
+                    new_block_sizes = {**entry["cfg"]["block_sizes"], -1: self.config.block_size}
+                    quant_cfg_list[i] = {
+                        **entry,
+                        "cfg": {**entry["cfg"], "block_sizes": new_block_sizes},
+                    }
+
         if self.config.quantize_mha:
-            quant_config["quant_cfg"].append(
+            quant_cfg_list.append(
                 {
                     "quantizer_name": "*[qkv]_bmm_quantizer",
                     "cfg": {"num_bits": (4, 3), "axis": None},
                 }
             )
+
+        quant_config = {**base_cfg, "quant_cfg": quant_cfg_list}
         set_quant_config_attr(
             quant_config,
             self.model_config.trt_high_precision_dtype.value,
@@ -617,8 +624,8 @@ def main() -> None:
                 logger.info(f"Quantizing backbone: {backbone_name}")
                 backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone)
 
-                # Full pipeline inference for calibration — cached modules
-                # (transformer, video_decoder) are exercised during __call__
+                # Calibration runs the full pipeline (not just `mod`), so the
+                # closure intentionally ignores the backbone argument.
                 def forward_loop(mod):
                     calibrator.run_calibration(batched_prompts)
 

modelopt/torch/quantization/nn/modules/quant_conv.py

@@ -15,6 +15,8 @@
 
 """Quantized convolution."""
 
+import warnings
+
 import torch.nn as nn
 
 from ... import tensor_quant
@@ -90,34 +92,23 @@ class _QuantConv3d(QuantLinearConvBase):
 
     def _should_use_implicit_gemm(self):
         """Check if both quantizers are NVFP4 and the implicit GEMM kernel is available."""
-        if hasattr(self, "_use_implicit_gemm"):
-            return self._use_implicit_gemm
-        result = False
-        if (
+        return (
             hasattr(self, "input_quantizer")
             and hasattr(self, "weight_quantizer")
             and _is_nvfp4_quantizer(self.input_quantizer)
             and _is_nvfp4_quantizer(self.weight_quantizer)
             and self.groups == 1
-        ):
-            try:
-                from modelopt.torch.kernels.conv.implicit_gemm_cuda import (
-                    conv3d_implicit_gemm_cuda,  # noqa: F401
-                )
-
-                result = True
-            except ImportError:
-                pass
-        self._use_implicit_gemm = result
-        return result
+        )
 
     def _implicit_gemm_forward(self, input):
         """Run NVFP4 implicit GEMM kernel. Input may already be padded."""
         from modelopt.torch.kernels.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
 
+        # _get_amax is an internal TensorQuantizer method with no public equivalent;
+        # block_sizes is a public property.
         act_amax = self.input_quantizer._get_amax(input)
         weight = _nvfp4_quantize_weight_along_k(self.weight, self.weight_quantizer)
-        fp4_block_size = self.input_quantizer._block_sizes.get(-1, 16)
+        fp4_block_size = self.input_quantizer.block_sizes.get(-1, 16)
 
         output = conv3d_implicit_gemm_cuda(
             input,
@@ -137,7 +128,18 @@ def forward(self, input, *args, **kwargs):
         if not self._should_use_implicit_gemm():
             return super().forward(input, *args, **kwargs)
 
+        if self.training:
+            warnings.warn(
+                "Implicit GEMM Conv3D kernel is inference-only and does not support training. "
+                "Falling back to the default cuDNN quantization path, which could produce "
+                "different numerics.",
+                stacklevel=2,
+            )
+            return super().forward(input, *args, **kwargs)
+
         # During calibration, only collect amax — use the faster cuDNN path.
+        # _if_calib/_if_quant are internal TensorQuantizer state with no public property;
+        # toggled via enable_calib()/disable_calib()/enable_quant()/disable_quant().
         if self.input_quantizer._if_calib and not self.input_quantizer._if_quant:
             return super().forward(input, *args, **kwargs)