feat: support torchao nvfp4 dynamic quantization

docs/user_guide/QUANTIZATION.md

@@ -16,6 +16,8 @@ Quantization is a powerful technique to reduce the memory footprint and computat
 |<span style="color:#c77dff;">int8_per_tensor</span>|quantize weights and activations to int8 (<span style="color:green;">dynamic quantization</span>) with tensorwise method.|<span style="color:#c77dff;">>=sm80</span>, Ampere or newer|
 |<span style="color:#c77dff;">int8_weight_only</span>|quantize <span style="color:green;">only weights</span> to int8, keep activations in full precision|<span style="color:#c77dff;">>=sm80</span>, Ampere or newer|
 |<span style="color:#c77dff;">int4_weight_only</span>|quantize <span style="color:green;">only weights</span> to int4, keep activations in full precision|<span style="color:#c77dff;">>=sm90</span>, Hopper or newer, TMA required|
+|<span style="color:#c77dff;">nvfp4</span>|TorchAO dynamic activation and NVFP4 weight quantization. The current integration uses Triton activation quantization and dynamic per-tensor scaling.|<span style="color:#c77dff;">>=sm100</span>, Blackwell or newer|
+|<span style="color:#c77dff;">nvfp4_weight_only</span>|TorchAO NVFP4 weight-only quantization with dynamic per-tensor weight scaling; activations remain in their original precision.|<span style="color:#c77dff;">>=sm100</span>, Blackwell or newer|
 |<span style="color:#c77dff;">svdq_int4_r{32...}</span>|post-training <span style="color:green;">SVDQuant (W4A4)</span> with calibration and checkpoint serialization for users who want the higher-accuracy PTQ workflow.|<span style="color:#c77dff;">>=sm80</span>, Ampere or newer, excluded Hopper (NO INT4 MMA)|
 |<span style="color:#c77dff;">svdq_nvfp4_r{32...}</span>|post-training <span style="color:green;">SVDQuant (W4A4)</span> with NVFP4 packed weights/activations, calibration, and checkpoint serialization. <span style="color:green;">Only</span> <span style="color:#c77dff;">svdq_kwargs["runtime_kernel"]="v1"</span> is currently supported.|<span style="color:#c77dff;">>=sm120</span>, Blackwell or newer|
 |<span style="color:#c77dff;">svdq_int4_r{32...}_dq</span>|quantize weights and activations to int4 with <span style="color:green;">SVDQuant dynamic quantization (W4A4)</span> without any calibration.|<span style="color:#c77dff;">>=sm80</span>, Ampere or newer, excluded Hopper (NO INT4 MMA)|
@@ -289,13 +291,42 @@ INT4 quantization can provide even better memory reduction compared to FP8 or IN
 Please note that users should also install <span style="color:#c77dff;">mslk</span> kernel library to enable INT8/INT4 quantization features. The <span style="color:#c77dff;">int4_weight_only</span> w4a16 compute kennel requires architectures >= <span style="color:#c77dff;">sm90</span> (Hopper or newer, TMA required). For older architectures, users can use <span style="color:#c77dff;">int8_weight_only</span> quantization for better compatibility. 
 
 ```bash
-# stable: mslk (change cu130 to cu129 if using CUDA 12.9), required torch>=2.11.0
-uv pip install torch==2.11.0 mslk --index-url https://download.pytorch.org/whl/cu130 --upgrade
-# nightly: mslk (change cu130 to cu129 if using CUDA 12.9), required torch>=2.11.0
-uv pip install --pre torch mslk --index-url https://download.pytorch.org/whl/nightly/cu130 --upgrade
+# stable: mslk (change cu13x to cu129 if using CUDA 12.9), torch>=2.11.0
+# mslk version matching: torch==2.11.x mslk==1.1.0; torch==2.12.x mslk==1.2.0
+pip install torch==2.11.0 mslk==1.1.0 --index-url https://download.pytorch.org/whl/cu130
+pip install torch==2.12.0 mslk==1.2.0 --index-url https://download.pytorch.org/whl/cu132
+# nightly: mslk (change cu13x to cu129 if using CUDA 12.9), required torch>=2.11.0
+pip install --pre torch mslk --index-url https://download.pytorch.org/whl/nightly/cu132
 ```
+
 In the case of <span style="color:#c77dff;">distributed inference</span> (context parallelism or tensor parallelism), we recommend users to use <span style="color:#c77dff;">float8 quantization</span> to avoid potential compatibility issues.
 
+
+## TorchAO NVFP4 Quantization
+
+<span style="color:#c77dff;">nvfp4</span> uses <span style="color:green;">NVFP4DynamicActivationNVFP4WeightConfig</span> from TorchAO's MX formats prototype. Cache-DiT currently fixes both <span style="color:#c77dff;">use_triton_kernel=True</span> and <span style="color:#c77dff;">use_dynamic_per_tensor_scale=True</span>. It is an online dynamic quantization path and is separate from the <span style="color:#c77dff;">svdq_nvfp4_*</span> SVDQuant path. (Required <span style="color:green;">mslk</span> kernels library)
+
+For example:
+
+```python
+import cache_dit
+from cache_dit import QuantizeConfig  
+
+cache_dit.enable_cache( 
+  pipe, cache_config=..., quantize_config=QuantizeConfig(quant_type="nvfp4"), 
+)
+```
+
+The current implementation requires Blackwell or newer GPUs (SM100+) and skips Linear layers whose two weight dimensions are not divisible by 16 or whose output dimensions are too small for the NVFP4 kernel. The supported CLI shortcut is:
+
+```bash
+python3 -m cache_dit.generate flux --nvfp4
+python3 -m cache_dit.generate flux --nvfp4 --compile
+python3 -m cache_dit.generate flux --nvfp4-weight-only
+```
+
+The model must use bfloat16 Linear weights. The Triton activation kernel can fall back internally when a runtime shape does not satisfy its kernel constraints.
+
 ## SVDQuant (W4A4) PTQ
 
 Cache-DiT provides a native SVDQuant PTQ workflow for W4A4 quantization (with high performance <span style="color:green;">W4A4 GEMM kernels</span> and an <span style="color:green;">easy-to-use</span> PTQ interface). The public API is intentionally small: build a <span style="color:#c77dff;">QuantizeConfig</span>, quantize with <span style="color:#c77dff;">cache_dit.quantize(...)</span>, then reload with <span style="color:#c77dff;">cache_dit.load(...)</span>. Cache-DiT now supports both <span style="color:#c77dff;">INT4</span> and <span style="color:#c77dff;">NVFP4</span> SVDQuant PTQ flows. We highly recommend using native SVDQuant support in Cache-DiT for W4A4 quantization, as it can provide high performance and better usability compared to other third-party low-bit quantization libraries.   
@@ -824,15 +855,15 @@ Conversion complete. Load the quantized model with:
 **Experimental weight-only smooth strategy:**
 
 ```bash
-# INT4
+# SVDQ INT4
 cache-dit-convert \
   --model-path black-forest-labs/FLUX.2-klein-4B \
   --save-dir ./FLUX.2-klein-4B-svdq \
   --quant-type svdq-int4-r256-dq \
   --svdq-smooth-strategy weight \
   --svdq-calibrate-precision medium
 
-# NVFP4
+# SVDQ NVFP4
 cache-dit-convert \
   --model-path black-forest-labs/FLUX.2-klein-4B \
   --save-dir ./FLUX.2-klein-4B-svdq-nvfp4 \
@@ -844,7 +875,7 @@ cache-dit-convert \
 **With v2 runtime kernel and verbose logging:**
 
 ```bash
-# INT4 only, NVFP4 DQ is not supported with v2 runtime kernel for now.
+# SVDQ INT4 only, SVDQ NVFP4 DQ is not supported with v2 runtime kernel for now.
 cache-dit-convert \
   --model-path /path/to/FLUX.1-dev \
   --save-dir ./FLUX.1-dev-svdq \

pyproject.toml

@@ -28,7 +28,7 @@ ray = [
 ]
 
 quantization = [
-    "torchao>=0.14.1",
+    "torchao>=0.17.0",
     "bitsandbytes>=0.48.1",
     "safetensors>=0.5.3",
 ]

src/cache_dit/_utils/utils.py

@@ -425,6 +425,8 @@ def get_args(parse: bool = True, ) -> argparse.ArgumentParser | argparse.Namespa
       "float8_per_tensor",
       "float8_per_block",
       "float8_weight_only",
+      "nvfp4",
+      "nvfp4_weight_only",
       "int8_per_row",
       "int8_per_tensor",
       "int8_weight_only",
@@ -516,6 +518,18 @@ def get_args(parse: bool = True, ) -> argparse.ArgumentParser | argparse.Namespa
     default=False,
     help="Enable int4 weight-only quantization for transformer",
   )
+  parser.add_argument(
+    "--nvfp4",
+    action="store_true",
+    default=False,
+    help="Enable TorchAO NVFP4 dynamic activation and weight quantization for transformer",
+  )
+  parser.add_argument(
+    "--nvfp4-weight-only",
+    action="store_true",
+    default=False,
+    help="Enable TorchAO NVFP4 weight-only quantization for transformer",
+  )
   parser.add_argument(
     "--svdq-int4-r32-dq",
     "--svdq-r32",
@@ -1183,38 +1197,36 @@ def get_base_args(parse: bool = True) -> argparse.Namespace | argparse.ArgumentP
 
 def maybe_postprocess_args(args: argparse.Namespace) -> argparse.Namespace:
   # Force enable quantization if quantize_type is specified
-  if args.float8_per_row:
-    args.quantize_type = "float8_per_row"
-  elif args.float8_per_tensor:
-    args.quantize_type = "float8_per_tensor"
-  elif args.float8_per_block:
-    args.quantize_type = "float8_per_block"
-  elif args.float8_weight_only:
-    args.quantize_type = "float8_weight_only"
-  elif args.int8_per_row:
-    args.quantize_type = "int8_per_row"
-  elif args.int8_per_tensor:
-    args.quantize_type = "int8_per_tensor"
-  elif args.int8_weight_only:
-    args.quantize_type = "int8_weight_only"
-  elif args.int4_weight_only:
-    args.quantize_type = "int4_weight_only"
-  elif args.svdq_int4_r32_dq:
-    args.quantize_type = "svdq_int4_r32_dq"
-  elif args.svdq_int4_r64_dq:
-    args.quantize_type = "svdq_int4_r64_dq"
-  elif args.svdq_int4_r128_dq:
-    args.quantize_type = "svdq_int4_r128_dq"
-  elif args.svdq_int4_r256_dq:
-    args.quantize_type = "svdq_int4_r256_dq"
-  elif args.svdq_nvfp4_r32_dq:
-    args.quantize_type = "svdq_nvfp4_r32_dq"
-  elif args.svdq_nvfp4_r64_dq:
-    args.quantize_type = "svdq_nvfp4_r64_dq"
-  elif args.svdq_nvfp4_r128_dq:
-    args.quantize_type = "svdq_nvfp4_r128_dq"
-  elif args.svdq_nvfp4_r256_dq:
-    args.quantize_type = "svdq_nvfp4_r256_dq"
+  quantize_shortcuts = [
+    ("float8_per_row", args.float8_per_row),
+    ("float8_per_tensor", args.float8_per_tensor),
+    ("float8_per_block", args.float8_per_block),
+    ("float8_weight_only", args.float8_weight_only),
+    ("int8_per_row", args.int8_per_row),
+    ("int8_per_tensor", args.int8_per_tensor),
+    ("int8_weight_only", args.int8_weight_only),
+    ("int4_weight_only", args.int4_weight_only),
+    ("nvfp4", args.nvfp4),
+    ("nvfp4_weight_only", args.nvfp4_weight_only),
+    ("svdq_int4_r32_dq", args.svdq_int4_r32_dq),
+    ("svdq_int4_r64_dq", args.svdq_int4_r64_dq),
+    ("svdq_int4_r128_dq", args.svdq_int4_r128_dq),
+    ("svdq_int4_r256_dq", args.svdq_int4_r256_dq),
+    ("svdq_nvfp4_r32_dq", args.svdq_nvfp4_r32_dq),
+    ("svdq_nvfp4_r64_dq", args.svdq_nvfp4_r64_dq),
+    ("svdq_nvfp4_r128_dq", args.svdq_nvfp4_r128_dq),
+    ("svdq_nvfp4_r256_dq", args.svdq_nvfp4_r256_dq),
+  ]
+  selected_shortcuts = [quant_type for quant_type, enabled in quantize_shortcuts if enabled]
+  if len(selected_shortcuts) > 1:
+    raise ValueError(f"Quantization shortcuts are mutually exclusive, got {selected_shortcuts}.")
+  if selected_shortcuts:
+    selected_quant_type = selected_shortcuts[0]
+    if args.quantize_type is not None and args.quantize_type != selected_quant_type:
+      raise ValueError(
+        f"--quantize-type {args.quantize_type} conflicts with --{selected_quant_type.replace('_', '-')}."
+      )
+    args.quantize_type = selected_quant_type
 
   if args.quantize_type is not None:
     args.quantize = True

src/cache_dit/quantization/config.py

@@ -333,7 +333,7 @@ class QuantizeConfig:
   backend: str | QuantizeBackend = QuantizeBackend.AUTO
   # Quantization type, currently support "float8_weight_only" and "float8_per_row",
   # "float8_per_tensor", "float8_per_block", "int8_per_row", "int8_per_tensor",
-  # "int8_weight_only", "int4_weight_only", etc.
+  # "int8_weight_only", "int4_weight_only", "nvfp4", "nvfp4_weight_only", etc.
   quant_type: str = "float8_per_row"
   # The layers specified in this variable will be excluded from quantization,
   # even if they are in the repeated blocks or not filtered out by filter_fn.

src/cache_dit/quantization/torchao/quantize_ao.py

@@ -18,6 +18,8 @@
   "float8_per_tensor",
   "float8_per_block",
   "float8_weight_only",
+  "nvfp4",
+  "nvfp4_weight_only",
   "int8_per_row",
   "int8_per_tensor",
   "int8_weight_only",
@@ -29,6 +31,8 @@
   "float8_per_row",
   "float8_per_block",
   "float8_weight_only",
+  "nvfp4",
+  "nvfp4_weight_only",
   "int8_per_tensor",
   "int8_per_row",
   "int8_weight_only",
@@ -198,6 +202,12 @@ def __post_init__(self):
         9,
       ), "FP8 requires Ada or newer GPUs (>=sm89), but got " + str(
         current_platform.get_device_capability())
+    if self.is_nvfp4():
+      assert current_platform.get_device_capability() >= (
+        10,
+        0,
+      ), "NVFP4 requires Blackwell or newer GPUs (>=sm100), but got " + str(
+        current_platform.get_device_capability())
 
   @staticmethod
   def from_config(
@@ -411,6 +421,12 @@ def is_float8_per_block(self) -> bool:
   def is_float8_weight_only(self) -> bool:
     return self.quant_type == "float8_weight_only"
 
+  def is_nvfp4(self) -> bool:
+    return self.quant_type in ("nvfp4", "nvfp4_weight_only")
+
+  def is_nvfp4_weight_only(self) -> bool:
+    return self.quant_type == "nvfp4_weight_only"
+
   def required_fallback(self) -> bool:
     # Currently, only support float8 per-tensor fallback for rowwise layers if
     # regional quantiztion is enabled. Not support fallback for int8/int4/weight-only
@@ -485,6 +501,7 @@ def _get_torchao_config(quant_type: str, **kwargs) -> AOBaseConfig:
         Float8DynamicActivationFloat8WeightConfig,
         PerRow,
       )
+      from torchao.quantization.quantize_.common import KernelPreference
 
       quant_config = Float8DynamicActivationFloat8WeightConfig(
         weight_dtype=kwargs.get(
@@ -496,6 +513,7 @@ def _get_torchao_config(quant_type: str, **kwargs) -> AOBaseConfig:
           torch.float8_e4m3fn,
         ),
         granularity=(PerRow(), PerRow()),
+        kernel_preference=KernelPreference.TORCH,
       )
     elif quant_type == "float8_per_tensor":
       from torchao.quantization import (
@@ -550,6 +568,20 @@ def _get_torchao_config(quant_type: str, **kwargs) -> AOBaseConfig:
         torch.float8_e4m3fn,
       ), )
 
+    elif quant_type == "nvfp4":
+      from torchao.prototype.mx_formats.inference_workflow import (
+        NVFP4DynamicActivationNVFP4WeightConfig, )
+
+      quant_config = NVFP4DynamicActivationNVFP4WeightConfig(
+        use_triton_kernel=True,
+        use_dynamic_per_tensor_scale=True,
+      )
+
+    elif quant_type == "nvfp4_weight_only":
+      from torchao.prototype.mx_formats.inference_workflow import NVFP4WeightOnlyConfig
+
+      quant_config = NVFP4WeightOnlyConfig(use_dynamic_per_tensor_scale=True)
+
     elif quant_type == "int8_per_row":
       from torchao.quantization import (
         Int8DynamicActivationInt8WeightConfig,
@@ -658,6 +690,20 @@ def _is_curr_plan_allow_to_quantize(name: str) -> bool:  # precision plan
       logger.debug(skip_reason)
       return False
 
+    if curr_quant_type in ("nvfp4", "nvfp4_weight_only"):
+      weight_shape = tuple(m.weight.shape)
+      if weight_shape[-2] % 16 != 0 or weight_shape[-1] % 16 != 0:
+        skip_reason = _skip_reason(f"weight_shape{weight_shape}%16!=0")
+        quant_ctx.skipped_map[curr_quant_type].append(skip_reason)
+        logger.debug(skip_reason)
+        return False
+      if (curr_quant_type == "nvfp4"
+          and (weight_shape[-2] <= 64 or (weight_shape[-1] <= 1024 and weight_shape[-2] <= 1024))):
+        skip_reason = _skip_reason(f"weight_shape{weight_shape} is too small")
+        quant_ctx.skipped_map[curr_quant_type].append(skip_reason)
+        logger.debug(skip_reason)
+        return False
+
     # check blockwise fp8 support for linear layers, if not supported,
     # skip quantization for that layer.
     if curr_quant_type in [