Low-bits Quantization

Overview

Quantization is a powerful technique to reduce the memory footprint and computational cost of deep learning models by representing weights and activations with lower precision data types. Cache-DiT supports various quantization methods, including FP8, INT8, INT4, and NVFP4 quantization, to help users achieve faster inference and lower memory usage while maintaining acceptable model performance.

Quantization Types	Description	Devices
float8_per_row	quantize weights and activations to float8 (dynamic quantization) with rowwise method. (recommended)	>=sm89, Ada, Hopper or newer
float8_per_tensor	quantize weights and activations to float8 (dynamic quantization) with tensorwise method.	>=sm89, Ada, Hopper or newer
float8_per_block	block-wise quantization weights and activations (dynamic quantization) to float8, which can provide better precision, activations's blocksize: (1, 128), weight's blocksize: (128, 128)	>=sm89, Ada, Hopper or newer
float8_weight_only	quantize only weights to float8, keep activations in full precision	>=sm89, Ada, Hopper or newer
int8_per_row	quantize weights and activations to int8 (dynamic quantization) with rowwise method.	>=sm80, Ampere or newer
int8_per_tensor	quantize weights and activations to int8 (dynamic quantization) with tensorwise method.	>=sm80, Ampere or newer
int8_weight_only	quantize only weights to int8, keep activations in full precision	>=sm80, Ampere or newer
int4_weight_only	quantize only weights to int4, keep activations in full precision	>=sm90, Hopper or newer, TMA required
nvfp4	TorchAO dynamic activation and NVFP4 weight quantization. The current integration uses Triton activation quantization and dynamic per-tensor scaling.	>=sm100, Blackwell or newer
nvfp4_weight_only	TorchAO NVFP4 weight-only quantization with dynamic per-tensor weight scaling; activations remain in their original precision.	>=sm100, Blackwell or newer
svdq_int4_r{32...}	post-training SVDQuant (W4A4) with calibration and checkpoint serialization for users who want the higher-accuracy PTQ workflow.	>=sm80, Ampere or newer, excluded Hopper (NO INT4 MMA)
svdq_nvfp4_r{32...}	post-training SVDQuant (W4A4) with NVFP4 packed weights/activations, calibration, and checkpoint serialization. Only svdq_kwargs["runtime_kernel"]="v1" is currently supported.	>=sm120, Blackwell or newer
svdq_int4_r{32...}_dq	quantize weights and activations to int4 with SVDQuant dynamic quantization (W4A4) without any calibration.	>=sm80, Ampere or newer, excluded Hopper (NO INT4 MMA)
svdq_nvfp4_r{32...}_dq	quantize weights and activations to NVFP4 with SVDQuant dynamic quantization (W4A4) without any calibration. Only svdq_kwargs["runtime_kernel"]="v1" is currently supported.	>=sm120, Blackwell or newer

FP8 Quantization

Currently, Cache-DiT supports online quantization with different quantization types via TorchAo backend (torchao>=0.17.0 required). Users can implement model quantization by calling quantize or pass a QuantizeConfig to enable_cache API. Please make sure to install the latest version of torchao before using quantization features.

# stable: torchao (change cu130 to cu129 if using CUDA 12.9)
uv pip install torchao --index-url https://download.pytorch.org/whl/cu130 --upgrade
# nightly: torchao (change cu130 to cu129 if using CUDA 12.9)
uv pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu130 --upgrade

For GPUs with low memory capacity, we recommend using float8_per_row or float8_per_block, as these methods cause almost no loss in precision. We also recommend enabling torch.compile for better performance with quantization. Supported quantization types including:

• float8_per_row: quantize both weights and activations to float8 (dynamic quantization) with rowwise method.
• float8_per_tensor: quantize both weights and activations to float8 (dynamic quantization) with tensorwise method.
• float8_per_block: block-wise quantization weights and activations (dynamic quantization) to float8, which can provide better precision, activations's blocksize: (1, 128), weight's blocksize: (128, 128). NOT supported for distributed inference for now.
• float8_weight_only: quantize only weights to float8, keep activations in full precision.

Here are some examples of how to use quantization with cache-dit. You can directly specify the quantization config in the enable_cache API.

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig

# quant_type: float8_per_row, float8_per_tensor, float8_per_block, float8_weight_only, 
# int8_per_row, int8_per_tensor, int8_weight_only, int4_weight_only, etc.
# Pass a QuantizeConfig to the `enable_cache` API.
cache_dit.enable_cache( 
  pipe, cache_config=DBCacheConfig(), # w/ default
  parallelism_config=ParallelismConfig(ulysses_size=2),
  quantize_config=QuantizeConfig(quant_type="float8_per_row"),
)

Users can also specify different quantization configs for different components. For example, quantize the transformer to float8_per_row and the text encoder to float8_weight_only.

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig

cache_dit.enable_cache( 
  pipe, cache_config=DBCacheConfig(), # w/ default
  parallelism_config=ParallelismConfig(ulysses_size=2),
  quantize_config=QuantizeConfig(
    components_to_quantize={
      "transformer": {
        "quant_type": "float8_per_row",
        "exclude_layers": ["embedder", "embed"],
      },
      "text_encoder": {
        "quant_type": "float8_weight_only",
        "exclude_layers": ["lm_head"],
      }
    }
  ),
)

Or, directly call the quantize API for more fine-grained control.

import cache_dit
from cache_dit import QuantizeConfig

cache_dit.quantize(
  pipe.transformer, 
  quantize_config=QuantizeConfig(quant_type="float8_per_row"),
)
cache_dit.quantize(
  pipe.text_encoder, 
  quantize_config=QuantizeConfig(quant_type="float8_weight_only"),
)

Please also enable torch.compile for better performance with quantization.

import cache_dit

cache_dit.set_compile_configs()
pipe.transformer = torch.compile(pipe.transformer)
pipe.text_encoder = torch.compile(pipe.text_encoder)

Users can set exclude_layers in QuantizeConfig to exclude some sensitive layers that are not robust to quantization, e.g., embedding layers. Layers that contain any of the keywords in theexclude_layers list will be excluded from quantization. For example:

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig

cache_dit.enable_cache( 
  pipe, cache_config=DBCacheConfig(), # w/ default
  parallelism_config=ParallelismConfig(ulysses_size=2),
  quantize_config=QuantizeConfig(
    quant_type="float8_per_row",
    exclude_layers=["embedder", "embed"],
  ),
)

By default, quant_type="float8_per_row" for better precision. Users can set it to "float8_per_tensor" to use per-tensor quantization for better performance on some hardware.

Regional Quantization

Cache-DiT also supports regional quantization, which allows users to quantize only the repeated blocks in a transformer. This can be useful for better balancing the precision and efficiency. Users can specify the blocks to be quantized via the regional_quantize and repeated_blocks arguments in QuantizeConfig. For example, to quantize repeated blocks of the Flux2's transformer:

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig

cache_dit.enable_cache( 
  pipe, cache_config=DBCacheConfig(), # w/ default
  parallelism_config=ParallelismConfig(ulysses_size=2),
  quantize_config=QuantizeConfig(
    quant_type="float8_per_row",
    # Default (True), only quantize the repeated blocks in transformer if the repeated_blocks is 
    # specified. If set to False, the whole transformer will be quantized.
    regional_quantize=True, 
    # Specify the block names for the transformer, cache-dit will automatically find the repeated 
    # blocks and quantize it inplace. The block names can be found in the model architecture, e.g., 
    # for FLUX.2, the block name is "Flux2TransformerBlock" and "Flux2SingleTransformerBlock".
    repeated_blocks=['Flux2TransformerBlock', 'Flux2SingleTransformerBlock'],
    # repeated_blocks will be detected automatically from diffusers' transformer class, namely:
    # default repeated_blocks = transformer._repeated_blocks if exists, else None (quantize 
    # the whole transformer.
  ),
)

FP8 Per-Tensor Fallback

The per_tensor_fallback option in Cache-DiT's quantization configuration allows users to enable a fallback mechanism for layers that do not support float8 per-row or per-block quantization. This is particularly useful in scenarios where tensor parallelism is applied, and certain layers (e.g., those applied with RowwiseParallel) may encounter memory layout mismatch errors when quantized to float8 per-row.

When per_tensor_fallback is set to True, if a layer cannot be quantized to float8 per-row or per-block, it will automatically fall back to float8 per-tensor quantization instead of raising an error. This ensures that the quantization process can continue smoothly without interruption, while still providing the benefits of reduced precision for supported layers.

To enable this feature, simply set the per_tensor_fallback flag to True (default) in the QuantizeConfig when calling the enable_cache API. Only support for float8 quantization for now. For example:

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig

cache_dit.enable_cache( 
  pipe, cache_config=DBCacheConfig(), # w/ default
  parallelism_config=ParallelismConfig(tp_size=2),
  quantize_config=QuantizeConfig(
    quant_type="float8_per_row",
    # Must be True to enable fp8 per-tensor fallback.
    regional_quantize=True, # default, True.
    repeated_blocks=['Flux2TransformerBlock', 'Flux2SingleTransformerBlock'],
    # Enable fallback to float8 per-tensor quantization, default to True
    # for better compatibility for layers that do not support float8 per-row 
    # quantization, e.g., layers with RowwiseParallel applied in tensor parallelism.
    per_tensor_fallback=True, 
  ),
)

For examples, without fp8 per-tensor fallback, the cache-dit will auto skip the layers that do not support float8 per-row quantization, and raise warning for those layers. The performance will be worse due to less layers being quantized. (quantize 88 layers, skip 56 layers)

# w/o fp8 per-tensor fallback, quantize 88 layers, skip 56 layers, performance downgrade.
torchrun --nproc_per_node=2 -m cache_dit.generate flux2_klein_9b_kv_edit \
   --parallel tp --compile --float8-per-row --q-verbose \
   --disable-per-tensor-fallback

-----------------------------------------------------------------------------------
Quantized        Region: ['Flux2TransformerBlock', 'Flux2SingleTransformerBlock']  |
Quantized Linear Layers: 88    float8_per_row     56 (skipped)                     |
Quantized Linear Layers: 88    (total)                                             |
Skipped   Linear Layers: 56    (total)                                             |
Linear           Layers: 144   (total)                                             |
-----------------------------------------------------------------------------------
------------------------------------------------------------------------------------
float8_per_row, skip: attn.to_out.0        : pattern<RowwiseParallel>: 8    layers  |
float8_per_row, skip: attn.to_add_out      : pattern<RowwiseParallel>: 8    layers  |
float8_per_row, skip: ff.linear_out        : pattern<RowwiseParallel>: 8    layers  |
float8_per_row, skip: ff_context.linear_out: pattern<RowwiseParallel>: 8    layers  |
float8_per_row, skip: attn.to_out          : pattern<RowwiseParallel>: 24   layers  |
------------------------------------------------------------------------------------

With fp8 per-tensor fallback enabled, those layers that do not support float8 per-row quantization will be quantized to float8 per-tensor instead, and the performance will be better due to more layers being quantized. (quantize 144 layers, skip 0 layer)

# w/ fp8 per-tensor fallback enabled, quantize 144 layers, skip 0 layer, better performance.
torchrun --nproc_per_node=2 -m cache_dit.generate flux2_klein_9b_kv_edit \
   --parallel tp --compile --float8-per-row --q-verbose  

# Default, enabled fp8 per-tensor fallback
-----------------------------------------------------------------------------------
Quantized        Region: ['Flux2TransformerBlock', 'Flux2SingleTransformerBlock']  |
Quantized Linear Layers: 88    float8_per_row     0 (skipped)                      |
Quantized Linear Layers: 56    float8_per_tensor  0 (skipped)                      |
Quantized Linear Layers: 144   (total)                                             |
Skipped   Linear Layers: 0     (total)                                             |
Linear           Layers: 144   (total)                                             |
-----------------------------------------------------------------------------------

(Hybrid) Precision Plan

The precision_plan option in QuantizeConfig allows users to specify different quantization types for matched layer-name patterns. It is useful when you want better control of the accuracy and performance trade-off for attention sub-layers (for example, keep to_k/to_v in float8_per_row while using float8_per_tensor for to_q/to_out). Please note:

• Layers not matched by precision_plan continue to use the base quant_type.
• precision_plan is only valid when regional_quantize=True. If regional quantization is disabled, precision plan will be ignored.
• precision_plan is compatible with per_tensor_fallback. If a selected plan type is not supported by a specific layer/hardware path (case: rowwise tensor parallel is used and the basic quantize type is float8_per_row), fallback logic still works automatically when enabled.

For example: (FLUX.2-Klein-9b-kv)

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig

cache_dit.enable_cache(
  pipe,
  cache_config=DBCacheConfig(),
  quantize_config=QuantizeConfig(
     # Default type for unmatched layers in transformer.
    quant_type="float8_per_row",
    regional_quantize=True,
    repeated_blocks=['Flux2TransformerBlock', 'Flux2SingleTransformerBlock'],
    per_tensor_fallback=True,
    precision_plan={
      "attn.to_q": "float8_per_tensor",  # match: **attn.to_q**, best performance. 
      "attn.to_k": "float8_weight_only", # match: **attn.to_k**, best precision.
      "attn.to_v": "float8_per_block",   # match: **attn.to_v**, better precision.
      "attn.to_out": "float8_per_row",   # match: **attn.to_out**, better precision.
    },
  ),
)
# python3 -m cache_dit.generate flux2_klein_9b_kv_edit --config quantize_plan.yaml --compile

Then, the output summary will show the quantization type for each layer, and users can verify the quantization plan is applied correctly.

-----------------------------------------------------------------------------------
Quantized        Region: ['Flux2TransformerBlock', 'Flux2SingleTransformerBlock']  |
Quantized Linear Layers: 96    float8_per_row     0 (skipped)                      |
Quantized Linear Layers: 32    float8_per_tensor  0 (skipped)                      |
Quantized Linear Layers: 8     float8_per_block   0 (skipped)                      |
Quantized Linear Layers: 8     float8_weight_only 0 (skipped)                      |
Quantized Linear Layers: 144   (total)                                             |
Skipped   Linear Layers: 0     (total)                                             |
Linear           Layers: 144   (total)                                             |
-----------------------------------------------------------------------------------

INT8/INT4 Quantization

In addition to FP8 quantization, Cache-DiT also supports INT8 and INT4 quantization for weights, which can further reduce the memory footprint of the model. Users can specify int8_per_row, int8_per_tensor, int8_weight_only, or int4_weight_only as the quantization type in the QuantizeConfig when calling the enable_cache API. For example:

import cache_dit
from cache_dit import DBCacheConfig, ParallelismConfig, QuantizeConfig  

cache_dit.enable_cache( 
  # Or "int8_per_tensor", "int8_weight_only", "int4_weight_only", etc.
  pipe, quantize_config=QuantizeConfig(quant_type="int8_per_row"), 
)

INT4 quantization can provide even better memory reduction compared to FP8 or INT8, but it may cause more precision loss. We recommend users to try different quantization types and choose the one that best fits their needs in terms of the trade-off between performance and precision. In most cases, float8 per-row can be a good choice for better memory reduction while maintaining acceptable precision.

Please note that users should also install mslk kernel library to enable INT8/INT4 quantization features. The int4_weight_only w4a16 compute kennel requires architectures >= sm90 (Hopper or newer, TMA required). For older architectures, users can use int8_weight_only quantization for better compatibility.

# 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 distributed inference (context parallelism or tensor parallelism), we recommend users to use float8 quantization to avoid potential compatibility issues.

TorchAO NVFP4 Quantization

nvfp4 uses NVFP4DynamicActivationNVFP4WeightConfig from TorchAO's MX formats prototype. Cache-DiT currently fixes both use_triton_kernel=True and use_dynamic_per_tensor_scale=True. It is an online dynamic quantization path and is separate from the svdq_nvfp4_* SVDQuant path. (Required mslk kernels library)

For example:

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:

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 W4A4 GEMM kernels and an easy-to-use PTQ interface). The public API is intentionally small: build a QuantizeConfig, quantize with cache_dit.quantize(...), then reload with cache_dit.load(...). Cache-DiT now supports both INT4 and NVFP4 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.

With Cache-DiT's SVDQuant support, users can easily apply W4A4 quantization to their models with just a few lines of code, without worrying about the underlying implementation details or compatibility issues. Before using SVDQuant, please make sure to build Cache-DiT from source with SVDQuant support. We may consider releasing pre-built Cache-DiT packages with SVDQuant support in the future.

For PTQ, the available quant types are svdq_int4_r{rank} and svdq_nvfp4_r{rank}. The INT4 path remains the general cross-architecture option for supported non-Hopper GPUs, while the NVFP4 path is currently validated on Blackwell (sm120) and requires svdq_kwargs={"runtime_kernel": "v1"}. Setting runtime_kernel="v2" with NVFP4 is rejected intentionally at config-validation time.

# Required: CUDA 13.0+, PyTorch 2.11+, Ubuntu 22.04+ (GLIBC 2.32+).
uv pip install -U cache-dit-cu13 # PyPI, stable release with SVDQ.
# Optional: just build Cache-DiT with SVDQuant support from source.
git clone https://github.com/vipshop/cache-dit.git && cd cache-dit
git submodule update --init --recursive --force # init submodules 
CACHE_DIT_BUILD_SVDQUANT=1 MAX_JOBS=32 uv pip install -e ".[quantization]"

The only thing users need to do is to prepare the calibration function (calibrate_fn) for collecting quantization statistics, and Cache-DiT will take care of the rest, including computing the quantization parameters, applying quantization, and loading the quantized weights for inference or further fine-tuning. For example:

Step 1: Load the pre-trained model in full precision and DON'T forget to move the model to "CUDA" for better calibration performance.

import torch
import cache_dit
from diffusers import Flux2KleinPipeline
from cache_dit import QuantizeConfig

# 1. Load the pre-trained model in "cuda"
pipe = Flux2KleinPipeline.from_pretrained(
  "black-forest-labs/FLUX.2-klein-4B",
  torch_dtype=torch.bfloat16,
).to("cuda")

Step 2: Prepare the calibration data and calibration function (calibrate_fn) for SVDQuant QuantizeConfig. Cache-DiT will automatically collect the quantization statistics during the calibration process, and use those statistics to compute the quantization parameters for each linear layer.

# 2.1 Prepare the calibration dataset for collecting quantization statistics.
calibration_prompts = ["A cute cat sitting on the beach.", ...]

# 2.2 Define the calibration function for SVDQuant PTQ.
def calibrate_fn(**_: object) -> None:
  with torch.inference_mode():
    for prompt in calibration_prompts:
      _ = pipe(
        prompt=prompt,
        height=1024,
        width=1024,
        num_inference_steps=4,
        generator=torch.Generator(device="cpu").manual_seed(0),
      )

Step 3: Build the QuantizeConfig for SVDQuant, and call cache_dit.quantize(...) to apply SVDQuant W4A4 quantization to the model. We have to wait some time for the quantization process to finish, as SVDQuant will perform SVD decomposition for each linear layer and compute the quantization parameters based on the collected quantization statistics. The quantized model will be saved to disk in the specified directory.

# 3. Build the QuantizeConfig for SVDQuant, and call `cache_dit.quantize(...)`.
quant_config = QuantizeConfig(
  quant_type="svdq_int4_r32", # _r{rank}, e.g., r16, r32, r64, r128, etc.
  calibrate_fn=calibrate_fn,
  serialize_to="./FLUX.2-klein-4B-svdq/",
  svdq_kwargs={"calibrate_precision": "low"},  # default, optional: low | medium | high
)
pipe.transformer = cache_dit.quantize(pipe.transformer, quant_config)

The NVFP4 PTQ flow uses the same public API, but should select an NVFP4 quant type and explicitly keep the runtime kernel on v1:

quant_config = QuantizeConfig(
  quant_type="svdq_nvfp4_r32",  # _r{rank}, e.g., r16, r32, r64, r128, etc.
  calibrate_fn=calibrate_fn,
  serialize_to="./FLUX.2-klein-4B-svdq-nvfp4/",
  svdq_kwargs={
    "calibrate_precision": "low",
    "runtime_kernel": "v1",
  },
)
pipe.transformer = cache_dit.quantize(pipe.transformer, quant_config)

calibrate_precision now defaults to low, which uses randomized torch.svd_lowrank and keeps the calibration math in float32. medium and high remain available as opt-in debugging or accuracy-investigation modes. On a real FLUX.2-klein-4B PTQ run (rank=32, 8 calibration prompts, 1024x1024), the quantized-model behavior was almost the same across all three modes; the practical difference was mostly PTQ cost. The speedup column below refers to quantization-time speedup relative to high.

calibrate_precision	low-rank route	quantization_s	speedup	PSNR
low	torch.svd_lowrank w/ float32	20.2258	18.02x	29.2528
medium	full torch.linalg.svd w/ float32	111.8944	3.26x	29.2454
high	float64 full SVD, CUDA gesvd	364.4877	1.00x	29.1880

In practice the PSNR values are almost the same, so low is the recommended production default. It keeps the best PTQ throughput while preserving the same loaded-model behavior seen from medium and high.

Step 4: The FLUX.2-klein-4B-svdq directory now contains: {quant_type}.safetensors (for example, svdq_int4_r32.safetensors or svdq_nvfp4_r32.safetensors) and quant_config.json. The quantized model can be loaded with cache_dit.load(...) for inference or further fine-tuning. For example:

# 4.1 First, keep in CPU for now.
pipe = Flux2KleinPipeline.from_pretrained(...) 
# 4.2 Load the quantized model for inference.
pipe.transformer = cache_dit.load(pipe.transformer, "./FLUX.2-klein-4B-svdq/")
pipe.to("cuda") # Transfer the model to GPU for inference.

For SVDQuant, serialize_to should point to a directory. Cache-DiT normalizes the checkpoint name to {quant_type}.safetensors, writes a sibling quant_config.json, and lets cache_dit.load(...) recover the real checkpoint path directly from that directory. Performance comparison between SVDQuant W4A4 and full precision BF16 for FLUX.2-klein-4B is listed below:

The full benchmark report for FLUX.2-klein-4B SVDQuant (W4A4) is listed below. The SVDQuant W4A4 quantized model can achieve around 2.5x memory reduction for transformer weights, around 1.7x speedup in latency, while maintaining good output quality (the output images are almost the same as the full precision model).

stage	resolution	avg_latency_s	peak_memory_gb	transformer_gb
float_reference	1024x1024	2.1302	17.3245	7.2188
memory_quantized	1024x1024	1.2520	12.3916	2.2761
loaded_quantized	1024x1024	1.2410	12.3883	2.2761

Notes: full_blocks means all transformer blocks are quantized; memory_quantized means the quantized model is in memory; loaded_quantized means the quantized model is reloaded from disk with cache_dit.load(...).

SVDQuant in Cache-DiT is compatible with torch.compile, and users can further speed up the quantized model with compilation. The performance comparison between eager and compiled mode for SVDQuant W4A4 is listed below:

stage	resolution	avg_latency_s	latency_delta	speedup_vs_eager
quantized + eager	1024x1024	1.2689	0.0000	1.0000x
quantized + compile	1024x1024	1.0161	-0.2528	1.2488x

Cache-DiT also has an NVFP4 PTQ path for Blackwell. On a RTX 5090 (sm120) validation run with FLUX.2-klein-4B, svdq_nvfp4_r32, runtime_kernel="v1", 512x512, 4 denoising steps, benchmark_runs=3, and warmup_runs=1, the full e2e example completed end-to-end for in-memory inference, checkpoint reload, and compiled quantized execution:

stage	resolution	latency_s	speedup	peak_memory_gb	psnr
baseline	512x512	0.3043	1.0000x	15.5402	-
memory_quantized	512x512	0.1725	1.7641x	10.7098	17.0449
loaded_quantized	512x512	0.1764	1.7251x	10.7127	15.5375
compiled_quantized	512x512	0.1430	2.1280x	10.7127	16.1652

In that validation, the loaded transformer reported precision == "nvfp4", confirming that the serialized PTQ checkpoint preserved the NVFP4 precision metadata rather than silently falling back to INT4.

Under the same svdq_nvfp4_r32 + runtime_kernel="v1" setup, rerunning the full e2e example at 1024x1024 on the same RTX 5090 (sm120) kept the same stage ordering while shifting the absolute runtime and memory envelope upward:

stage	resolution	latency_s	speedup	peak_memory_gb	psnr
baseline	1024x1024	0.9703	1.0000x	17.3245	-
memory_quantized	1024x1024	0.5754	1.6862x	12.4963	20.1529
loaded_quantized	1024x1024	0.5851	1.6584x	12.4969	19.9014
compiled_quantized	1024x1024	0.4742	2.0458x	12.4969	21.2930

For this 1024x1024 run, the transformer CUDA footprint again dropped from 7.2188 GiB in the baseline to 2.3850 GiB in all three quantized stages, while the compiled stage remained the fastest steady-state path.

How to use SVDQ in Nunchaku with Cache-DiT?

FAQ: Can I use Cache-DiT's SVDQuant support for Nunchaku int4 models?

Yes, Cache-DiT also supports the caching and context parallelism scheme for int4 models from official Nunchaku library. Users can leverage Cache-DiT to speed up Nunchaku 4-bits W4A4 models. But we recommend users to use Cache-DiT's native SVDQuant support for better compatibility and usability.

import cache_dit
from diffusers import QwenImagePipeline
from nunchaku import NunchakuQwenImageTransformer2DModel

transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(
  f"path-to/svdq-int4_r32-qwen-image.safetensors"
)
pipe = QwenImagePipeline.from_pretrained(
  "Qwen/Qwen-Image", transformer=transformer, torch_dtype=torch.bfloat16,
).to("cuda")

cache_dit.enable_cache(pipe, cache_config=..., parallelism_config=...)

SVDQuant (W4A4) DQ

Cache-DiT also supports SVDQ dynamic quantization through quant types ending with _dq, such as svdq_int4_r32_dq, svdq_int4_r128_dq, svdq_nvfp4_r32_dq, and svdq_nvfp4_r128_dq. By default, this workflow keeps the low-rank SVD branch and, by default, fixes the smooth factor to 1, so users do not need to provide a calibration callback or a serialization directory.

For NVFP4 DQ, the current runtime restriction is the same as PTQ: keep svdq_kwargs["runtime_kernel"]="v1". INT4 DQ can still use the broader INT4 runtime options where supported, but NVFP4 runtime_kernel="v2" is intentionally unsupported for now.

Under the hood, cache-dit now treats this as an explicit backend-local smooth-strategy choice: DQ uses the identity smooth vector by default, while PTQ continues to use activation-derived smoothing. The public DQ behavior does not change, but this separation makes it easier to add future opt-in smooth strategies without redefining the meaning of existing _dq quant types.

The default DQ usage keeps the identity smooth vector and does not require any extra svdq_kwargs:

import cache_dit
from cache_dit import QuantizeConfig

# 0. use 'enable_cache' API with the default identity-smooth DQ.
cache_dit.enable_cache(
  pipe,
  quantize_config=QuantizeConfig(
    quant_type="svdq_int4_r32_dq",
  ),
)

# 1. or directly call the `quantize` API for more fine-grained control.
pipe.transformer = cache_dit.quantize(
  pipe.transformer,
  QuantizeConfig(
    quant_type="svdq_int4_r32_dq",
  ),
)

The NVFP4 DQ path uses the same flow, but should choose an NVFP4 _dq type and keep the runtime kernel on v1:

pipe.transformer = cache_dit.quantize(
  pipe.transformer,
  QuantizeConfig(
    quant_type="svdq_nvfp4_r32_dq",
    svdq_kwargs={"runtime_kernel": "v1"},
  ),
)

The avaliable DQ smooth strategies are not limited to the identity smooth vector. We have also implemented some experimental heuristics that derive non-trivial smooth factors from weight statistics alone, without any activation calibration. These heuristics are inspired by the mathematical form of the PTQ smooth factor, but they intentionally ignore the activation term and only use weight statistics to compute the smooth factors. For example, users can set the smooth_strategy in svdq_kwargs to weight, weight_inv, or few_shot to enable these non-default DQ modes:

import cache_dit
from cache_dit import QuantizeConfig

cache_dit.enable_cache(
  pipe,
  quantize_config=QuantizeConfig(
    quant_type="svdq_int4_r32_dq",
    svdq_kwargs={"smooth_strategy": "weight"},
  ),
)

These experimental modes still avoid activation calibration, but derive non-trivial per-channel smooth factors from weight statistics alone. They are currently available as opt-in Python API settings rather than dedicated CLI shortcuts. For the shared generate helpers, Cache-DiT also exposes --svdq-{int4|nvfp4}-r{32|64|128|256}-dq, as quick INT or NVFP4 DQ entry points. The mathematical difference between the DQ smooth strategies is listed below.

identity: use the unit smooth vector directly.

\[ s_i = 1 \]

This keeps the current zero-calibration DQ behavior unchanged. The low-rank decomposition is applied directly to the original weight matrix, so the entire DQ approximation is driven only by the weight tensor and the target rank.

weight: replace the activation term with a constant proxy and build a purely weight-derived smooth factor. In the current implementation, PTQ uses the following smooth scale:

\[ s_i = \frac{a_i^\alpha}{w_i^{1-\alpha}} \]

The weight-only DQ heuristic starts from:

\[ \hat{s}_i = w_i^{-(1-\alpha)} \]

Here w_i denotes the per-input-channel weight span. Cache-DiT then stabilizes this raw heuristic by geometric-mean normalization and explicit clipping:

\[ s_i = \operatorname{clip}\!\left(\frac{\hat{s}_i}{\exp\left(\frac{1}{n}\sum_j \log \hat{s}_j\right)},\ 0.25,\ 4.0\right) \]

Intuitively, identity does not redistribute any quantization difficulty across channels, while weight tries to equalize channels using weight statistics alone. Because it has no access to real activation outliers, it should be treated as an experimental heuristic rather than a replacement for PTQ smoothing.

weight_inv: use the same weight-only proxy, but flip the exponent direction so that channels with larger weight span receive larger smooth factors.

\[ \hat{s}_i^{\mathrm{inv}} = w_i^{1-\alpha} \]

Cache-DiT then applies the same geometric-mean normalization and clipping.

Intuitively, weight_inv does the opposite trade-off of weight: it intentionally makes the weight tensor harder to quantize on large-span channels, but it may reduce activation quantization difficulty because those activation channels are divided by a larger smooth factor at runtime. For examples:

# default: identity smooth strategy
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy identity --svdq-calibrate-precision low

# experimental: weight-only heuristic
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy weight --svdq-calibrate-precision low

# experimental: bias the trade-off toward easier activation quantization
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy weight_inv --svdq-calibrate-precision low

# high-rank are recommended for better precision while using SVDQ DQ, e.g, 256.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r256_dq \
  --svdq-smooth-strategy identity --svdq-calibrate-precision low

# higher SVD decomposition precision may be helpful for better DQ accuracy.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r256_dq \
  --svdq-smooth-strategy identity --svdq-calibrate-precision medium

Here are some preliminary results on the impact of identity, weight-only, and weight_inv heuristics DQ smooth strategies for a SVDQ DQ sweep on FLUX.2-klein-4B.

baseline (bf16)	identity (default)	weight-only heuristic	weight_inv heuristic
2.13s	1.28s	1.28s	1.28s
-	PSNR: 28.71	PSNR: 29.62	PSNR: 29.35

few_shot: delay SVDQ DQ materialization until the runtime has observed a small number of real inference forwards. Cache-DiT will keep the transformer in float mode during the first few_shot_steps root-module forwards, stream per-layer activation spans by layer name, and after the last observed forward finishes it will relax those activation spans according to a configurable few-shot policy, recompute activation-derived smooth scales, quantize the eligible linear layers in place, and use the quantized transformer from the next forward onward.

This mode exposes four backend-local knobs:

• few_shot_steps: number of root-module forwards to observe before quantization. This means transformer/module forward passes, not top-level pipeline invocations, and the count is cumulative on the same armed module. Default: 1. That cumulative behavior matters in diffusion pipelines because one pipeline run usually executes the transformer many times. For example, if the same pipe instance runs a 50-step sample once and then runs another 50-step sample again, the few-shot controller sees roughly 100 root-module forwards in total rather than resetting back to 50 for the second run.
• few_shot_relax_factor: maximum multiplicative factor applied to activation spans during few-shot relaxation. It must satisfy few_shot_relax_factor >= 1.0, so the observed activation span is preserved or amplified, never reduced. Default: 1.5. In practice, values around 1.5-2.5 are usually the safer starting range; larger values such as 4.0 can still oversmooth or blur image outputs when the outlier prior becomes too strong.
• few_shot_relax_top_ratio: fraction of the largest activation spans used to define the few-shot threshold. Default: 0.25.
• few_shot_relax_strategy: how the relax factor is assigned across channels before smooth-scale recomputation. Default: auto. Available values: fixed, top, auto, stable_auto, power, log, and rank.

If the observed activation-span vector is denoted as a and the weight-span vector is denoted as w, Cache-DiT first computes the threshold corresponding to 1 - few_shot_relax_top_ratio on a, relaxes the activation span, and then recomputes the smooth scale:

\[ s_i^{\mathrm{relaxed}} = \frac{\left(a_i^{\mathrm{relaxed}}\right)^\alpha}{w_i^{1-\alpha}} \]

It then applies one of the following relax strategies:

For auto, Cache-DiT also relaxes the channels below the threshold, but uses a smaller multiplier for smaller activation spans and caps the largest channels at few_shot_relax_factor.

\[ m_i = 1 + \left(\text{relax} - 1\right) \cdot \operatorname{clip}\left(\frac{a_i - a_{\min}}{\tau - a_{\min}}, 0, 1\right) \]

\[ a_i^{\mathrm{relaxed}} = a_i \cdot m_i \]

For stable_auto, Cache-DiT keeps the same magnitude-aware linear ramp as auto, but snaps that normalized response to a small number of evenly spaced buckets before applying the final relax multiplier. In the current implementation, the bucketized response is:

\[ r_i^{\mathrm{stable}} = \operatorname{round}(8 \cdot r_i^{\mathrm{auto}}) / 8 \]

This does not make the entire runtime path fully deterministic, but it does make the relax policy less sensitive to small first-forward activation-stat fluctuations.

For top:

\[ a_i^{\mathrm{relaxed}} = \begin{cases} a_i \cdot \text{relax}, & a_i \geq \tau \\ a_i, & a_i < \tau \end{cases} \]

Channels outside the selected top-ratio are kept unchanged.

For fixed:

\[ a_i^{\mathrm{relaxed}} = a_i \]

This disables the relax step entirely and uses the original activation-derived few-shot smooth scale after recomputation from the unchanged activation span. It is useful as a control/baseline when you want few-shot activation collection without any extra channel-wise amplification. At runtime, this strategy ignores few_shot_relax_factor and few_shot_relax_top_ratio.

Here tau is the quantile threshold and a_min is the minimum activation span in the layer. This makes the relax factor increase monotonically with the observed few-shot activation span while keeping it bounded by the configured maximum.

Internally, Cache-DiT first constructs a strategy-specific normalized response in the unit interval [0, 1], then maps it to the activation-span multiplier:

\[ m_i = 1 + r_i \cdot \left(\text{relax} - 1\right) \]

We intentionally clamp the intermediate response to [0, 1] before this affine map so the final activation-span multiplier is always bounded inside [1, relax]. Because relax is constrained to be at least 1, every channel is either unchanged or amplified; no channel can be shrunk by the few-shot relax step. Since smooth scale is recomputed from the activation span raised to alpha, the induced smooth-scale multiplier is bounded more conservatively inside [1, relax^alpha].

That does not mean arbitrarily large factors are safe. A larger factor tells the quantizer to treat the already-large activation-span channels as if they were even more likely to encounter future activation outliers. When that prior becomes too strong, the recomputed smoothing step over-allocates protection to those channels, which often shows up as visibly softer or blurrier image details. If you need a stronger head-channel bias, it is usually better to keep few_shot_relax_factor moderate and change few_shot_relax_strategy to power or reduce few_shot_relax_top_ratio, rather than pushing the factor directly to 4.0.

baseline (bf16)	DQ only	few-shot 1 steps + auto	few-shot 4 steps + auto

The remaining non-trivial strategies reuse the same bounded activation-span multiplier range [1, few_shot_relax_factor], but change how the normalized response is built before that final affine map:

• stable_auto: uses the same linear response as auto, but bucketizes it into a small number of levels so repeated few-shot runs are less sensitive to tiny activation-stat changes.
• power: uses a convex squared response.

\[ r_i^2 \]

This makes the largest activation spans receive much stronger amplification while smaller channels stay closer to 1. - log: uses a concave log response, so mid/high activation-span channels get lifted earlier than in the linear auto strategy. - rank: ignores the absolute activation-span gaps and constructs the response from channel rank percentiles, which can be more robust when a layer has a few extreme outliers.

import cache_dit
from cache_dit import QuantizeConfig

cache_dit.enable_cache(
  pipe,
  quantize_config=QuantizeConfig(
    quant_type="svdq_int4_r32_dq",
    svdq_kwargs={
      "smooth_strategy": "few_shot",
      "few_shot_steps": 1,
      "few_shot_relax_factor": 1.5,
      "few_shot_relax_top_ratio": 0.25,
      "few_shot_relax_strategy": "auto",
      # Optional: request compile after runtime quantization finishes.
      "few_shot_auto_compile": True,
    },
  ),
)

When few_shot_auto_compile is enabled, helper flows may defer transformer compilation until the runtime quantization step finishes. This means the warmup forwards stay in float mode, while the next forward after quantization runs on the quantized + compiled transformer.

If no owner-level deferred compile callback is registered, Cache-DiT falls back to compiling the quantized root module in-place with nn.Module.compile(). When the root module is a diffusers transformer that exposes compile_repeated_blocks(), that regional compile path is still preferred.

The same workflow is available from the generate CLI through --quantize-type or shortcut flags:

# few-shot activation-derived DQ: sample one runtime forward, then quantize.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-relax-factor 1.5 --svdq-few-shot-relax-top-ratio 0.25 \
  --svdq-few-shot-relax-strategy auto

# few-shot fixed: keep the observed activation span unchanged, then recompute smooth scale.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-relax-strategy fixed

# few-shot auto relax: larger activation spans receive larger relax factors, capped at 1.5.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-relax-factor 1.5 --svdq-few-shot-relax-top-ratio 0.25 \
  --svdq-few-shot-relax-strategy auto

# few-shot stable_auto relax: keep the same auto trend, but bucketize it for better stability.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-relax-factor 1.5 --svdq-few-shot-relax-top-ratio 0.25 \
  --svdq-few-shot-relax-strategy stable_auto

# few-shot power relax: focus the extra amplification on the largest channels.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-relax-factor 1.5 --svdq-few-shot-relax-top-ratio 0.25 \
  --svdq-few-shot-relax-strategy power

# few-shot rank relax: assign relax factors from channel rank percentiles.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-relax-factor 1.5 --svdq-few-shot-relax-top-ratio 0.25 \
  --svdq-few-shot-relax-strategy rank

# few-shot + deferred compile after runtime quantization completes.
python -m cache_dit.generate flux2_klein_4b --quantize-type svdq_int4_r128_dq \
  --svdq-smooth-strategy few_shot --svdq-few-shot-steps 1 \
  --svdq-few-shot-compile

# Tip: increase `--svdq-few-shot-steps` from 1 to a larger number if you want to collect more
# transformer/module forwards before few-shot materialization. The counter is cumulative on the
# same armed module, so two runs of the same 50-step pipeline already contribute about 100
# forwards in total. Larger values can improve few-shot DQ accuracy, but they also increase the
# runtime quantization latency because more warmup forwards are needed before quantization happens.

Unlike the PTQ workflow, SVDQ dynamic quantization is in-memory only. It does not require calibrate_fn or serialize_to, and it does not support load() from a serialized checkpoint. The few-shot mode also keeps this in-memory-only contract: it delays materialization to runtime, but still does not emit a serialized checkpoint.

At the backend-config level, this means DQ currently supports identity by default, experimental weight / weight_inv heuristics, and the runtime-activated few_shot mode as explicit opt-ins. In the CLI, this is controlled by --svdq-smooth-strategy, whose default value is identity. PTQ keeps the activation-derived strategy and serialized checkpoint workflow. We will support more smooth strategies in the future, and this separation makes it easier to add those as opt-in extensions without redefining the meaning of existing _dq quant types.

The CLI also exposes --svdq-calibrate-precision (alias --svdq-calib) for the SVDQ decomposition math. For DQ, the default remains low to preserve the shipped fast path, but users can now explicitly select medium or high when they want a different accuracy / quantization-time trade-off. The few-shot CLI also exposes --svdq-few-shot-steps, --svdq-few-shot-relax-factor, --svdq-few-shot-relax-top-ratio, --svdq-few-shot-relax-strategy, and --svdq-few-shot-compile. In the shared helper flow, top-level --compile is treated as a compatibility alias for few-shot runs, while --svdq-few-shot-compile is the precise switch that means "compile only after runtime quantization completes". The parser also accepts top_q4 as an alias of top, but fixed, top, auto, stable_auto, power, log, and rank are the canonical strategy names.

SVDQ Converter CLI

Cache-DiT ships a standalone converter CLI that loads a pretrained diffusion pipeline and converts it to an SVDQ W4A4 quantized format in a single step. The converter currently supports SVDQ dynamic quantization (_dq) only — it does not require a calibration dataset or a calibration callback.

The converter is available via the cache-dit-convert console script or as a runnable module:

# Console script (installed with cache-dit):
cache-dit-convert --model-path /path/to/FLUX.1-dev \
  --save-dir ./FLUX.1-dev-svdq \
  --quant-type svdq-int4-r128-dq

# Equivalent Python module invocation:
python3 -m cache_dit.quantization.converter \
  --model-path /path/to/FLUX.1-dev \
  --save-dir ./FLUX.1-dev-svdq \
  --quant-type svdq-int4-r128-dq

CLI reference

Argument	Required	Default	Description
--model-path	yes	—	Path to the pretrained diffusion model directory (e.g. /path/to/FLUX.1-dev).
--save-dir	yes	—	Directory where the quantized checkpoint ({quant_type}.safetensors) and quant_config.json are saved.
--quant-type	yes	—	SVDQ quant type, accepting both hyphen (svdq-int4-r128-dq) and underscore (svdq_int4_r128_dq) forms. Currently only _dq types are supported.
--torch-dtype	no	bfloat16	Torch dtype used when loading the float pipeline. One of: float16, bfloat16, float32.
--device	no	cuda	Device to run quantization on.
--svdq-smooth-strategy	no	identity	SVDQ DQ smooth strategy. One of: identity, weight, weight_inv, few_shot. See the SVDQuant DQ section for details.
--svdq-calibrate-precision	no	low	Precision plan for SVDQ decomposition math. One of: low, medium, high.
--svdq-runtime-kernel	no	v1	Packed runtime GEMM kernel used by SVDQW4A4Linear. One of: v1, v2.
--verbose	no	false	Print detailed quantization information including per-layer skip reasons.

Usage examples

Basic conversion — identity smooth (zero calibration):

cache-dit-convert \
  --model-path black-forest-labs/FLUX.2-klein-4B \
  --save-dir ./FLUX.2-klein-4B-svdq \
  --quant-type svdq-int4-r128-dq

Output:

Loading pipeline from black-forest-labs/FLUX.2-klein-4B ...
Pipeline loaded: Flux2KleinPipeline
Transformer memory before quantize: 7.22 GiB
Quantizing transformer (svdq_int4_r128_dq) and saving to ./FLUX.2-klein-4B-svdq ...
Transformer memory after quantize: 2.28 GiB (3.2x reduction)
Quantized checkpoint saved to ./FLUX.2-klein-4B-svdq/svdq_int4_r128_dq.safetensors
Quant config saved to ./FLUX.2-klein-4B-svdq/quant_config.json
Conversion complete. Load the quantized model with:
  cache_dit.load(transformer, './FLUX.2-klein-4B-svdq/')

Experimental weight-only smooth strategy:

# 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

# SVDQ NVFP4
cache-dit-convert \
  --model-path black-forest-labs/FLUX.2-klein-4B \
  --save-dir ./FLUX.2-klein-4B-svdq-nvfp4 \
  --quant-type svdq-nvfp4-r256-dq \
  --svdq-smooth-strategy weight \
  --svdq-calibrate-precision medium 

With v2 runtime kernel and verbose logging:

# 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 \
  --quant-type svdq-int4-r64-dq \
  --svdq-runtime-kernel v2 \
  --verbose

Loading the converted model

After conversion, the save-dir directory contains {quant_type}.safetensors and quant_config.json. Load the quantized model for inference with cache_dit.load:

import torch
import cache_dit
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained(...)
pipe.transformer = cache_dit.load(pipe.transformer, "./FLUX.2-klein-4B-svdq/")
pipe.to("cuda")
image = pipe("A cat sitting on the beach.", height=1024, width=1024).images[0]

You can also pass the safetensors path directly:

pipe.transformer = cache_dit.load(
    pipe.transformer,
    "./FLUX.2-klein-4B-svdq/svdq_int4_r128_dq.safetensors"
)

SVDQ Runtime Kernel

Cache-DiT's SVDQuant support also includes a high-performance W4A4 GEMM kernel for efficient runtime execution of SVDQ-quantized models. This kernel is optimized for Ada and Ampere architectures and can provide slightly better performance compared to existing INT4 GEMM kernels.

Users can enable the SVDQuant runtime by setting the runtime_kernel argument in QuantizeConfig to v2. For example:

quant_config = QuantizeConfig(
  quant_type="svdq_int4_r128_dq", # _r{rank}, e.g., r16, r32, r64, r128, etc.
  svdq_kwargs={"runtime_kernel": "v2"},  # enable SVDQ W4A4 v2 GEMM kernel for better performance.
)

Quick examples for enabling SVDQ runtime kernel from the generate CLI:

# v1 baseline: 11.81s -> v2 runtime: 11.58s 
python3 -m cache_dit.generate flux --svdq-int4-r64-dq --compile # v1, baseline
python3 -m cache_dit.generate flux --svdq-int4-r64-dq --svdq-runtime v2 --compile # v2

SVDQ with Fused MLP

When SVDQ quantizes a diffusers transformer whose FeedForward blocks use plain GELU activation, the default execution path runs three GPU kernels per MLP block: the first quantized GEMM, a separate GELU activation kernel, and the second quantized GEMM. Enabling fused MLP combines the first GEMM and GELU into a single kernel via svdq_gemm_w4a4_ext, eliminating one kernel launch per block.

The feature is controlled by svdq_kwargs["fused_mlp"] and works automatically with most diffusers transformer families — no per-model configuration is needed.

Quick start (CLI)

# Add --svdq-fused-mlp to any SVDQ generate command:
python -m cache_dit.generate longcat_image_edit --svdq-int4-r128-dq --svdq-fused-mlp

[06-09 02:56:25] [Cache-DiT] No caching or parallelism is applied.
[06-09 02:56:25] [Cache-DiT] Quantizing transformer module: LongCatImageTransformer2DModel to svdq_int4_r128_dq ...
[06-09 02:57:24] [Cache-DiT] ✅ Applying pass fused_gelu_mlp to 20 target(s).
[06-09 02:57:24] [Cache-DiT] ✅ Applying pass fused_gelu_proj to 20 target(s).
[06-09 02:57:24] [Cache-DiT] --------------------------------------------------------------------------------------------
[06-09 02:57:24] [Cache-DiT] SVDQuant    Region: ['LongCatImageTransformerBlock', 'LongCatImageSingleTransformerBlock']  |
[06-09 02:57:24] [Cache-DiT] SVDQuant      Type: svdq_int4_r128_dq                                                       |
[06-09 02:57:24] [Cache-DiT] SVDQuant      Rank: 128                                                                     |
[06-09 02:57:24] [Cache-DiT] Quantized   Layers: 260 / 260                                                               |
[06-09 02:57:24] [Cache-DiT] Skipped     Layers: 0                                                                       |
[06-09 02:57:24] [Cache-DiT] Linear      Layers: 260                                                                     |
[06-09 02:57:24] [Cache-DiT] --------------------------------------------------------------------------------------------

Baseline	SVDQ QD Rank=128 + Few-Shot	SVDQ QD Rank=128 + Few-Shot + Fused MLP
89s	58s	56s

Quick start (Python API — dynamic quantization)

import torch
import cache_dit
from diffusers import FluxPipeline
from cache_dit.quantization import QuantizeConfig

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
).to("cuda")

quant_config = QuantizeConfig(
    quant_type="svdq_int4_r32_dq",
    svdq_kwargs={"fused_mlp": True},
)
pipe.transformer = cache_dit.load(pipe.transformer, quant_config)

image = pipe("A cat holding a sign that says hello world").images[0]
image.save("flux_fused_mlp.png")

Quick start (Python API — PTQ with serialized checkpoint)

quant_config = QuantizeConfig(
    quant_type="svdq_int4_r32",
    serialize_to="./flux-svdq/",
    svdq_kwargs={"fused_mlp": True},
    calibrate_fn=my_calibrate_fn,
)
cache_dit.quantize(pipe.transformer, quant_config)

# Later, at inference time:
pipe.transformer = cache_dit.load(
    pipe.transformer,
    "./flux-svdq/svdq_int4_r32.safetensors",
)

When fused_mlp is enabled, cache-dit applies two complementary passes:

Pass	Targets	Fusion
fused_gelu_mlp	Standard FeedForward double blocks	fc1 + GELU + fc2 (qout path, no fp16 HBM write)
fused_gelu_proj	Single-stream blocks with concat MLP	fc1 + GELU only (fp16 output, concat unchanged)

Both passes use generic structural detection — they work with most diffusers transformers (FLUX, SD3, PixArt, HunyuanVideo, Wan, Cosmos, Bria, QwenImage, Chroma, Motif Video, and many more) without per-model code changes.