From 5b1a358020a8ea6340f67534a60520639f444686 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Sat, 18 Apr 2026 07:42:57 +0000 Subject: [PATCH 1/4] feat: NVFP4 Conv3D implicit GEMM kernel with end-to-end integration - Move CUDA implicit GEMM kernel from experimental/ to modelopt/torch/quantization/src/conv/ - Extend QuantConv to dispatch into the implicit GEMM kernel for NVFP4 - Add diffusers plugin hook for Wan Conv3D - Add unit, GPU kernel, and example integration tests - Update examples/diffusers/quantization/ for the E2E flow Signed-off-by: Jingyu Xin --- CHANGELOG.rst | 1 + .../diffusers/quantization/calibration.py | 13 +- .../diffusers/quantization/models_utils.py | 45 +- .../quantization/pipeline_manager.py | 85 ++- examples/diffusers/quantization/quantize.py | 136 ++-- .../diffusers/quantization/quantize_config.py | 3 +- examples/diffusers/quantization/utils.py | 18 +- experimental/conv/README.md | 105 ---- experimental/conv/bench_implicit_gemm.py | 208 ------- .../quantization/nn/modules/quant_conv.py | 79 ++- .../plugins/diffusion/diffusers.py | 35 ++ .../torch/quantization/src/conv/README.md | 140 +++++ .../src}/conv/implicit_gemm_binding.cpp | 0 .../src}/conv/implicit_gemm_cuda.py | 0 .../src}/conv/implicit_gemm_kernel.cu | 0 tests/examples/diffusers/conftest.py | 31 + tests/examples/diffusers/test_diffusers.py | 160 +++++ .../test_export_diffusers_hf_ckpt.py | 65 ++ .../kernels}/test_implicit_gemm.py | 585 +++++++++++++++--- .../plugins/test_diffusers_wan_conv3d.py | 129 ++++ .../torch/quantization/test_quant_conv.py | 141 +++++ 21 files changed, 1469 insertions(+), 510 deletions(-) delete mode 100644 experimental/conv/README.md delete mode 100644 experimental/conv/bench_implicit_gemm.py create mode 100644 modelopt/torch/quantization/src/conv/README.md rename {experimental => modelopt/torch/quantization/src}/conv/implicit_gemm_binding.cpp (100%) rename {experimental => modelopt/torch/quantization/src}/conv/implicit_gemm_cuda.py (100%) rename {experimental => modelopt/torch/quantization/src}/conv/implicit_gemm_kernel.cu (100%) create mode 100644 tests/examples/diffusers/conftest.py rename {experimental/conv => tests/gpu/torch/quantization/kernels}/test_implicit_gemm.py (64%) create mode 100644 tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 80dea0e43e4..20a677d0a0b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,7 @@ Changelog - Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md `_ 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 performant layerwise calibration for large models that don't fit on GPU (e.g. DeepSeek-R1, Kimi-K2). See `modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml `_ for usage. Layerwise calibration also supports PTQ with intermediate progress saving — useful when long PTQ runs get hit with Slurm timeouts. See `modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml `_ for usage. +- Add implicit GEMM CUDA kernel for Conv3D with fused NVFP4 fake quantization (``modelopt.torch.quantization.src.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** diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index 7464216cc83..27b1ec22436 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -108,11 +108,12 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: def _run_wan_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] ) -> None: + extra_params = self.pipeline_manager.config.extra_params kwargs = {} kwargs["negative_prompt"] = extra_args["negative_prompt"] - kwargs["height"] = extra_args["height"] - kwargs["width"] = extra_args["width"] - kwargs["num_frames"] = extra_args["num_frames"] + kwargs["height"] = extra_params.get("height", extra_args["height"]) + kwargs["width"] = extra_params.get("width", extra_args["width"]) + kwargs["num_frames"] = extra_params.get("num_frames", extra_args["num_frames"]) kwargs["guidance_scale"] = extra_args["guidance_scale"] if "guidance_scale_2" in extra_args: kwargs["guidance_scale_2"] = extra_args["guidance_scale_2"] @@ -154,7 +155,11 @@ def _run_ltx2_calibration(self, prompt_batch: list[str], extra_args: dict[str, A "images": extra_params.get("images", []), "tiling_config": extra_params.get("tiling_config", TilingConfig.default()), } - self.pipe(prompt=prompt, **kwargs) + decoded_video, decoded_audio = self.pipe(prompt=prompt, **kwargs) + # vae_decode_video returns a lazy generator — consume it so the + # video decoder's forward() actually runs during calibration. + for _ in decoded_video: + pass def _run_ltx_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 1e90c547148..b59744282f6 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -33,7 +33,9 @@ from utils import ( filter_func_default, filter_func_flux_dev, + filter_func_ltx2_vae, filter_func_ltx_video, + filter_func_wan_vae, filter_func_wan_video, ) @@ -54,31 +56,30 @@ class ModelType(str, Enum): WAN22_T2V_5b = "wan2.2-t2v-5b" -def get_model_filter_func(model_type: ModelType) -> Callable[[str], bool]: - """ - Get the appropriate filter function for a given model type. +_FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { + ModelType.FLUX_DEV: filter_func_flux_dev, + ModelType.FLUX2_DEV: filter_func_flux_dev, + ModelType.LTX_VIDEO_DEV: filter_func_ltx_video, + ModelType.LTX2: filter_func_ltx_video, + ModelType.WAN22_T2V_14b: filter_func_wan_video, + ModelType.WAN22_T2V_5b: filter_func_wan_video, +} - Args: - model_type: The model type enum +_VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { + (ModelType.LTX2, "video_decoder"): filter_func_ltx2_vae, + (ModelType.WAN22_T2V_14b, "vae"): filter_func_wan_vae, + (ModelType.WAN22_T2V_5b, "vae"): filter_func_wan_vae, +} - Returns: - A filter function appropriate for the model type - """ - filter_func_map = { - ModelType.FLUX_DEV: filter_func_flux_dev, - ModelType.FLUX_SCHNELL: filter_func_default, - ModelType.FLUX2_DEV: filter_func_flux_dev, - ModelType.SDXL_BASE: filter_func_default, - ModelType.SDXL_TURBO: filter_func_default, - ModelType.SD3_MEDIUM: filter_func_default, - ModelType.SD35_MEDIUM: filter_func_default, - ModelType.LTX_VIDEO_DEV: filter_func_ltx_video, - ModelType.LTX2: filter_func_ltx_video, - ModelType.WAN22_T2V_14b: filter_func_wan_video, - ModelType.WAN22_T2V_5b: filter_func_wan_video, - } - return filter_func_map.get(model_type, filter_func_default) +def get_model_filter_func( + model_type: ModelType, backbone_name: str = "transformer" +) -> Callable[[str], bool]: + """Get the appropriate filter function for a given model type and backbone.""" + vae_func = _VAE_FILTER_FUNC_MAP.get((model_type, backbone_name)) + if vae_func is not None: + return vae_func + return _FILTER_FUNC_MAP.get(model_type, filter_func_default) # Model registry with HuggingFace model IDs diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index 45f2fb63987..f62aeffca98 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -42,6 +42,7 @@ def __init__(self, config: ModelConfig, logger: logging.Logger): self.pipe: Any | None = None self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling self._transformer: torch.nn.Module | None = None + self._video_decoder: torch.nn.Module | None = None @staticmethod def create_pipeline_from( @@ -58,23 +59,20 @@ def create_pipeline_from( Raises: ValueError: If model type is unsupported """ - try: - pipeline_cls = MODEL_PIPELINE[model_type] - if pipeline_cls is None: - raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.") - model_id = ( - MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path - ) - pipe = pipeline_cls.from_pretrained( - model_id, - torch_dtype=torch_dtype, - use_safetensors=True, - **MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}), - ) - pipe.set_progress_bar_config(disable=True) - return pipe - except Exception as e: - raise e + pipeline_cls = MODEL_PIPELINE[model_type] + if pipeline_cls is None: + raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.") + model_id = ( + MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path + ) + pipe = pipeline_cls.from_pretrained( + model_id, + torch_dtype=torch_dtype, + use_safetensors=True, + **MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}), + ) + pipe.set_progress_bar_config(disable=True) + return pipe def create_pipeline(self) -> Any: """ @@ -157,42 +155,32 @@ def setup_device(self) -> None: self.logger.info("Enabling VAE tiling for LTX-Video") self.pipe.vae.enable_tiling() - def get_backbone(self) -> torch.nn.Module: - """ - Get the backbone model (transformer or UNet). - - Returns: - Backbone model module - """ - if not self.pipe: - raise RuntimeError("Pipeline not created. Call create_pipeline() first.") - - backbone_pairs = list(self.iter_backbones()) - if len(backbone_pairs) == 1: - return backbone_pairs[0][1] - return torch.nn.ModuleList([module for _, module in backbone_pairs]) - def iter_backbones(self) -> Iterator[tuple[str, torch.nn.Module]]: """ - Yield backbone modules by name, based on a backbone spec. - - Yields: - (backbone_name, module) pairs + Yield (backbone_name, module) pairs. """ if not self.pipe: raise RuntimeError("Pipeline not created. Call create_pipeline() first.") names = list(self.config.backbone) + if not names: + raise RuntimeError("No backbone names provided.") if self.config.model_type == ModelType.LTX2: - self._ensure_ltx2_transformer_cached() - name = names[0] if names else "transformer" - yield name, self._transformer + for name in names: + if name == "video_decoder": + self._ensure_ltx2_video_decoder_cached() + yield name, self._video_decoder + elif name == "transformer": + self._ensure_ltx2_transformer_cached() + yield name, self._transformer + else: + raise ValueError( + f"Unsupported LTX-2 backbone name '{name}'. " + "Expected 'transformer' or 'video_decoder'." + ) return - if not names: - raise RuntimeError("No backbone names provided.") - for name in names: module = getattr(self.pipe, name, None) if module is None: @@ -207,6 +195,16 @@ def _ensure_ltx2_transformer_cached(self) -> None: self.pipe.stage_1_model_ledger.transformer = lambda: transformer self._transformer = transformer + def _ensure_ltx2_video_decoder_cached(self) -> None: + if not self.pipe: + raise RuntimeError("Pipeline not created. Call create_pipeline() first.") + if self._video_decoder is None: + video_decoder = self.pipe.stage_1_model_ledger.video_decoder() + # Cache it so subsequent calls return the same (quantized) instance + self.pipe.stage_1_model_ledger.video_decoder = lambda: video_decoder + self.pipe.stage_2_model_ledger.video_decoder = lambda: video_decoder + self._video_decoder = video_decoder + def _create_ltx2_pipeline(self) -> Any: params = dict(self.config.extra_params) checkpoint_path = params.pop("checkpoint_path", None) @@ -261,7 +259,6 @@ def _create_ltx2_pipeline(self) -> Any: return TI2VidTwoStagesPipeline(**pipeline_kwargs) def print_quant_summary(self): - backbone_pairs = list(self.iter_backbones()) - for name, backbone in backbone_pairs: + for name, backbone in self.iter_backbones(): self.logger.info(f"{name} quantization info:") mtq.print_quant_summary(backbone) diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 953888f62b1..2a3c947a2d6 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -116,33 +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 + 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, @@ -158,6 +173,7 @@ def quantize_model( backbone: torch.nn.Module, quant_config: Any, forward_loop: callable, # type: ignore[valid-type] + backbone_name: str = "transformer", ) -> torch.nn.Module: """ Apply quantization to the model. @@ -166,15 +182,18 @@ def quantize_model( backbone: Model backbone to quantize quant_config: Quantization configuration forward_loop: Forward pass function for calibration + backbone_name: Name of the backbone being quantized """ self.logger.info("Checking for LoRA layers...") check_lora(backbone) - self.logger.info("Starting model quantization...") + self.logger.info(f"Starting model quantization for {backbone_name}...") mtq.quantize(backbone, quant_config, forward_loop) # Get model-specific filter function - model_filter_func = get_model_filter_func(self.model_config.model_type) - self.logger.info(f"Using filter function for {self.model_config.model_type.value}") + model_filter_func = get_model_filter_func(self.model_config.model_type, backbone_name) + self.logger.info( + f"Using filter function for {self.model_config.model_type.value}/{backbone_name}" + ) self.logger.info("Disabling specific quantizers...") mtq.disable_quantizer(backbone, model_filter_func) @@ -221,20 +240,27 @@ def _has_conv_layers(self, model: torch.nn.Module) -> bool: return True return False - def save_checkpoint(self, backbone: torch.nn.Module) -> None: + def save_checkpoint( + self, + backbone: torch.nn.Module, + backbone_name: str | None = None, + ) -> None: """ Save quantized model checkpoint. Args: backbone: The quantized backbone module to save (must be the same instance that was passed to mtq.quantize, as it carries the _modelopt_state). + backbone_name: Optional name for the backbone file (defaults to "backbone"). """ if not self.config.quantized_torch_ckpt_path: return ckpt_path = self.config.quantized_torch_ckpt_path ckpt_path.mkdir(parents=True, exist_ok=True) - target_path = ckpt_path / "backbone.pt" + filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" + target_path = ckpt_path / filename + self.logger.info(f"Saving backbone to {target_path}") mto.save(backbone, str(target_path)) @@ -292,14 +318,19 @@ def restore_checkpoint(self) -> None: if self.pipeline_manager is None: raise RuntimeError("Pipeline manager is required for per-backbone checkpoints.") - backbone = self.pipeline_manager.get_backbone() - if restore_path.exists() and restore_path.is_dir(): - source_path = restore_path / "backbone.pt" + if not restore_path.exists() or not restore_path.is_dir(): + raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}") + + for backbone_name, backbone in self.pipeline_manager.iter_backbones(): + source_path = restore_path / f"{backbone_name}.pt" if not source_path.exists(): - raise FileNotFoundError(f"Backbone checkpoint not found: {source_path}") - self.logger.info(f"Restoring backbone from {source_path}") + raise FileNotFoundError( + f"Checkpoint not found for '{backbone_name}' in {restore_path}" + ) + self.logger.info(f"Restoring {backbone_name} from {source_path}") mto.restore(backbone, str(source_path)) - self.logger.info("Backbone checkpoints restored successfully") + + self.logger.info("Checkpoints restored successfully") # TODO: should not do the any data type def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> None: @@ -368,9 +399,9 @@ def create_argument_parser() -> argparse.ArgumentParser: nargs="+", default=None, help=( - "Model backbone(s) in the DiffusionPipeline to work on. " - "Provide one name or multiple names separated by space or comma. " - "If not provided use default based on model type." + "Model backbone(s) in the DiffusionPipeline to quantize. " + "Provide one or more names (e.g., 'transformer', 'video_decoder'). " + "If not provided, uses default based on model type." ), ) model_group.add_argument( @@ -448,6 +479,12 @@ def create_argument_parser() -> argparse.ArgumentParser: action="store_true", help="Compress quantized weights to reduce memory footprint (FP8/FP4 only)", ) + quant_group.add_argument( + "--block-size", + type=int, + default=16, + help="Block size for NVFP4 quantization (default: 16)", + ) calib_group = parser.add_argument_group("Calibration Configuration") calib_group.add_argument("--batch-size", type=int, default=2, help="Batch size for calibration") @@ -535,6 +572,7 @@ def main() -> None: lowrank=args.lowrank, quantize_mha=args.quantize_mha, compress=args.compress, + block_size=args.block_size, ) if args.prompts_file is not None: @@ -571,7 +609,6 @@ def main() -> None: pipe = pipeline_manager.create_pipeline() pipeline_manager.setup_device() - backbone = pipeline_manager.get_backbone() export_manager = ExportManager(export_config, logger, pipeline_manager) if export_config.restore_from and export_config.restore_from.exists(): @@ -581,37 +618,48 @@ def main() -> None: logger.info("Initializing calibration...") calibrator = Calibrator(pipeline_manager, calib_config, model_config.model_type, logger) batched_prompts = calibrator.load_and_batch_prompts() - quantizer = Quantizer(quant_config, model_config, logger) - backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone) - # Pipe loads the ckpt just before the inference. - def forward_loop(mod): - calibrator.run_calibration(batched_prompts) + for backbone_name, backbone in pipeline_manager.iter_backbones(): + logger.info(f"Quantizing backbone: {backbone_name}") + backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone) - quantizer.quantize_model(backbone, backbone_quant_config, forward_loop) + # 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) - # Compress model weights if requested (only for FP8/FP4) - if quant_config.compress: - logger.info("Compressing model weights to reduce memory footprint...") - mtq.compress(backbone) - logger.info("Model compression completed") + quantizer.quantize_model( + backbone, + backbone_quant_config, + forward_loop, + backbone_name=backbone_name, + ) - export_manager.save_checkpoint(backbone) + # Compress model weights if requested (only for FP8/FP4) + if quant_config.compress: + logger.info(f"Compressing {backbone_name} weights...") + mtq.compress(backbone) + logger.info(f"{backbone_name} compression completed") - # TODO (Jingyu): To update this function, as we are focusing more on the torch deployment side. - check_conv_and_mha( - backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha - ) + # For VAE backbones, skip check_conv_and_mha — the whole point + # of VAE quantization is to quantize Conv layers. + if backbone_name not in ("video_decoder", "vae"): + check_conv_and_mha( + backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha + ) + + export_manager.save_checkpoint(backbone, backbone_name) pipeline_manager.print_quant_summary() - export_manager.export_onnx( - pipe, - backbone, - model_config.model_type, - quant_config.format, - ) + for backbone_name, backbone in pipeline_manager.iter_backbones(): + export_manager.export_onnx( + pipe, + backbone, + model_config.model_type, + quant_config.format, + ) export_manager.export_hf_ckpt(pipe, model_config) diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index 606debf9e32..a92dd4e8147 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -79,6 +79,7 @@ class QuantizationConfig: lowrank: int = 32 # SVDQuant lowrank quantize_mha: bool = False compress: bool = False + block_size: int = 16 # NVFP4 block size def validate(self) -> None: """Validate configuration consistency.""" @@ -120,7 +121,7 @@ class ModelConfig: model_type: ModelType = ModelType.FLUX_DEV model_dtype: dict[str, torch.dtype] = field(default_factory=lambda: {"default": torch.float16}) - backbone: str = "" + backbone: list[str] = field(default_factory=list) trt_high_precision_dtype: DataType = DataType.HALF override_model_path: Path | None = None cpu_offloading: bool = False diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index 0c38fc28606..d102e83e068 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -46,7 +46,7 @@ def filter_func_default(name: str) -> bool: def check_conv_and_mha(backbone, if_fp4, quantize_mha): for name, module in backbone.named_modules(): - if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and if_fp4: + if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d)) and if_fp4: module.weight_quantizer.disable() module.input_quantizer.disable() @@ -87,6 +87,22 @@ def filter_func_flux_dev(name: str) -> bool: return pattern.match(name) is not None +def filter_func_ltx2_vae(name: str) -> bool: + """Filter for LTX-2 VAE: keeps only conv1/conv2 in up_blocks resnets.""" + keep = re.compile(r".*up_blocks\.\d+\.resnets\.\d+\.conv[12](?:\.|$)") + return not keep.match(name) + + +def filter_func_wan_vae(name: str) -> bool: + """Filter for Wan 2.2 VAE: keeps only conv1/conv2 in resnet blocks.""" + keep = re.compile( + r".*(down_blocks\.\d+\.(?:resnets\.\d+\.)?conv[12]" + r"|mid_block\.resnets\.\d+\.conv[12]" + r"|up_blocks\.\d+\.resnets\.\d+\.conv[12])(?:\.|$)" + ) + return not keep.match(name) + + def filter_func_wan_video(name: str) -> bool: """Filter function specifically for WAN-Video models.""" pattern = re.compile( diff --git a/experimental/conv/README.md b/experimental/conv/README.md deleted file mode 100644 index 65b7cc5563e..00000000000 --- a/experimental/conv/README.md +++ /dev/null @@ -1,105 +0,0 @@ -# Conv3D Implicit GEMM (Experimental) - -Experimental Conv3D kernel prototype using implicit GEMM, with optional fused FP4 fake quantization for activations. - -This code is kept under `experimental/` by design and is **not** part of the stable `modelopt.torch.quantization` API. - -## Model Support - -| Model/Framework | Supported | Notes | -|-----------------|-----------|-------| -| Video diffusion VAE Conv3D layers | Tested | Validated on VAE encoder/decoder Conv3D layers in video diffusion models | -| Generic LLM backbones | No | Conv3D path is not relevant | -| End-to-end ModelOpt PTQ/QAT pipeline | No | Not wired into formal quantization/export/compress flows | - -## Deployment - -| Framework | Supported | Notes | -|-----------|-----------|-------| -| TensorRT-LLM | No | No formal export integration for this kernel path | -| vLLM | No | No integration | -| SGLang | No | No integration | -| PyTorch runtime (CUDA) | Yes (experimental) | JIT-compiles CUDA extension on first use | - -## Usage - -```python -import torch - -from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda -from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op - -x = torch.randn(1, 128, 21, 60, 106, device="cuda") -w = torch.randn(512, 128, 3, 3, 3, device="cuda") -block_size = 128 - -# Without FP4 activation quantization (drop-in-style Conv3D call) -out = conv3d_implicit_gemm_cuda(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) - -# Optional FP4 block quantization of weights along the GEMM K dimension. -# The kernel's A-tile (activations) is quantized along K = Cin*kD*kH*kW, -# so weights must be flattened to [Cout, K] before quantizing to match. -Cout, Cin = w.shape[:2] -K = Cin * w.shape[2] * w.shape[3] * w.shape[4] -w_flat = w.reshape(Cout, K) -w_q_flat = dynamic_block_quantize_op( - w_flat, - block_size, - w_flat.abs().max().unsqueeze(0), - 4, # num_bits - 2, # exponent_bits - 8, # scale_num_bits - 4, # scale_exponent_bits -) -w_q = w_q_flat.reshape_as(w) - -# With FP4 activation fake quantization -out_q = conv3d_implicit_gemm_cuda( - x, - w_q, - stride=(1, 1, 1), - padding=(1, 1, 1), - act_amax=x.abs().max().unsqueeze(0), - quant_act=True, - fp4_block_size=block_size, # 16, 32, 64, 128, or 256 -) -``` - -## API - -Function: `conv3d_implicit_gemm_cuda(...)` from `experimental/conv/implicit_gemm_cuda.py` - -| Parameter | Description | -|-----------|-------------| -| `x` | Input tensor `[N, Cin, D, H, W]` | -| `w` | Weight tensor `[Cout, Cin, kD, kH, kW]` | -| `bias` | Optional bias `[Cout]` | -| `stride` | Convolution stride `(D, H, W)` | -| `padding` | Convolution padding `(D, H, W)` | -| `dilation` | Convolution dilation `(D, H, W)` | -| `act_amax` | Activation abs-max scalar tensor (required when `quant_act=True`) | -| `quant_act` | Enable FP4 fake quantization on activations | -| `fp4_block_size` | FP4 quantization block size (`16`, `32`, `64`, `128`, or `256`) | - -## Status - -Current state: **Prototype** - -Known limitations: - -- API is unstable and may change without notice. -- Not registered in core quantization module registries. -- Not covered by formal export/compress integration. -- CUDA extension compile latency on first invocation. -- Validation and performance coverage are limited to local experiments. - -## Notes - -- The CUDA kernel is JIT-compiled on first call (can take several seconds). -- Output shape matches `torch.nn.functional.conv3d`. -- FP4 path applies quantize-dequantize in-kernel for activation tiles. - -## References - -- Implicit GEMM-based convolution design patterns in GPU kernels. -- ModelOpt FP4-related quantization utilities in `modelopt.torch.quantization.tensor_quant`. diff --git a/experimental/conv/bench_implicit_gemm.py b/experimental/conv/bench_implicit_gemm.py deleted file mode 100644 index 164c0744674..00000000000 --- a/experimental/conv/bench_implicit_gemm.py +++ /dev/null @@ -1,208 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Latency benchmark: implicit GEMM (quant / non-quant) vs cuDNN conv3d. - -Usage: - python -m experimental.conv.bench_implicit_gemm - python -m experimental.conv.bench_implicit_gemm --shapes wan22 - python -m experimental.conv.bench_implicit_gemm --shapes all --warmup 20 --iters 100 -""" - -import argparse - -import torch -import torch.nn.functional as F - -# --------------------------------------------------------------------------- -# Benchmark shapes -# --------------------------------------------------------------------------- - -# (name, N, Cin, D, H, W, Cout, kD, kH, kW, stride, padding, dilation) -SHAPES = { - "small": [ - ("small_16x32_3x3x3", 1, 16, 8, 8, 8, 32, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ], - "medium": [ - ("med_64x128_3x3x3", 1, 64, 16, 32, 32, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ("med_128x256_3x3x3", 1, 128, 8, 16, 16, 256, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ("med_128x128_1x3x3", 1, 128, 16, 32, 32, 128, 1, 3, 3, (1, 1, 1), (0, 1, 1), (1, 1, 1)), - ], - "wan22": [ - ("wan22_128x512", 1, 128, 21, 60, 106, 512, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ("wan22_512x512", 1, 512, 21, 60, 106, 512, 1, 1, 1, (1, 1, 1), (0, 0, 0), (1, 1, 1)), - ("wan22_512x128", 1, 512, 21, 60, 106, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ], - "stride": [ - ("stride2_64x128", 1, 64, 16, 32, 32, 128, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)), - ("stride2_128x256", 1, 128, 16, 32, 32, 256, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)), - ], -} - - -def get_shapes(name: str): - """Return list of benchmark shapes by name or all shapes.""" - if name == "all": - result = [] - for v in SHAPES.values(): - result.extend(v) - return result - return SHAPES[name] - - -# --------------------------------------------------------------------------- -# Timing utility -# --------------------------------------------------------------------------- - - -def bench_fn(fn, warmup: int, iters: int) -> float: - """Benchmark a callable, return median time in ms.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - times = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - torch.cuda.synchronize() - times.append(start.elapsed_time(end)) - - times.sort() - return times[len(times) // 2] # median - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -def run_benchmark(shapes_name: str, warmup: int, iters: int, fp4_block_size: int): - """Run latency benchmark for the given shapes.""" - from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda - - shapes = get_shapes(shapes_name) - - # Header - print(f"\n{'=' * 100}") - print( - f"Conv3D Latency Benchmark | warmup={warmup} iters={iters} fp4_block_size={fp4_block_size}" - ) - print(f"GPU: {torch.cuda.get_device_name()}") - print(f"{'=' * 100}") - print( - f"{'Shape':<25} {'M':>10} {'K':>8} {'N':>6} " - f"{'cuDNN':>9} {'GEMM':>9} {'GEMM+FP4':>9} " - f"{'GEMM/cuDNN':>11} {'FP4/cuDNN':>10}" - ) - print("-" * 100) - - for name, n, cin, d, h, w, cout, kd, kh, kw, stride, padding, dilation in shapes: - torch.manual_seed(42) - x = torch.randn(n, cin, d, h, w, device="cuda", dtype=torch.float32) - weight = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32) - act_amax = x.abs().max().unsqueeze(0) - - # Compute GEMM dimensions for display - sd, sh, sw = stride - dd, dh, dw = dilation - pd, ph, pw = padding - od = (d + 2 * pd - dd * (kd - 1) - 1) // sd + 1 - oh = (h + 2 * ph - dh * (kh - 1) - 1) // sh + 1 - ow = (w + 2 * pw - dw * (kw - 1) - 1) // sw + 1 - gemm_m = n * od * oh * ow - gemm_k = cin * kd * kh * kw - gemm_n = cout - - # cuDNN (torch.nn.functional.conv3d) - t_cudnn = bench_fn( - lambda: F.conv3d(x, weight, stride=stride, padding=padding, dilation=dilation), - warmup, - iters, - ) - - # Implicit GEMM (non-quantized) - t_gemm = bench_fn( - lambda: conv3d_implicit_gemm_cuda( - x, - weight, - stride=stride, - padding=padding, - dilation=dilation, - quant_act=False, - fp4_block_size=fp4_block_size, - ), - warmup, - iters, - ) - - # Implicit GEMM (FP4 quantized) - t_fp4 = bench_fn( - lambda: conv3d_implicit_gemm_cuda( - x, - weight, - stride=stride, - padding=padding, - dilation=dilation, - act_amax=act_amax, - quant_act=True, - fp4_block_size=fp4_block_size, - ), - warmup, - iters, - ) - - ratio_gemm = t_gemm / t_cudnn - ratio_fp4 = t_fp4 / t_cudnn - - print( - f"{name:<25} {gemm_m:>10,} {gemm_k:>8,} {gemm_n:>6,} " - f"{t_cudnn:>8.3f}ms {t_gemm:>8.3f}ms {t_fp4:>8.3f}ms " - f"{ratio_gemm:>10.2f}x {ratio_fp4:>9.2f}x" - ) - - print(f"{'=' * 100}") - print("Ratios > 1.0x mean slower than cuDNN; < 1.0x mean faster.") - print() - - -def main(): - """Entry point for the benchmark CLI.""" - parser = argparse.ArgumentParser(description="Conv3D latency benchmark") - parser.add_argument( - "--shapes", - default="all", - choices=[*list(SHAPES.keys()), "all"], - help="Which shape set to benchmark (default: all)", - ) - parser.add_argument("--warmup", type=int, default=20, help="Warmup iterations") - parser.add_argument("--iters", type=int, default=100, help="Benchmark iterations") - parser.add_argument( - "--fp4-block-size", - type=int, - default=128, - choices=[128, 256], - help="FP4 block size (default: 128)", - ) - args = parser.parse_args() - - run_benchmark(args.shapes, args.warmup, args.iters, args.fp4_block_size) - - -if __name__ == "__main__": - main() diff --git a/modelopt/torch/quantization/nn/modules/quant_conv.py b/modelopt/torch/quantization/nn/modules/quant_conv.py index 44f0ae663c5..ed165556249 100644 --- a/modelopt/torch/quantization/nn/modules/quant_conv.py +++ b/modelopt/torch/quantization/nn/modules/quant_conv.py @@ -15,8 +15,12 @@ """Quantized convolution.""" +import warnings + import torch.nn as nn +from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + from ... import tensor_quant from .quant_module import QuantLinearConvBase, QuantModuleRegistry, _LegacyQuantLinearConvBaseMixin @@ -62,12 +66,85 @@ class QuantConv2d(_LegacyQuantLinearConvBaseMixin, nn.Conv2d): default_quant_desc_weight = _QuantConv2d.default_quant_desc_weight +def _is_nvfp4_quantizer(quantizer) -> bool: + """Check if a TensorQuantizer is configured for NVFP4 dynamic block quantization.""" + return ( + quantizer.num_bits == (2, 1) + and quantizer.block_sizes is not None + and quantizer.block_sizes.get("scale_bits") == (4, 3) + and quantizer.block_sizes.get("type") == "dynamic" + ) + + +def _nvfp4_quantize_weight_along_k(weight, weight_quantizer): + """Apply NVFP4 fake quantization to Conv3D weight along the GEMM K dimension.""" + w_flat = weight.reshape(weight.shape[0], -1) + return weight_quantizer(w_flat).reshape_as(weight) + + @QuantModuleRegistry.register({nn.Conv3d: "nn.Conv3d"}) class _QuantConv3d(QuantLinearConvBase): - """Quantized 3D convolution.""" + """Quantized 3D convolution. + + For NVFP4, uses a fused implicit GEMM kernel with activation FP4 quantization + inside the kernel. For all other configs, the default cuDNN path is used. + """ default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL + def _should_use_implicit_gemm(self): + """Check if both quantizers are NVFP4 and the implicit GEMM kernel is available.""" + 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 + ) + + def _implicit_gemm_forward(self, input): + """Run NVFP4 implicit GEMM kernel. Input may already be padded.""" + # _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) + + output = conv3d_implicit_gemm_cuda( + input, + weight, + bias=self.bias, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + act_amax=act_amax, + quant_act=self.input_quantizer.is_enabled, + fp4_block_size=fp4_block_size, + ) + return self.output_quantizer(output) + + def forward(self, input, *args, **kwargs): + """Forward with implicit GEMM for NVFP4, default path otherwise.""" + 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) + + return self._implicit_gemm_forward(input) + class QuantConv3d(_LegacyQuantLinearConvBaseMixin, nn.Conv3d): """Quantized 3D convolution.""" diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f9ae55b3e2d..f2f6a702479 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -63,6 +63,7 @@ QuantModuleRegistry, TensorQuantizer, ) +from ...nn.modules.quant_conv import _QuantConv3d from ..custom import _QuantFunctionalMixin onnx_dtype_map = { @@ -278,3 +279,37 @@ def symbolic( high_precision_flag, disable_fp8_mha, ) + + +# WanCausalConv3d quantization support (diffusers VAE) +try: + from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + + @QuantModuleRegistry.register({WanCausalConv3d: "WanCausalConv3d"}) + class _QuantDiffusersWanCausalConv3d(_QuantConv3d): + """Quantized WanCausalConv3d — applies causal padding before quantized conv.""" + + def forward(self, x, cache_x=None): + # Apply WanCausalConv3d-specific causal padding + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + # NVFP4 implicit GEMM path (self.padding is (0,0,0) since padding already applied) + if self._should_use_implicit_gemm(): + if not (self.input_quantizer._if_calib and not self.input_quantizer._if_quant): + return self._implicit_gemm_forward(x) + + # Default quantized conv path (skip WanCausalConv3d.forward to avoid double-padding) + with self.quantize_weight(): + input = self.input_quantizer(x) + output = torch.nn.Conv3d.forward(self, input) + if isinstance(output, tuple): + return (self.output_quantizer(output[0]), *output[1:]) + return self.output_quantizer(output) + +except ImportError: + pass diff --git a/modelopt/torch/quantization/src/conv/README.md b/modelopt/torch/quantization/src/conv/README.md new file mode 100644 index 00000000000..37580b69a80 --- /dev/null +++ b/modelopt/torch/quantization/src/conv/README.md @@ -0,0 +1,140 @@ +# Conv3D Implicit GEMM + +Conv3D kernel using implicit GEMM with BF16 WMMA tensor cores and optional fused FP4 (E2M1) fake quantization. + +This kernel is integrated into `modelopt.torch.quantization` via `_QuantConv3d` — when NVFP4 quantization is applied to an `nn.Conv3d` layer through ModelOpt PTQ, the implicit GEMM path is used automatically. We have only tested it on VAE Conv3D layers from video generation models (e.g. Wan2.2). + +## Requirements + +- **GPU:** SM80+ (Ampere or newer) for BF16 WMMA tensor cores +- **PyTorch:** CUDA toolkit with JIT C++ extension support (`torch.utils.cpp_extension`) +- **Grouped convolution is not supported** (groups must be 1) + +## Data Types + +| Stage | Precision | +|-------|-----------| +| Input / output tensors | FP32, FP16, or BF16 (dtype is preserved) | +| Internal compute | BF16 via WMMA m16n16k16 tensor cores | +| Accumulation | FP32 | +| FP4 activation quantization | E2M1 values, FP8 E4M3 scales | + +## Integration with ModelOpt Quantization + +When NVFP4 quantization is configured on a `Conv3d` layer via ModelOpt PTQ, the implicit GEMM kernel is used automatically during quantized inference. The integration is in `_QuantConv3d` (`modelopt/torch/quantization/nn/modules/quant_conv.py`): + +- During **calibration**, the standard cuDNN path is used (faster). +- During **quantized inference** with NVFP4 input and weight quantizers, the kernel fuses activation FP4 quantization inside the GEMM. +- For all other quantization configs, the default cuDNN path is used as fallback. + +## Usage + +```python +import torch + +from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda +from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op + +x = torch.randn(1, 128, 21, 60, 106, device="cuda") +w = torch.randn(512, 128, 3, 3, 3, device="cuda") +block_size = 128 + +# Without FP4 activation quantization (drop-in-style Conv3D call) +out = conv3d_implicit_gemm_cuda(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + +# Optional FP4 block quantization of weights along the GEMM K dimension. +# The kernel's A-tile (activations) is quantized along K = Cin*kD*kH*kW, +# so weights must be flattened to [Cout, K] before quantizing to match. +Cout, Cin = w.shape[:2] +K = Cin * w.shape[2] * w.shape[3] * w.shape[4] +w_flat = w.reshape(Cout, K) +w_q_flat = dynamic_block_quantize_op( + w_flat, + block_size, + w_flat.abs().max().unsqueeze(0), + 4, # num_bits + 2, # exponent_bits + 8, # scale_num_bits + 4, # scale_exponent_bits +) +w_q = w_q_flat.reshape_as(w) + +# With FP4 activation fake quantization +out_q = conv3d_implicit_gemm_cuda( + x, + w_q, + stride=(1, 1, 1), + padding=(1, 1, 1), + act_amax=x.abs().max().unsqueeze(0), + quant_act=True, + fp4_block_size=block_size, # 16, 32, 64, 128, or 256 +) +``` + +## API + +### `conv3d_implicit_gemm_cuda` + +`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda` + +| Parameter | Description | +|-----------|-------------| +| `x` | Input tensor `[N, Cin, D, H, W]` | +| `w` | Weight tensor `[Cout, Cin, kD, kH, kW]` | +| `bias` | Optional bias `[Cout]` | +| `stride` | Convolution stride `(D, H, W)` | +| `padding` | Convolution padding `(D, H, W)` | +| `dilation` | Convolution dilation `(D, H, W)` | +| `act_amax` | Activation abs-max scalar tensor (required when `quant_act=True`) | +| `quant_act` | Enable FP4 fake quantization on activations | +| `fp4_block_size` | FP4 quantization block size (`16`, `32`, `64`, `128`, or `256`) | + +### `fp4_fake_quant` + +`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant` + +Standalone FP4 (E2M1) blockwise fake quantization with FP8 E4M3 scale quantization. Uses the same CUDA device functions as the fused path inside the GEMM kernel. + +| Parameter | Description | +|-----------|-------------| +| `x` | Input tensor (any shape; `numel` must be divisible by `block_size`) | +| `global_amax` | Scalar tensor — global abs max for scale computation | +| `block_size` | Number of elements per FP4 quantization block (default `16`) | + +## Testing + +```bash +# Run tests (requires GPU) +python -m pytest tests/gpu/torch/quantization/kernels/test_implicit_gemm.py -v +``` + +## Status + +Current state: **Integrated** (registered in `QuantModuleRegistry`, auto-dispatched for NVFP4 Conv3D) + +Known limitations: + +- CUDA extension compile latency on first invocation (~seconds). +- Grouped convolution (`groups > 1`) is not supported. In the ModelOpt E2E flow, `_QuantConv3d` automatically falls back to the default cuDNN path for grouped convolutions. +- BF16 rounding error accumulates with the K dimension — expect max abs diff scaling roughly as `sqrt(K)` compared to cuDNN FP32. +- Inference only (`@torch.no_grad`) — not suitable for QAT backward pass. + +## Notes + +- The CUDA kernel is JIT-compiled on first call via `torch.utils.cpp_extension.load()`. +- Output shape matches `torch.nn.functional.conv3d`. +- FP4 path applies quantize-dequantize in-kernel for activation tiles (no extra global memory pass). +- Tile config: BLOCK_M=64, BLOCK_N=64, BLOCK_K=256, 8 warps (256 threads), ~70 KB shared memory per block. + +## Files + +| File | Role | +|------|------| +| `implicit_gemm_cuda.py` | Python API and JIT compilation | +| `implicit_gemm_kernel.cu` | CUDA kernel (BF16 WMMA + FP4 quantization) | +| `implicit_gemm_binding.cpp` | PyTorch C++ extension binding | + +## References + +- Implicit GEMM-based convolution design patterns in GPU kernels. +- ModelOpt FP4-related quantization utilities in `modelopt.torch.quantization.tensor_quant`. diff --git a/experimental/conv/implicit_gemm_binding.cpp b/modelopt/torch/quantization/src/conv/implicit_gemm_binding.cpp similarity index 100% rename from experimental/conv/implicit_gemm_binding.cpp rename to modelopt/torch/quantization/src/conv/implicit_gemm_binding.cpp diff --git a/experimental/conv/implicit_gemm_cuda.py b/modelopt/torch/quantization/src/conv/implicit_gemm_cuda.py similarity index 100% rename from experimental/conv/implicit_gemm_cuda.py rename to modelopt/torch/quantization/src/conv/implicit_gemm_cuda.py diff --git a/experimental/conv/implicit_gemm_kernel.cu b/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu similarity index 100% rename from experimental/conv/implicit_gemm_kernel.cu rename to modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu diff --git a/tests/examples/diffusers/conftest.py b/tests/examples/diffusers/conftest.py new file mode 100644 index 00000000000..8893d188d9e --- /dev/null +++ b/tests/examples/diffusers/conftest.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + + +@pytest.fixture(scope="session") +def tiny_wan22_path(tmp_path_factory): + """Create a tiny Wan 2.2 (14B-style) pipeline and return its path. + + Built once per session and shared across all tests that need it. + """ + try: + from _test_utils.torch.diffusers_models import create_tiny_wan22_pipeline_dir + except ImportError: + pytest.skip("Wan 2.2 diffusers models not available (requires diffusers with WanPipeline)") + + tmp_path = tmp_path_factory.mktemp("wan22") + return str(create_tiny_wan22_pipeline_dir(tmp_path)) diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 5bc8f981ecd..9af6481d918 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -21,6 +21,27 @@ from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm +# Tiny video model args — override MODEL_DEFAULTS for fast CI +_WAN22_TINY_EXTRA_PARAMS = [ + "--extra-param", + "height=16", + "--extra-param", + "width=16", + "--extra-param", + "num_frames=5", +] + +_WAN22_FAST_CALIB_ARGS = [ + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + "--model-dtype", + "BFloat16", +] + class DiffuserModel(NamedTuple): dtype: str @@ -151,6 +172,145 @@ def test_diffusers_quantization( model.inference(tmp_path) +def _run_wan22_quantize( + tiny_wan22_path: str, + tmp_path: Path, + format_type: str, + quant_algo: str, + collect_method: str, + *, + model: str = "wan2.2-t2v-14b", + backbone: str | None = None, + ckpt_suffix: str | None = None, +) -> None: + """Run quantize.py for Wan 2.2 with the tiny model.""" + suffix = ckpt_suffix or format_type + ckpt_path = str(tmp_path / f"wan22_{suffix}.pt") + cmd_args = [ + "python", + "quantize.py", + "--model", + model, + "--override-model-path", + tiny_wan22_path, + "--format", + format_type, + "--quant-algo", + quant_algo, + "--collect-method", + collect_method, + "--trt-high-precision-dtype", + "BFloat16", + "--quantized-torch-ckpt-save-path", + ckpt_path, + *_WAN22_FAST_CALIB_ARGS, + *_WAN22_TINY_EXTRA_PARAMS, + ] + if backbone is not None: + cmd_args.extend(["--backbone", backbone]) + run_example_command(cmd_args, "diffusers/quantization") + + +def _run_wan22_restore( + tiny_wan22_path: str, + tmp_path: Path, + format_type: str, + quant_algo: str, + collect_method: str, + *, + model: str = "wan2.2-t2v-14b", + backbone: str | None = None, + ckpt_suffix: str | None = None, +) -> None: + """Restore a Wan 2.2 quantized checkpoint.""" + suffix = ckpt_suffix or format_type + ckpt_path = str(tmp_path / f"wan22_{suffix}.pt") + cmd_args = [ + "python", + "quantize.py", + "--model", + model, + "--override-model-path", + tiny_wan22_path, + "--format", + format_type, + "--quant-algo", + quant_algo, + "--collect-method", + collect_method, + "--trt-high-precision-dtype", + "BFloat16", + "--restore-from", + ckpt_path, + *_WAN22_FAST_CALIB_ARGS, + *_WAN22_TINY_EXTRA_PARAMS, + ] + if backbone is not None: + cmd_args.extend(["--backbone", backbone]) + run_example_command(cmd_args, "diffusers/quantization") + + +def test_wan22_int8_smoothquant(tiny_wan22_path: str, tmp_path: Path) -> None: + """Wan 2.2 INT8 SmoothQuant: quantize + restore.""" + _run_wan22_quantize(tiny_wan22_path, tmp_path, "int8", "smoothquant", "min-mean") + _run_wan22_restore(tiny_wan22_path, tmp_path, "int8", "smoothquant", "min-mean") + + +@pytest.mark.parametrize( + ("format_type", "quant_algo"), + [ + pytest.param("fp8", "max", marks=minimum_sm(89)), + pytest.param("fp4", "max", marks=minimum_sm(89)), + ], + ids=["wan22_fp8_max", "wan22_fp4_max"], +) +def test_wan22_fp8_fp4( + tiny_wan22_path: str, tmp_path: Path, format_type: str, quant_algo: str +) -> None: + """Wan 2.2 FP8/FP4: quantize + restore (requires SM89+).""" + _run_wan22_quantize(tiny_wan22_path, tmp_path, format_type, quant_algo, "default") + _run_wan22_restore(tiny_wan22_path, tmp_path, format_type, quant_algo, "default") + + +@pytest.mark.parametrize( + ("format_type", "quant_algo"), + [ + pytest.param("fp8", "max", marks=minimum_sm(89)), + pytest.param("fp4", "max", marks=minimum_sm(89)), + ], + ids=["wan22_vae_fp8_max", "wan22_vae_fp4_max"], +) +def test_wan22_vae_fp8_fp4( + tiny_wan22_path: str, tmp_path: Path, format_type: str, quant_algo: str +) -> None: + """Wan 2.2 VAE FP8/FP4 quantization: quantize + restore. + + Exercises the ``WanCausalConv3d``/Conv3D NVFP4 implicit-GEMM dispatch path + on the Wan VAE end-to-end through ``quantize.py``. + """ + ckpt_suffix = f"vae_{format_type}" + _run_wan22_quantize( + tiny_wan22_path, + tmp_path, + format_type, + quant_algo, + "default", + model="wan2.2-t2v-5b", + backbone="vae", + ckpt_suffix=ckpt_suffix, + ) + _run_wan22_restore( + tiny_wan22_path, + tmp_path, + format_type, + quant_algo, + "default", + model="wan2.2-t2v-5b", + backbone="vae", + ckpt_suffix=ckpt_suffix, + ) + + @pytest.mark.parametrize( ("model_name", "model_path", "torch_compile"), [ diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 6db1eaeb68d..5dca512b5b9 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -21,6 +21,16 @@ from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm +# Tiny video model args — override MODEL_DEFAULTS for fast CI +_WAN22_TINY_EXTRA_PARAMS = [ + "--extra-param", + "height=16", + "--extra-param", + "width=16", + "--extra-param", + "num_frames=5", +] + class DiffuserHfExportModel(NamedTuple): name: str @@ -128,3 +138,58 @@ def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path) weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin")) assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" + + +@pytest.mark.parametrize( + ("format_type", "quant_algo", "collect_method"), + [ + ("int8", "smoothquant", "min-mean"), + pytest.param("fp8", "max", "default", marks=minimum_sm(89)), + ], + ids=["wan22_int8_smoothquant", "wan22_fp8_max"], +) +def test_wan22_hf_ckpt_export( + tiny_wan22_path: str, + tmp_path: Path, + format_type: str, + quant_algo: str, + collect_method: str, +) -> None: + """Quantize tiny Wan 2.2 and export to HF checkpoint.""" + hf_ckpt_dir = tmp_path / f"wan22_{format_type}_hf_ckpt" + cmd_args = [ + "python", + "quantize.py", + "--model", + "wan2.2-t2v-14b", + "--override-model-path", + tiny_wan22_path, + "--format", + format_type, + "--quant-algo", + quant_algo, + "--collect-method", + collect_method, + "--model-dtype", + "BFloat16", + "--trt-high-precision-dtype", + "BFloat16", + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + "--hf-ckpt-dir", + str(hf_ckpt_dir), + *_WAN22_TINY_EXTRA_PARAMS, + ] + run_example_command(cmd_args, "diffusers/quantization") + + assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}" + + config_files = list(hf_ckpt_dir.rglob("config.json")) + assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}" + + weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin")) + assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" diff --git a/experimental/conv/test_implicit_gemm.py b/tests/gpu/torch/quantization/kernels/test_implicit_gemm.py similarity index 64% rename from experimental/conv/test_implicit_gemm.py rename to tests/gpu/torch/quantization/kernels/test_implicit_gemm.py index af52660e42e..56ceaacc01f 100644 --- a/experimental/conv/test_implicit_gemm.py +++ b/tests/gpu/torch/quantization/kernels/test_implicit_gemm.py @@ -28,7 +28,7 @@ @pytest.fixture(scope="module") def cuda_conv3d(): """Import and return the CUDA implicit GEMM conv3d function.""" - from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda return conv3d_implicit_gemm_cuda @@ -303,9 +303,9 @@ def test_deterministic(self, cuda_conv3d): @pytest.fixture(scope="module") -def cuda_fp4_quant(): - """Import FP4 fake quant for reference comparisons.""" - from experimental.conv.implicit_gemm_cuda import fp4_fake_quant +def cuda_fp4(): + """Import and return the CUDA FP4 fake quant function.""" + from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant return fp4_fake_quant @@ -361,7 +361,7 @@ def test_quant_deterministic(self, cuda_conv3d, fp4_block_size): assert torch.equal(out1, out2), f"Non-deterministic for fp4_block_size={fp4_block_size}" @pytest.mark.parametrize("fp4_block_size", [16, 32, 64, 128, 256]) - def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4_quant, fp4_block_size): + def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4, fp4_block_size): """Compare fused kernel vs unfused: fp4(im2col) @ fp4(weight). Uses a shape where K is a multiple of 256 so all K-tiles are full @@ -382,9 +382,9 @@ def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4_quant, fp4_block im2col = x.permute(0, 2, 3, 4, 1).reshape(-1, cin) # [M, K] # 2. FP4 fake-quant both matrices along K with the same block_size - im2col_q = cuda_fp4_quant(im2col, act_amax, fp4_block_size) + im2col_q = cuda_fp4(im2col, act_amax, fp4_block_size) w_flat = w.reshape(cout, cin).transpose(0, 1).contiguous() # [K, Cout] - w_flat_q = cuda_fp4_quant(w_flat, w_amax, fp4_block_size) + w_flat_q = cuda_fp4(w_flat, w_amax, fp4_block_size) # 3. Matmul (in BF16 to match kernel's WMMA path) ref_out = (im2col_q.bfloat16() @ w_flat_q.bfloat16()).float() @@ -421,81 +421,33 @@ def test_smaller_block_less_error(self, cuda_conv3d): """Smaller FP4 block sizes should generally produce lower quantization error. Finer-grained blocks capture local ranges better, reducing quant error vs cuDNN. - Test monotonicity: error(16) <= error(32) <= ... <= error(256) (with some tolerance). - Reports detailed accuracy metrics for each block size vs cuDNN baseline. + Test monotonicity on a medium config: error(16) <= error(64) <= error(256) (with 1.2x slack). """ torch.manual_seed(42) - # Test multiple shapes to get a comprehensive picture - configs = [ - ("Small K=432", 1, 16, 8, 8, 8, 32, 3, 3, 3), - ("Medium K=1728", 1, 64, 8, 8, 8, 64, 3, 3, 3), - ("Large K=3456", 1, 128, 5, 8, 8, 256, 3, 3, 3), - ] + # Medium K=1728: Cin=64, 3x3x3 kernel + cin, cout = 64, 64 + x = torch.randn(1, cin, 8, 8, 8, device="cuda", dtype=torch.float32) + w = torch.randn(cout, cin, 3, 3, 3, device="cuda", dtype=torch.float32) + act_amax = x.abs().max().unsqueeze(0) + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1)) block_sizes = [16, 32, 64, 128, 256] - all_errors = {} - - for desc, n, cin, d, h, w_s, cout, kd, kh, kw in configs: - x = torch.randn(n, cin, d, h, w_s, device="cuda", dtype=torch.float32) - w = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32) - act_amax = x.abs().max().unsqueeze(0) - k_size = cin * kd * kh * kw - - ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1)) - ref_abs_mean = ref.abs().mean().item() - - # Also compute no-quant baseline (BF16 rounding only) - out_nq = cuda_conv3d( + errors = {} + for bs in block_sizes: + out = cuda_conv3d( x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1), - quant_act=False, - ) - nq_diff = (out_nq - ref).abs() - - print( - f"\n {desc} (K={k_size}), output range [{ref.min().item():.1f}, {ref.max().item():.1f}]" - ) - print( - f" {'Block Size':>10} | {'Max Diff':>10} | {'Mean Diff':>10} | {'RMSE':>10} | {'Rel Err%':>8}" - ) - print(f" {'-' * 10}-+-{'-' * 10}-+-{'-' * 10}-+-{'-' * 10}-+-{'-' * 8}") - print( - f" {'no-quant':>10} | {nq_diff.max().item():>10.4f} | " - f"{nq_diff.mean().item():>10.6f} | " - f"{((out_nq - ref) ** 2).mean().sqrt().item():>10.4f} | " - f"{nq_diff.mean().item() / ref_abs_mean * 100:>7.3f}%" + act_amax=act_amax, + quant_act=True, + fp4_block_size=bs, ) + errors[bs] = (out - ref).abs().mean().item() - errors = {} - for bs in block_sizes: - out = cuda_conv3d( - x, - w, - stride=(1, 1, 1), - padding=(1, 1, 1), - dilation=(1, 1, 1), - act_amax=act_amax, - quant_act=True, - fp4_block_size=bs, - ) - diff = (out - ref).abs() - max_d = diff.max().item() - mean_d = diff.mean().item() - rmse = ((out - ref) ** 2).mean().sqrt().item() - rel_err = mean_d / ref_abs_mean * 100 - errors[bs] = mean_d - print( - f" {bs:>10} | {max_d:>10.4f} | {mean_d:>10.6f} | " - f"{rmse:>10.4f} | {rel_err:>7.3f}%" - ) - all_errors[desc] = errors - - # Monotonicity check on the medium config - errors = all_errors["Medium K=1728"] + # Monotonicity: smaller blocks should have equal or lower error for smaller, larger in [(16, 64), (16, 256), (32, 256), (64, 256)]: assert errors[smaller] <= errors[larger] * 1.2, ( f"Expected error({smaller})={errors[smaller]:.6f} <= " @@ -590,14 +542,6 @@ def test_quant_realistic_shape(self, cuda_conv3d, fp4_block_size): # ============================================================================= -@pytest.fixture(scope="module") -def cuda_fp4(): - """Import and return the CUDA FP4 fake quant function.""" - from experimental.conv.implicit_gemm_cuda import fp4_fake_quant - - return fp4_fake_quant - - def _py_fp4_fake_quant_ref(x_flat, global_amax, block_size): """Pure Python reference for FP4 fake quant (no BF16 rounding). @@ -909,8 +853,8 @@ def _modelopt_dynamic_block_quantize_available(): class TestFP4FakeQuantVsModelopt: """Compare experimental CUDA FP4 fake quant against all modelopt FP4 implementations. - This ensures the standalone FP4 kernel in experimental/conv produces the same - results as the official modelopt quantization paths: + This ensures the standalone FP4 kernel produces the same results as the + other modelopt quantization paths: 1. Triton fp4_fake_quant_block (Hopper+ dynamic blockwise) 2. cuda_ext_mx.fused_amax_convert (CUDA extension fallback) 3. dynamic_block_quantize_op (high-level API that dispatches to either) @@ -1065,3 +1009,484 @@ def test_vs_triton_input_dtypes(self, cuda_fp4, dtype): # BF16/FP16 input rounding may cause small diffs tol = 1e-2 if dtype != torch.float32 else 1e-5 assert max_diff < tol, f"dtype={dtype}: experimental vs Triton max diff: {max_diff:.6e}" + + +# ============================================================================= +# Input Validation / Error Path Tests +# ============================================================================= + + +class TestConv3dInputValidation: + """Verify error paths raise appropriate exceptions.""" + + def test_invalid_fp4_block_size(self, cuda_conv3d): + """fp4_block_size not in {16, 32, 64, 128, 256} should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, 3, device="cuda") + with pytest.raises(ValueError, match="fp4_block_size"): + cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), fp4_block_size=7) + + def test_non_5d_input(self, cuda_conv3d): + """Non-5D tensors should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, device="cuda") # 4D + w = torch.randn(8, 4, 3, 3, 3, device="cuda") + with pytest.raises(ValueError, match="5D"): + cuda_conv3d(x, w) + + def test_non_5d_weight(self, cuda_conv3d): + """Non-5D weight should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, device="cuda") # 4D + with pytest.raises(ValueError, match="5D"): + cuda_conv3d(x, w) + + def test_grouped_conv_error(self, cuda_conv3d): + """Mismatched Cin (groups > 1) should raise ValueError.""" + x = torch.randn(1, 8, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, 3, device="cuda") # Cin=4 != x.Cin=8 + with pytest.raises(ValueError, match="Grouped convolution"): + cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + + def test_quant_act_without_amax(self, cuda_conv3d): + """quant_act=True without act_amax should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, 3, device="cuda") + with pytest.raises(ValueError, match="act_amax"): + cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=True, act_amax=None) + + def test_fp4_numel_not_divisible(self, cuda_fp4): + """fp4_fake_quant should error when numel is not divisible by block_size.""" + inp = torch.randn(17, device="cuda") + amax = torch.tensor([1.0], device="cuda") + with pytest.raises(AssertionError, match="divisible"): + cuda_fp4(inp, amax, block_size=16) + + +# ============================================================================= +# Input Dtype Tests +# ============================================================================= + + +class TestConv3dInputDtypes: + """Verify conv3d works with non-float32 inputs and preserves output dtype.""" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_preservation(self, cuda_conv3d, dtype): + """Output dtype should match input dtype.""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda", dtype=dtype) + w = torch.randn(32, 16, 3, 3, 3, device="cuda", dtype=dtype) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_correctness(self, cuda_conv3d, dtype): + """Non-float32 inputs should produce correct results (vs F.conv3d in float32).""" + torch.manual_seed(42) + x_fp32 = torch.randn(1, 16, 8, 8, 8, device="cuda") + w_fp32 = torch.randn(32, 16, 3, 3, 3, device="cuda") + x = x_fp32.to(dtype) + w = w_fp32.to(dtype) + + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + ref = F.conv3d(x_fp32, w_fp32, stride=(1, 1, 1), padding=(1, 1, 1)) + + # Both BF16 input rounding and internal BF16 WMMA contribute to error + max_diff = (out.float() - ref).abs().max().item() + k_size = 16 * 27 + scaled_atol = ATOL * (k_size / 1000.0) ** 0.5 * 2 # extra slack for input rounding + assert max_diff < scaled_atol, f"dtype={dtype}: max diff {max_diff:.4f} > {scaled_atol:.4f}" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_quant_path(self, cuda_conv3d, dtype): + """FP4 quantized path should also work with non-float32 inputs.""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda", dtype=dtype) + w = torch.randn(32, 16, 3, 3, 3, device="cuda", dtype=dtype) + act_amax = x.float().abs().max().unsqueeze(0) + + out = cuda_conv3d( + x, + w, + stride=(1, 1, 1), + padding=(1, 1, 1), + act_amax=act_amax, + quant_act=True, + fp4_block_size=16, + ) + assert out.dtype == dtype + assert not torch.isnan(out).any() + assert out.abs().max() > 0 + + +# ============================================================================= +# Non-Contiguous Input Tests +# ============================================================================= + + +class TestConv3dNonContiguous: + """Verify kernel handles non-contiguous tensors (via internal .contiguous() calls).""" + + def test_non_contiguous_input(self, cuda_conv3d): + """Permuted (non-contiguous) input should produce correct results.""" + torch.manual_seed(42) + # Create non-contiguous tensor via permute + permute back + x_base = torch.randn(1, 8, 8, 8, 16, device="cuda") + x = x_base.permute(0, 4, 1, 2, 3) # [1, 16, 8, 8, 8] but non-contiguous + assert not x.is_contiguous() + + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + x_contig = x.contiguous() + + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + ref = cuda_conv3d(x_contig, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.equal(out, ref), "Non-contiguous input produced different results" + + def test_non_contiguous_weight(self, cuda_conv3d): + """Transposed (non-contiguous) weight should produce correct results.""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + # Create non-contiguous weight + w_base = torch.randn(16, 32, 3, 3, 3, device="cuda") + w = w_base.transpose(0, 1) # [32, 16, ...] but non-contiguous + assert not w.is_contiguous() + + w_contig = w.contiguous() + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + ref = cuda_conv3d(x, w_contig, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.equal(out, ref), "Non-contiguous weight produced different results" + + +# ============================================================================= +# Combined Conv Parameter Tests +# ============================================================================= + + +class TestConv3dCombinedParams: + """Test combinations of stride + dilation + padding that were never combined.""" + + def test_stride_and_dilation(self, cuda_conv3d): + """Stride > 1 and dilation > 1 simultaneously.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (2, 2, 2)) + + def test_asymmetric_stride_and_padding(self, cuda_conv3d): + """Asymmetric stride with asymmetric padding.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 2, 2), (0, 1, 2), (1, 1, 1)) + + def test_all_non_default(self, cuda_conv3d): + """Non-default stride + padding + dilation all at once.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 2, 1), (1, 0, 1), (1, 2, 1)) + + def test_bias_with_stride(self, cuda_conv3d): + """Bias with non-default stride.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + b = torch.randn(32, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, b, (2, 2, 2), (1, 1, 1), (1, 1, 1)) + + def test_bias_with_dilation(self, cuda_conv3d): + """Bias with non-default dilation.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + b = torch.randn(32, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, b, (1, 1, 1), (2, 2, 2), (2, 2, 2)) + + +# ============================================================================= +# FP4 Quantized Path: Advanced Conv Params +# ============================================================================= + + +def _run_quant_smoke_test(cuda_conv3d, x, w, bias, stride, padding, dilation, fp4_block_size=16): + """Helper: run FP4-quantized conv and verify basic sanity.""" + act_amax = x.abs().max().unsqueeze(0) + out = cuda_conv3d( + x, + w, + bias=bias, + stride=stride, + padding=padding, + dilation=dilation, + act_amax=act_amax, + quant_act=True, + fp4_block_size=fp4_block_size, + ) + ref = F.conv3d(x, w, bias=bias, stride=stride, padding=padding, dilation=dilation) + assert out.shape == ref.shape, f"Shape mismatch: {out.shape} vs {ref.shape}" + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + # Quantized output should be in a reasonable range relative to reference + if ref.abs().max() > 0: + ratio = out.abs().max().item() / ref.abs().max().item() + assert 0.01 < ratio < 100, f"Output magnitude ratio {ratio:.2f} is unreasonable" + return out + + +class TestConv3dFP4QuantAdvanced: + """FP4 quantized path with non-trivial stride, dilation, and kernel shapes.""" + + def test_quant_with_stride(self, cuda_conv3d): + """FP4 quant with stride=(2,2,2).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (1, 1, 1)) + + def test_quant_with_dilation(self, cuda_conv3d): + """FP4 quant with dilation=(2,2,2).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (2, 2, 2), (2, 2, 2)) + + def test_quant_with_asymmetric_kernel(self, cuda_conv3d): + """FP4 quant with 1x3x3 kernel.""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 16, 16, device="cuda") + w = torch.randn(32, 16, 1, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 1, 1), (1, 1, 1)) + + def test_quant_with_stride_and_dilation(self, cuda_conv3d): + """FP4 quant with both stride>1 and dilation>1.""" + torch.manual_seed(42) + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (2, 2, 2)) + + def test_quant_with_no_padding(self, cuda_conv3d): + """FP4 quant with padding=(0,0,0).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + def test_quant_bias_reference(self, cuda_conv3d, cuda_fp4): + """FP4 quant + bias: verify bias is added correctly by comparing with/without. + + The difference between bias and no-bias output should equal the bias broadcast. + """ + torch.manual_seed(42) + cin, cout = 256, 64 + x = torch.randn(1, cin, 4, 4, 4, device="cuda") + w = torch.randn(cout, cin, 1, 1, 1, device="cuda") + b = torch.randn(cout, device="cuda") + act_amax = x.abs().max().unsqueeze(0) + + kwargs = { + "stride": (1, 1, 1), + "padding": (0, 0, 0), + "dilation": (1, 1, 1), + "act_amax": act_amax, + "quant_act": True, + "fp4_block_size": 16, + } + out_bias = cuda_conv3d(x, w, bias=b, **kwargs) + out_no_bias = cuda_conv3d(x, w, bias=None, **kwargs) + + # Difference should be the bias broadcast over spatial dims + diff = out_bias - out_no_bias # [1, Cout, D, H, W] + expected_bias = b.view(1, -1, 1, 1, 1).expand_as(diff) + assert torch.allclose(diff, expected_bias, atol=1e-5), ( + f"Bias diff mismatch: max {(diff - expected_bias).abs().max().item():.6e}" + ) + + +# ============================================================================= +# Zero / Degenerate Input Tests +# ============================================================================= + + +class TestConv3dZeroInputs: + """Tests with zero and degenerate inputs.""" + + def test_zero_input(self, cuda_conv3d): + """Zero activation tensor should produce zero (or bias-only) output.""" + x = torch.zeros(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.allclose(out, ref, atol=1e-5), f"Max diff: {(out - ref).abs().max().item()}" + + def test_zero_weight(self, cuda_conv3d): + """Zero weight tensor should produce zero output.""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + w = torch.zeros(32, 16, 3, 3, 3, device="cuda") + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.allclose(out, torch.zeros_like(out), atol=1e-5) + + def test_zero_input_quant(self, cuda_conv3d): + """Zero input with FP4 quant should not produce NaN.""" + x = torch.zeros(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + # act_amax=0 is a tricky edge case — the kernel's scale guard should handle it + act_amax = torch.tensor([1e-10], device="cuda") # near-zero but not exactly 0 + out = cuda_conv3d( + x, + w, + stride=(1, 1, 1), + padding=(1, 1, 1), + act_amax=act_amax, + quant_act=True, + fp4_block_size=16, + ) + assert not torch.isnan(out).any(), "Zero input with quant produced NaN" + + +# ============================================================================= +# Numerical Stability Tests +# ============================================================================= + + +class TestConv3dNumericalStability: + """Test with extreme value ranges.""" + + def test_large_values(self, cuda_conv3d): + """Large input values (randn * 100).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") * 100 + w = torch.randn(32, 16, 3, 3, 3, device="cuda") * 100 + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + # With large values, BF16 rounding error scales proportionally + rel_err = (out - ref).abs().max().item() / ref.abs().max().item() + assert rel_err < 0.05, f"Relative error {rel_err:.4f} too high for large values" + + def test_small_values(self, cuda_conv3d): + """Small input values (randn * 1e-3).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") * 1e-3 + w = torch.randn(32, 16, 3, 3, 3, device="cuda") * 1e-3 + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert out.shape == ref.shape + # Small values: absolute error is small, relative error may be larger due to BF16 + max_diff = (out - ref).abs().max().item() + assert max_diff < 1e-5, f"Max diff {max_diff:.6e} for small values" + + def test_uniform_input(self, cuda_conv3d): + """Uniform input (all ones) — exposes accumulation patterns.""" + x = torch.ones(1, 16, 8, 8, 8, device="cuda") + w = torch.ones(32, 16, 3, 3, 3, device="cuda") + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + k_size = 16 * 27 + scaled_atol = ATOL * (k_size / 1000.0) ** 0.5 + max_diff = (out - ref).abs().max().item() + assert max_diff < scaled_atol, f"Uniform input: max diff {max_diff:.4f}" + + +# ============================================================================= +# Exact Block Boundary Tests +# ============================================================================= + + +class TestConv3dExactBoundaries: + """Shapes that land exactly on BLOCK_M=64, BLOCK_N=64, BLOCK_K=256 boundaries.""" + + def test_m_exact_128(self, cuda_conv3d): + """M = 128 = 2 * BLOCK_M (exactly 2 M-tiles, no remainder).""" + # batch=1, output 4x4x8 = 128 with kernel 1x1x1 + x = torch.randn(1, 32, 4, 4, 8, device="cuda") + w = torch.randn(64, 32, 1, 1, 1, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + def test_k_exact_512(self, cuda_conv3d): + """K = 512 = 2 * BLOCK_K (exactly 2 K-tiles, no remainder).""" + # Cin=512, kernel 1x1x1 -> K=512 + x = torch.randn(1, 512, 4, 4, 4, device="cuda") + w = torch.randn(64, 512, 1, 1, 1, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + def test_cout_exact_64(self, cuda_conv3d): + """Cout = 64 = 1 * BLOCK_N (exactly 1 N-tile).""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(64, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (1, 1, 1), (1, 1, 1)) + + def test_all_exact_multiples(self, cuda_conv3d): + """M=64, N=64, K=256 — single tile in each dimension.""" + # batch=1, Cin=256, kernel 1x1x1 -> K=256; output 4x4x4=64; Cout=64 + x = torch.randn(1, 256, 4, 4, 4, device="cuda") + w = torch.randn(64, 256, 1, 1, 1, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + +# ============================================================================= +# FP4 Fake Quant: Shape and Edge Case Tests +# ============================================================================= + + +class TestFP4FakeQuantShapes: + """Test fp4_fake_quant with multi-dimensional inputs.""" + + def test_3d_shape_preservation(self, cuda_fp4): + """3D input should preserve shape after quantization.""" + inp = torch.randn(4, 8, 32, device="cuda") # numel=1024 + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + assert out.shape == (4, 8, 32) + + def test_4d_shape_preservation(self, cuda_fp4): + """4D input should preserve shape.""" + inp = torch.randn(2, 4, 8, 16, device="cuda") # numel=1024 + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + assert out.shape == (2, 4, 8, 16) + + def test_5d_shape_preservation(self, cuda_fp4): + """5D input (like a Conv3D activation) should preserve shape.""" + inp = torch.randn(1, 4, 4, 4, 16, device="cuda") # numel=1024 + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + assert out.shape == (1, 4, 4, 4, 16) + + def test_multidim_correctness(self, cuda_fp4): + """Multi-dim quantization should equal flatten -> quant -> reshape.""" + torch.manual_seed(42) + inp = torch.randn(4, 8, 32, device="cuda") + amax = inp.abs().max().unsqueeze(0) + + out_3d = cuda_fp4(inp, amax, block_size=16) + out_flat = cuda_fp4(inp.reshape(-1), amax, block_size=16).reshape(4, 8, 32) + assert torch.equal(out_3d, out_flat) + + +class TestFP4FakeQuantEdgeCases: + """Edge cases for fp4_fake_quant.""" + + def test_very_large_values(self, cuda_fp4): + """Very large input values should saturate to max E2M1 level, not produce NaN.""" + inp = torch.tensor([1e6, -1e6, 5e5, -5e5, 1e4, -1e4, 100, -100], device="cuda") + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=8) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_very_small_values(self, cuda_fp4): + """Very small input values should quantize to zero or near-zero.""" + inp = torch.tensor([1e-8, -1e-8, 1e-10, -1e-10, 1e-6, -1e-6, 0, 0], device="cuda") + amax = torch.tensor([1.0], device="cuda") + out = cuda_fp4(inp, amax, block_size=8) + assert not torch.isnan(out).any() + # Very small values relative to amax should quantize to ~0 + assert out.abs().max() < 1e-3 + + def test_uniform_block(self, cuda_fp4): + """All-same-value block.""" + inp = torch.full((16,), 3.0, device="cuda") + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + # All elements are the same, so they should all quantize to the same E2M1 level + assert (out == out[0]).all(), f"Uniform block produced non-uniform output: {out}" + + def test_near_zero_amax(self, cuda_fp4): + """Very small global_amax should not produce NaN/Inf.""" + inp = torch.randn(16, device="cuda") * 1e-8 + amax = torch.tensor([1e-10], device="cuda") + out = cuda_fp4(inp, amax, block_size=16) + assert not torch.isnan(out).any(), "Near-zero amax produced NaN" + assert not torch.isinf(out).any(), "Near-zero amax produced Inf" diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py new file mode 100644 index 00000000000..aa2f32829d3 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the quantized WanCausalConv3d wrapper in the diffusers plugin. + +WanCausalConv3d applies asymmetric causal padding before the underlying Conv3D, +so the quantized subclass has to replicate that padding logic around the +quantized forward. These tests run on CPU with the NVFP4 dispatch disabled (by +only exercising the default quantized path); the NVFP4 implicit-GEMM path is +covered by the GPU tests under ``tests/gpu/torch/quantization/kernels``. +""" + +import pytest +import torch + +pytest.importorskip("diffusers") + +from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + +# Triggers registration of _QuantDiffusersWanCausalConv3d. +import modelopt.torch.quantization.plugins.diffusion.diffusers # noqa: F401 +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantDiffusersWanCausalConv3d + + +def _make_quantized(in_ch: int = 4, out_ch: int = 6, padding=(1, 1, 1)) -> torch.nn.Module: + m = WanCausalConv3d(in_ch, out_ch, kernel_size=3, padding=padding) + m.eval() + # Convert via the registry so the generated class picks up WanCausalConv3d + # in its MRO (needed for ``nn.Conv3d._conv_forward`` to be reachable). + return QuantModuleRegistry.convert(m) + + +def _strip_quant_state(state_dict: dict) -> dict: + return { + k: v + for k, v in state_dict.items() + if not any( + k.startswith(p) for p in ("input_quantizer", "weight_quantizer", "output_quantizer") + ) + } + + +class TestQuantWanCausalConv3dRegistration: + def test_registered(self): + assert WanCausalConv3d in QuantModuleRegistry + # Registry-returned class is a generated subclass of our dm class. + assert issubclass(QuantModuleRegistry[WanCausalConv3d], _QuantDiffusersWanCausalConv3d) + + def test_convert_preserves_type_identity(self): + m = _make_quantized() + mro = [c.__name__ for c in type(m).__mro__] + assert "WanCausalConv3d" in mro + assert "_QuantDiffusersWanCausalConv3d" in mro + + def test_no_implicit_gemm_without_nvfp4(self): + # Default quantizer config is INT8 — must NOT route through implicit GEMM. + m = _make_quantized() + assert not m._should_use_implicit_gemm() + + +class TestQuantWanCausalConv3dForward: + """Exercise the default quantized path (NVFP4 kernel is GPU-only). + + We disable the quantizers and assert output matches an unquantized + ``WanCausalConv3d`` with the same weights — this verifies the causal-padding + logic in the overridden ``forward`` is preserved after conversion. + """ + + @pytest.mark.parametrize("padding", [(1, 1, 1), (2, 0, 0), (0, 1, 1)]) + def test_matches_unquantized_no_cache(self, padding): + torch.manual_seed(0) + m = _make_quantized(in_ch=4, out_ch=6, padding=padding) + m.input_quantizer.disable() + m.weight_quantizer.disable() + + m_ref = WanCausalConv3d(4, 6, kernel_size=3, padding=padding) + m_ref.eval() + m_ref.load_state_dict(_strip_quant_state(m.state_dict()), strict=False) + + x = torch.randn(1, 4, 5, 6, 6) + out = m(x) + out_ref = m_ref(x) + assert torch.allclose(out, out_ref, atol=1e-5), ( + f"Max diff: {(out - out_ref).abs().max().item()}" + ) + + def test_matches_unquantized_with_cache_x(self): + """cache_x is the temporal-cache branch used during causal decoding.""" + torch.manual_seed(0) + m = _make_quantized(padding=(1, 1, 1)) # _padding[4] == 2 > 0 → cache path active + m.input_quantizer.disable() + m.weight_quantizer.disable() + + m_ref = WanCausalConv3d(4, 6, kernel_size=3, padding=(1, 1, 1)) + m_ref.eval() + m_ref.load_state_dict(_strip_quant_state(m.state_dict()), strict=False) + + x = torch.randn(1, 4, 5, 6, 6) + cache_x = torch.randn(1, 4, 1, 6, 6) + out = m(x, cache_x=cache_x) + out_ref = m_ref(x, cache_x=cache_x) + assert torch.allclose(out, out_ref, atol=1e-5) + + def test_output_quantizer_applied(self): + """Enabling the output quantizer must change the forward output.""" + torch.manual_seed(0) + m = _make_quantized() + m.input_quantizer.disable() + m.weight_quantizer.disable() + # Output quantizer is disabled by default; enable it and check it takes effect. + x = torch.randn(1, 4, 3, 5, 5) + out_disabled = m(x) + m.output_quantizer.enable() + out_enabled = m(x) + # INT8 default config clamps and rounds; at least some elements differ. + assert not torch.allclose(out_disabled, out_enabled) diff --git a/tests/unit/torch/quantization/test_quant_conv.py b/tests/unit/torch/quantization/test_quant_conv.py index be872de206d..ee1d449d6e6 100644 --- a/tests/unit/torch/quantization/test_quant_conv.py +++ b/tests/unit/torch/quantization/test_quant_conv.py @@ -15,6 +15,8 @@ """Tests of QuantConv module.""" +import warnings + import pytest import torch import torch.nn.functional as F @@ -23,11 +25,19 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn.modules import quant_conv +from modelopt.torch.quantization.nn.modules.quant_conv import ( + _is_nvfp4_quantizer, + _nvfp4_quantize_weight_along_k, +) from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer NUM_IN_CHANNELS = 3 NUM_OUT_CHANNELS = 5 +_NVFP4_CFG = QuantizerAttributeConfig( + num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)} +) + class TestQuantConvND: @pytest.mark.parametrize( @@ -316,3 +326,134 @@ def test_against_unquantized(self, conv_cls, nn_conv_cls, input_shape): output = conv(test_input) assert torch.allclose(quant_output, output) + + +class TestQuantConv3dNVFP4: + """Tests for the NVFP4 implicit GEMM dispatch path in ``_QuantConv3d``. + + The CUDA kernel itself is GPU-only (covered in ``tests/gpu``); here we only + exercise the CPU-side predicate and fallback branches that currently lack + coverage. + """ + + @staticmethod + def _make_nvfp4_conv3d(groups: int = 1, bias: bool = False) -> quant_conv.QuantConv3d: + return quant_conv.QuantConv3d( + NUM_IN_CHANNELS * groups, + NUM_OUT_CHANNELS * groups, + kernel_size=3, + groups=groups, + bias=bias, + quant_desc_input=_NVFP4_CFG, + quant_desc_weight=_NVFP4_CFG, + ) + + def test_is_nvfp4_quantizer_true(self): + q = TensorQuantizer(_NVFP4_CFG) + assert _is_nvfp4_quantizer(q) + + def test_is_nvfp4_quantizer_false_for_int8(self): + q = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + assert not _is_nvfp4_quantizer(q) + + def test_is_nvfp4_quantizer_false_for_fp8(self): + # FP8 E4M3: num_bits == (4, 3), block_sizes is None + q = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + assert not _is_nvfp4_quantizer(q) + + def test_is_nvfp4_quantizer_false_for_static_block(self): + static_cfg = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + q = TensorQuantizer(static_cfg) + assert not _is_nvfp4_quantizer(q) + + def test_should_use_implicit_gemm_true(self): + m = self._make_nvfp4_conv3d() + assert m._should_use_implicit_gemm() + + def test_should_use_implicit_gemm_false_groups_gt_1(self): + m = self._make_nvfp4_conv3d(groups=NUM_IN_CHANNELS) + # Groups > 1 disqualifies the fused kernel even with NVFP4 quantizers. + assert not m._should_use_implicit_gemm() + + def test_should_use_implicit_gemm_false_non_nvfp4(self): + # INT8 per-tensor config — default Conv3D quantizers, no NVFP4. + m = quant_conv.QuantConv3d(NUM_IN_CHANNELS, NUM_OUT_CHANNELS, kernel_size=3, bias=False) + assert not m._should_use_implicit_gemm() + + def test_nvfp4_quantize_weight_along_k_reshape(self): + """Verify weight is flattened/restored along the K (input) dimension. + + Uses a stub quantizer so this test stays CPU-only (the real NVFP4 dynamic + quantizer requires a CUDA tensor). + """ + + def identity_quantizer(x): + # K-dim must be the last axis when passed to the quantizer so NVFP4's + # block-wise scaling aligns with the GEMM reduction axis. + assert x.dim() == 2 and x.shape[0] == NUM_OUT_CHANNELS + return x + + w = torch.randn(NUM_OUT_CHANNELS, NUM_IN_CHANNELS, 3, 3, 3) + qw = _nvfp4_quantize_weight_along_k(w, identity_quantizer) + assert qw.shape == w.shape + assert torch.equal(qw, w) + + def test_forward_non_nvfp4_matches_unquantized(self): + # Disabled quantizers: forward must match plain conv3d exactly. + m = quant_conv.QuantConv3d(NUM_IN_CHANNELS, NUM_OUT_CHANNELS, kernel_size=3, bias=False) + m.input_quantizer.disable() + m.weight_quantizer.disable() + x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4) + out = m(x) + ref = F.conv3d(x, m.weight) + assert torch.allclose(out, ref) + + def test_forward_nvfp4_training_warns_and_falls_back(self): + """Training mode must fall back to the default (cuDNN) path with a warning. + + The implicit-GEMM kernel is inference-only; this exercises the CPU-visible + training-fallback branch. We disable the quantizers so the default-path + ``super().forward()`` does not try to run NVFP4 dynamic quantization + (which requires CUDA). + """ + m = self._make_nvfp4_conv3d() + # NVFP4 predicate reads configuration (num_bits/block_sizes), not enable + # state, so disabling the quantizers still routes through the NVFP4 branch. + m.input_quantizer.disable() + m.weight_quantizer.disable() + assert m._should_use_implicit_gemm() + m.train() + x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + out = m(x) + assert out.shape == (1, NUM_OUT_CHANNELS, 2, 2, 2) + assert any("inference-only" in str(w.message) for w in caught), ( + f"Expected an 'inference-only' warning, got: {[str(w.message) for w in caught]}" + ) + + def test_forward_nvfp4_calib_only_uses_default_path(self): + """Calibration-only mode must NOT try to invoke the CUDA kernel. + + When the input quantizer is in calibration mode without quant enabled, + the forward must fall back to the default path. This exercises the + calib-only early-return branch; we assert that the output matches the + default-path output exactly (i.e. the implicit-GEMM path wasn't used). + """ + m = self._make_nvfp4_conv3d() + m.eval() + # Match the state toggled by TensorQuantizer.disable_quant()/enable_calib(). + m.input_quantizer.disable_quant() + m.input_quantizer.enable_calib() + m.weight_quantizer.disable_quant() + m.weight_quantizer.enable_calib() + assert m.input_quantizer._if_calib and not m.input_quantizer._if_quant + + x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4) + out = m(x) + # Default path with quant disabled should equal plain conv3d. + ref = F.conv3d(x, m.weight.detach()) + assert torch.allclose(out, ref, atol=1e-5) From 2dd43d085bb2abc49d1274152c3aca1fced422f5 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Sat, 18 Apr 2026 07:45:26 +0000 Subject: [PATCH 2/4] Update the codecov.yml Signed-off-by: Jingyu Xin --- .github/codecov.yml | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/.github/codecov.yml b/.github/codecov.yml index b4ac8367690..24756fdcbb2 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -11,15 +11,3 @@ coverage: target: auto threshold: 1% # Allow atmost 1% coverage drop from main branch. patch: false - -# Exclude GPU-only Triton kernel files from ALL codecov calculations (project -# and patch checks, all flags). Rationale: these files are dominated by -# @triton.jit kernel bodies that CPU unit tests cannot exercise. GPU tests -# cover them end-to-end (see tests/gpu/torch/sparsity/attention_sparsity/) but -# the `gpu`-flag upload may race with the PR status check, so relying on flag -# combination alone leaves the project check flaky. Dropping these files here -# makes the check deterministic — local `pytest --cov` and GPU runs still -# measure them; only the codecov PR status ignores them. -ignore: - - "modelopt/torch/kernels/triton_fa.py" - - "modelopt/torch/kernels/hf_triton_attention.py" From eaaaa0f664a27a8a183e8d200493062d44285446 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Sat, 18 Apr 2026 08:26:18 +0000 Subject: [PATCH 3/4] Update the test case Signed-off-by: Jingyu Xin --- .../torch/quantization/plugins/test_diffusers_wan_conv3d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py index aa2f32829d3..dcd8f24ab13 100644 --- a/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py +++ b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py @@ -25,7 +25,7 @@ import pytest import torch -pytest.importorskip("diffusers") +pytest.importorskip("onnx") from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d From c77aaed028c7094cf5fe0d0de95514f42ac99ac1 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Sat, 18 Apr 2026 09:16:22 +0000 Subject: [PATCH 4/4] Update the test case & add the cuda_arch >= 800 check & update readme Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 14 + .../torch/quantization/src/conv/README.md | 1 + .../src/conv/implicit_gemm_kernel.cu | 5 + tests/examples/diffusers/test_diffusers.py | 247 +++++++----------- .../test_export_diffusers_hf_ckpt.py | 114 ++++---- 5 files changed, 185 insertions(+), 196 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 84f248bfb15..ac14d982279 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -117,6 +117,20 @@ python quantize.py \ --hf-ckpt-dir ./hf_ckpt ``` +#### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM) + +The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/quantization/src/conv/README.md`](../../modelopt/torch/quantization/src/conv/README.md) for kernel details. + +```sh +python quantize.py \ + --model {wan2.2-t2v-14b|wan2.2-t2v-5b} \ + --backbone vae \ + --format fp4 --quant-algo max --collect-method default \ + --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \ + --batch-size 1 --calib-size 32 --n-steps 30 \ + --quantized-torch-ckpt-save-path ./wan22_vae_fp4.pt +``` + #### [LTX-2](https://github.com/Lightricks/LTX-2) FP4 > [!WARNING] diff --git a/modelopt/torch/quantization/src/conv/README.md b/modelopt/torch/quantization/src/conv/README.md index 37580b69a80..6b14fd5953b 100644 --- a/modelopt/torch/quantization/src/conv/README.md +++ b/modelopt/torch/quantization/src/conv/README.md @@ -125,6 +125,7 @@ Known limitations: - Output shape matches `torch.nn.functional.conv3d`. - FP4 path applies quantize-dequantize in-kernel for activation tiles (no extra global memory pass). - Tile config: BLOCK_M=64, BLOCK_N=64, BLOCK_K=256, 8 warps (256 threads), ~70 KB shared memory per block. +- The kernel body is guarded by `#if __CUDA_ARCH__ >= 800` so it compiles as an empty stub when nvcc targets pre-Ampere archs (PyTorch's default `-gencode` list can include sm_75, which lacks BF16 WMMA fragments). Dispatch is enforced at runtime by `_get_cuda_module()` via `_MIN_SM_MAJOR = 8`. ## Files diff --git a/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu b/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu index a3b40f48481..10d20c2e379 100644 --- a/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu +++ b/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu @@ -145,6 +145,10 @@ __global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2) const float *__restrict__ act_amax, int Cin, int Dp, int Hp, int Wp, int Cout, int OD, int OH, int OW, int kD, int kH, int kW, int sd, int sh, int sw, int dd, int dh, int dw, int M, int K) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + // BF16 WMMA fragments are only available on sm_80+. On earlier archs this + // kernel compiles as an empty stub; the runtime gate in _get_cuda_module() + // prevents dispatching it on unsupported hardware. // Derived constants constexpr int NUM_WARPS = WARPS_M * WARPS_N; constexpr int NUM_THREADS = NUM_WARPS * 32; @@ -469,6 +473,7 @@ __global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2) y[m_idx * Cout + n_idx] = result; } } +#endif // __CUDA_ARCH__ >= 800 } // ============================================================================= diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 9af6481d918..5b117b41b3f 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -21,27 +21,6 @@ from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm -# Tiny video model args — override MODEL_DEFAULTS for fast CI -_WAN22_TINY_EXTRA_PARAMS = [ - "--extra-param", - "height=16", - "--extra-param", - "width=16", - "--extra-param", - "num_frames=5", -] - -_WAN22_FAST_CALIB_ARGS = [ - "--calib-size", - "2", - "--batch-size", - "1", - "--n-steps", - "2", - "--model-dtype", - "BFloat16", -] - class DiffuserModel(NamedTuple): dtype: str @@ -172,143 +151,117 @@ def test_diffusers_quantization( model.inference(tmp_path) -def _run_wan22_quantize( - tiny_wan22_path: str, - tmp_path: Path, - format_type: str, - quant_algo: str, - collect_method: str, - *, - model: str = "wan2.2-t2v-14b", - backbone: str | None = None, - ckpt_suffix: str | None = None, -) -> None: - """Run quantize.py for Wan 2.2 with the tiny model.""" - suffix = ckpt_suffix or format_type - ckpt_path = str(tmp_path / f"wan22_{suffix}.pt") - cmd_args = [ - "python", - "quantize.py", - "--model", - model, - "--override-model-path", - tiny_wan22_path, - "--format", - format_type, - "--quant-algo", - quant_algo, - "--collect-method", - collect_method, - "--trt-high-precision-dtype", - "BFloat16", - "--quantized-torch-ckpt-save-path", - ckpt_path, - *_WAN22_FAST_CALIB_ARGS, - *_WAN22_TINY_EXTRA_PARAMS, - ] - if backbone is not None: - cmd_args.extend(["--backbone", backbone]) - run_example_command(cmd_args, "diffusers/quantization") +class Wan22Model(NamedTuple): + model: str + backbone: str | None + format_type: str + quant_algo: str + collect_method: str + def _ckpt_path(self, tmp_path: Path) -> str: + stem = self.model.replace("wan2.2-t2v-", "") + parts = [stem, *([self.backbone] if self.backbone else []), self.format_type] + return str(tmp_path / f"wan22_{'_'.join(parts)}.pt") -def _run_wan22_restore( - tiny_wan22_path: str, - tmp_path: Path, - format_type: str, - quant_algo: str, - collect_method: str, - *, - model: str = "wan2.2-t2v-14b", - backbone: str | None = None, - ckpt_suffix: str | None = None, -) -> None: - """Restore a Wan 2.2 quantized checkpoint.""" - suffix = ckpt_suffix or format_type - ckpt_path = str(tmp_path / f"wan22_{suffix}.pt") - cmd_args = [ - "python", - "quantize.py", - "--model", - model, - "--override-model-path", - tiny_wan22_path, - "--format", - format_type, - "--quant-algo", - quant_algo, - "--collect-method", - collect_method, - "--trt-high-precision-dtype", - "BFloat16", - "--restore-from", - ckpt_path, - *_WAN22_FAST_CALIB_ARGS, - *_WAN22_TINY_EXTRA_PARAMS, - ] - if backbone is not None: - cmd_args.extend(["--backbone", backbone]) - run_example_command(cmd_args, "diffusers/quantization") + def _common_args(self, tiny_wan22_path: str) -> list[str]: + cmd_args = [ + "python", + "quantize.py", + "--model", + self.model, + "--override-model-path", + tiny_wan22_path, + "--format", + self.format_type, + "--quant-algo", + self.quant_algo, + "--collect-method", + self.collect_method, + "--model-dtype", + "BFloat16", + "--trt-high-precision-dtype", + "BFloat16", + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + # Tiny video dims — override MODEL_DEFAULTS for fast CI. + "--extra-param", + "height=16", + "--extra-param", + "width=16", + "--extra-param", + "num_frames=5", + ] + if self.backbone is not None: + cmd_args.extend(["--backbone", self.backbone]) + return cmd_args + def quantize(self, tiny_wan22_path: str, tmp_path: Path) -> None: + run_example_command( + [ + *self._common_args(tiny_wan22_path), + "--quantized-torch-ckpt-save-path", + self._ckpt_path(tmp_path), + ], + "diffusers/quantization", + ) -def test_wan22_int8_smoothquant(tiny_wan22_path: str, tmp_path: Path) -> None: - """Wan 2.2 INT8 SmoothQuant: quantize + restore.""" - _run_wan22_quantize(tiny_wan22_path, tmp_path, "int8", "smoothquant", "min-mean") - _run_wan22_restore(tiny_wan22_path, tmp_path, "int8", "smoothquant", "min-mean") + def restore(self, tiny_wan22_path: str, tmp_path: Path) -> None: + run_example_command( + [*self._common_args(tiny_wan22_path), "--restore-from", self._ckpt_path(tmp_path)], + "diffusers/quantization", + ) +# The VAE (``AutoencoderKLWan``) is shared between Wan 2.2 14B and 5B, so the +# Conv3D NVFP4 implicit-GEMM dispatch exercises the same kernel either way; we +# parametrize both ``--model`` values to also cover the ``quantize.py`` dispatch +# for each. @pytest.mark.parametrize( - ("format_type", "quant_algo"), + "wan_model", [ - pytest.param("fp8", "max", marks=minimum_sm(89)), - pytest.param("fp4", "max", marks=minimum_sm(89)), + Wan22Model("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"), + pytest.param( + Wan22Model("wan2.2-t2v-14b", None, "fp8", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-14b", None, "fp4", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-14b", "vae", "fp8", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-14b", "vae", "fp4", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-5b", "vae", "fp8", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-5b", "vae", "fp4", "max", "default"), + marks=minimum_sm(89), + ), ], - ids=["wan22_fp8_max", "wan22_fp4_max"], -) -def test_wan22_fp8_fp4( - tiny_wan22_path: str, tmp_path: Path, format_type: str, quant_algo: str -) -> None: - """Wan 2.2 FP8/FP4: quantize + restore (requires SM89+).""" - _run_wan22_quantize(tiny_wan22_path, tmp_path, format_type, quant_algo, "default") - _run_wan22_restore(tiny_wan22_path, tmp_path, format_type, quant_algo, "default") - - -@pytest.mark.parametrize( - ("format_type", "quant_algo"), - [ - pytest.param("fp8", "max", marks=minimum_sm(89)), - pytest.param("fp4", "max", marks=minimum_sm(89)), + ids=[ + "wan22_14b_transformer_int8_smoothquant", + "wan22_14b_transformer_fp8_max", + "wan22_14b_transformer_fp4_max", + "wan22_14b_vae_fp8_max", + "wan22_14b_vae_fp4_max", + "wan22_5b_vae_fp8_max", + "wan22_5b_vae_fp4_max", ], - ids=["wan22_vae_fp8_max", "wan22_vae_fp4_max"], ) -def test_wan22_vae_fp8_fp4( - tiny_wan22_path: str, tmp_path: Path, format_type: str, quant_algo: str -) -> None: - """Wan 2.2 VAE FP8/FP4 quantization: quantize + restore. - - Exercises the ``WanCausalConv3d``/Conv3D NVFP4 implicit-GEMM dispatch path - on the Wan VAE end-to-end through ``quantize.py``. - """ - ckpt_suffix = f"vae_{format_type}" - _run_wan22_quantize( - tiny_wan22_path, - tmp_path, - format_type, - quant_algo, - "default", - model="wan2.2-t2v-5b", - backbone="vae", - ckpt_suffix=ckpt_suffix, - ) - _run_wan22_restore( - tiny_wan22_path, - tmp_path, - format_type, - quant_algo, - "default", - model="wan2.2-t2v-5b", - backbone="vae", - ckpt_suffix=ckpt_suffix, - ) +def test_wan22_quantization(wan_model: Wan22Model, tiny_wan22_path: str, tmp_path: Path) -> None: + wan_model.quantize(tiny_wan22_path, tmp_path) + wan_model.restore(tiny_wan22_path, tmp_path) @pytest.mark.parametrize( diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 5dca512b5b9..a5c81d36937 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -21,16 +21,6 @@ from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm -# Tiny video model args — override MODEL_DEFAULTS for fast CI -_WAN22_TINY_EXTRA_PARAMS = [ - "--extra-param", - "height=16", - "--extra-param", - "width=16", - "--extra-param", - "num_frames=5", -] - class DiffuserHfExportModel(NamedTuple): name: str @@ -140,51 +130,77 @@ def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path) assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" +class Wan22HfExportModel(NamedTuple): + model: str + backbone: str | None + format_type: str + quant_algo: str + collect_method: str + + def _suffix(self) -> str: + stem = self.model.replace("wan2.2-t2v-", "") + parts = [stem, *([self.backbone] if self.backbone else []), self.format_type] + return "_".join(parts) + + def quantize_and_export_hf(self, tiny_wan22_path: str, tmp_path: Path) -> Path: + hf_ckpt_dir = tmp_path / f"wan22_{self._suffix()}_hf_ckpt" + cmd_args = [ + "python", + "quantize.py", + "--model", + self.model, + "--override-model-path", + tiny_wan22_path, + "--format", + self.format_type, + "--quant-algo", + self.quant_algo, + "--collect-method", + self.collect_method, + "--model-dtype", + "BFloat16", + "--trt-high-precision-dtype", + "BFloat16", + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + # Tiny video dims — override MODEL_DEFAULTS for fast CI. + "--extra-param", + "height=16", + "--extra-param", + "width=16", + "--extra-param", + "num_frames=5", + "--hf-ckpt-dir", + str(hf_ckpt_dir), + ] + if self.backbone is not None: + cmd_args.extend(["--backbone", self.backbone]) + run_example_command(cmd_args, "diffusers/quantization") + return hf_ckpt_dir + + @pytest.mark.parametrize( - ("format_type", "quant_algo", "collect_method"), + "wan_model", [ - ("int8", "smoothquant", "min-mean"), - pytest.param("fp8", "max", "default", marks=minimum_sm(89)), + Wan22HfExportModel("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"), + pytest.param( + Wan22HfExportModel("wan2.2-t2v-14b", None, "fp8", "max", "default"), + marks=minimum_sm(89), + ), + ], + ids=[ + "wan22_14b_transformer_int8_smoothquant", + "wan22_14b_transformer_fp8_max", ], - ids=["wan22_int8_smoothquant", "wan22_fp8_max"], ) def test_wan22_hf_ckpt_export( - tiny_wan22_path: str, - tmp_path: Path, - format_type: str, - quant_algo: str, - collect_method: str, + wan_model: Wan22HfExportModel, tiny_wan22_path: str, tmp_path: Path ) -> None: - """Quantize tiny Wan 2.2 and export to HF checkpoint.""" - hf_ckpt_dir = tmp_path / f"wan22_{format_type}_hf_ckpt" - cmd_args = [ - "python", - "quantize.py", - "--model", - "wan2.2-t2v-14b", - "--override-model-path", - tiny_wan22_path, - "--format", - format_type, - "--quant-algo", - quant_algo, - "--collect-method", - collect_method, - "--model-dtype", - "BFloat16", - "--trt-high-precision-dtype", - "BFloat16", - "--calib-size", - "2", - "--batch-size", - "1", - "--n-steps", - "2", - "--hf-ckpt-dir", - str(hf_ckpt_dir), - *_WAN22_TINY_EXTRA_PARAMS, - ] - run_example_command(cmd_args, "diffusers/quantization") + hf_ckpt_dir = wan_model.quantize_and_export_hf(tiny_wan22_path, tmp_path) assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}"