--- title: nvfp4_qat_config ---

NVFP4 quantization-aware (QAD) linear method, inference path.

Quantizes every targeted linear's weight to NVFP4 once at load time and runs each forward as a registered flashinfer-backed FP4 matmul. The original fp16/bf16 weight is popped immediately after quantization so the half-precision copy does not keep occupying GPU memory — that's what lets a Wan-2.1 pipeline stay fully resident on a single GPU without any CPU offloading.

The quantize / matmul custom ops are owned by :mod:fastvideo.layers.quantization.nvfp4_config and registered under the fastvideo_fp4:: namespace. We reuse them here for two reasons:

  1. Re-registering the same op name in a second module would raise.
  2. The registered ops have register_fake shape/dtype kernels, which is what makes the inference pipeline's per-block torch.compile trace through without graph breaks. Calling raw flashinfer functions (the old behavior of this file, plus a @torch.compile on apply) graph-breaks at every quantize and every matmul.

For QAT training, see nvfp4_qat_train_config which keeps the weight trainable and fake-quantizes on the fly via a straight-through estimator.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig

QuantizationConfig()

Base class for quantization configs.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_name

get_name() -> QuantizationMethods

Name of the quantization method.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_min_capability

get_min_capability() -> int

Minimum GPU capability to support the quantization method.

E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_config_filenames

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.from_config

from_config(config: dict[str, Any]) -> QuantizationConfig

Create a config class from the model's quantization config.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.override_quantization_method

override_quantization_method(hf_quant_cfg, user_quant) -> QuantizationMethods | None

Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_from_keys

get_from_keys(config: dict[str, Any], keys: list[str]) -> Any

Get a value from the model's quantization config.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_from_keys_or

get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any

Get a optional value from the model's quantization config.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_quant_method

get_quant_method(layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None

Get the quantize method to use for the quantized layer.

Parameters

layer: torch.nn.Module

The layer for the quant method.

prefix: str

The full name of the layer in the state dict

Returns: The quantize method. None if the given layer doesn't support quant method.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.validate_runtime

validate_runtime(device: torch.device) -> None

Fail fast on a device the scheme cannot execute on.

Called once by loaders that select a serialized checkpoint scheme, before any weight is read. The default accepts every device; schemes with kernel or capability requirements override it.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizationConfig.get_cache_scale

get_cache_scale(name: str) -> str | None

fastvideo.layers.quantization.nvfp4_qat_config.QuantizeMethodBase

QuantizeMethodBase()

Base class for different quantized methods.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizeMethodBase.create_weights

create_weights(layer: torch.nn.Module, *weight_args, **extra_weight_attrs)

Create weights for a layer.

The weights will be set as attributes of the layer.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizeMethodBase.apply

apply(layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor

Apply the weights in layer to the input tensor.

Expects create_weights to have been called before on the layer.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizeMethodBase.embedding

embedding(layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor

Gather embeddings in the layer based on indices in the input tensor.

Expects create_weights to have been called before on the layer.

fastvideo.layers.quantization.nvfp4_qat_config.QuantizeMethodBase.process_weights_after_loading

process_weights_after_loading(layer: nn.Module) -> None

Process the weight after loading.

This can be used for example, to transpose weights for computation.

fastvideo.layers.quantization.nvfp4_qat_config.set_weight_attrs

set_weight_attrs(weight: torch.Tensor, weight_attrs: dict[str, Any] | None)

Set attributes on a weight tensor.

This method is used to set attributes on a weight tensor. This method will not overwrite existing attributes.

Parameters

weight: torch.Tensor

The weight tensor.

weight_attrs: dict[str, Any] | None

A dictionary of attributes to set on the weight tensor.

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATQuantizeMethod

NVFP4QATQuantizeMethod()

Inference-only NVFP4 linear method with weight popping.

The dense weight parameter is materialized at load time only so that :func:convert_model_to_fp4 can read it once; the loader then removes it via mod._parameters.pop('weight'). From that point forward, apply reads only _fp4_weight / _fp4_weight_scale / _weight_global_sf.

Source
class NVFP4QATQuantizeMethod(QuantizeMethodBase):
    """Inference-only NVFP4 linear method with weight popping.

    The dense ``weight`` parameter is materialized at load time only so
    that :func:`convert_model_to_fp4` can read it once; the loader then
    removes it via ``mod._parameters.pop('weight')``. From that point
    forward, ``apply`` reads only ``_fp4_weight`` / ``_fp4_weight_scale``
    / ``_weight_global_sf``.
    """

    def __init__(self) -> None:
        super().__init__()
        self.weight_fp4 = None
        self.weight_scale = None
        # Static input global scale factor. Matches the FastVideo-Quantization
        # production path; recomputing it per-call via a ``.max()`` reduction
        # (the previous behavior) adds a sync point, costs a kernel launch,
        # and produces a data-dependent value that prevents CUDA-graph
        # capture under ``torch.compile(mode='reduce-overhead')``.
        self.x_global_sf = torch.tensor(1.0, device="cuda", dtype=torch.float32)

    def create_weights(self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int],
                       input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs) -> None:
        weight = Parameter(torch.empty(
            sum(output_partition_sizes),
            input_size_per_partition,
            dtype=params_dtype,
        ),
                           requires_grad=False)
        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
        layer.register_parameter("weight", weight)
        set_weight_attrs(weight, extra_weight_attrs)

    def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
        # ``_fp4_weight`` carries the (out, in/2) packed fp4 weight, so its
        # row count is the output dim even after the dense weight is popped.
        out_dim = layer._fp4_weight.shape[0]
        original_shape = x.shape

        assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
        x = x.view(-1, x.shape[-1])
        x_global_sf = self.x_global_sf
        x_fp4, x_scale = _nvfp4_quantize(
            x,
            x_global_sf,
            sfLayout=_layout_128x4(),
            do_shuffle=False,
        )

        weight_fp4 = layer._fp4_weight
        weight_scale = layer._fp4_weight_scale
        weight_global_sf = layer._weight_global_sf

        out = _mm_fp4(
            x_fp4,
            weight_fp4.T,
            x_scale,
            weight_scale.T,
            1.0 / (x_global_sf * weight_global_sf),
            torch.bfloat16,
            None,
            backend="cutlass",
        )

        if bias is not None:
            out = out + bias
        out = out.view(*original_shape[:-1], out_dim)
        return out

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATQuantizeMethod.create_weights

create_weights(layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs) -> None
Source
def create_weights(self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int],
                   input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs) -> None:
    weight = Parameter(torch.empty(
        sum(output_partition_sizes),
        input_size_per_partition,
        dtype=params_dtype,
    ),
                       requires_grad=False)
    set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
    layer.register_parameter("weight", weight)
    set_weight_attrs(weight, extra_weight_attrs)

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATQuantizeMethod.apply

apply(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor
Source
def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
    # ``_fp4_weight`` carries the (out, in/2) packed fp4 weight, so its
    # row count is the output dim even after the dense weight is popped.
    out_dim = layer._fp4_weight.shape[0]
    original_shape = x.shape

    assert x.dtype in (torch.bfloat16, torch.float16), (f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
    x = x.view(-1, x.shape[-1])
    x_global_sf = self.x_global_sf
    x_fp4, x_scale = _nvfp4_quantize(
        x,
        x_global_sf,
        sfLayout=_layout_128x4(),
        do_shuffle=False,
    )

    weight_fp4 = layer._fp4_weight
    weight_scale = layer._fp4_weight_scale
    weight_global_sf = layer._weight_global_sf

    out = _mm_fp4(
        x_fp4,
        weight_fp4.T,
        x_scale,
        weight_scale.T,
        1.0 / (x_global_sf * weight_global_sf),
        torch.bfloat16,
        None,
        backend="cutlass",
    )

    if bias is not None:
        out = out + bias
    out = out.view(*original_shape[:-1], out_dim)
    return out

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig

NVFP4QATConfig(target_layers: tuple[str, ...] | None = None)

NVFP4 (Wan-style) linear quantization, inference.

Parameters

target_layers: tuple[str, ...] | None

Substrings matched against each linear layer's prefix. A layer is quantized if any substring is contained in its prefix. Defaults to the standard Wan attention + FFN projections (:data:DEFAULT_FP4_LAYERS).

Source
class NVFP4QATConfig(QuantizationConfig):
    """NVFP4 (Wan-style) linear quantization, inference.

    Args:
        target_layers: Substrings matched against each linear layer's
            prefix. A layer is quantized if any substring is contained in
            its prefix. Defaults to the standard Wan attention + FFN
            projections (:data:`DEFAULT_FP4_LAYERS`).
    """

    def __init__(self, target_layers: tuple[str, ...] | None = None) -> None:
        super().__init__()
        self.target_layers = (tuple(target_layers) if target_layers else DEFAULT_FP4_LAYERS)

    def get_name(self) -> str:
        return "nvfp4_qat"

    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        return [torch.bfloat16, torch.float16]

    @classmethod
    def get_min_capability(cls) -> int:
        return 100

    @staticmethod
    def get_config_filenames() -> list[str]:
        return []

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> NVFP4QATConfig:
        target_layers = config.get("target_layers")
        if target_layers is not None:
            target_layers = tuple(target_layers)
        return cls(target_layers=target_layers)

    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        from fastvideo.layers.linear import LinearBase
        if isinstance(layer, LinearBase) and any(name in prefix for name in self.target_layers):
            return NVFP4QATQuantizeMethod()
        return None

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig.get_name

get_name() -> str
Source
def get_name(self) -> str:
    return "nvfp4_qat"

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]
Source
def get_supported_act_dtypes(self) -> list[torch.dtype]:
    return [torch.bfloat16, torch.float16]

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig.get_min_capability

get_min_capability() -> int
Source
@classmethod
def get_min_capability(cls) -> int:
    return 100

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig.get_config_filenames

get_config_filenames() -> list[str]
Source
@staticmethod
def get_config_filenames() -> list[str]:
    return []

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig.from_config

from_config(config: dict[str, Any]) -> NVFP4QATConfig
Source
@classmethod
def from_config(cls, config: dict[str, Any]) -> NVFP4QATConfig:
    target_layers = config.get("target_layers")
    if target_layers is not None:
        target_layers = tuple(target_layers)
    return cls(target_layers=target_layers)

fastvideo.layers.quantization.nvfp4_qat_config.NVFP4QATConfig.get_quant_method

get_quant_method(layer: torch.nn.Module, prefix: str)
Source
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
    from fastvideo.layers.linear import LinearBase
    if isinstance(layer, LinearBase) and any(name in prefix for name in self.target_layers):
        return NVFP4QATQuantizeMethod()
    return None

fastvideo.layers.quantization.nvfp4_qat_config.convert_model_to_fp4

convert_model_to_fp4(model: torch.nn.Module) -> None

Prequantize every FP4-tagged linear and drop its dense weight.

Walks the module tree, and for each layer whose quant_method is an :class:NVFP4QATQuantizeMethod, computes the NVFP4 packed weight / scale / global-scale buffers, then pops the original fp16/bf16 weight parameter so it no longer occupies GPU memory.

Source
def convert_model_to_fp4(model: torch.nn.Module) -> None:
    """Prequantize every FP4-tagged linear and drop its dense weight.

    Walks the module tree, and for each layer whose ``quant_method`` is
    an :class:`NVFP4QATQuantizeMethod`, computes the NVFP4 packed weight
    / scale / global-scale buffers, then pops the original fp16/bf16
    ``weight`` parameter so it no longer occupies GPU memory.
    """
    SfLayout, _, _ = _require_flashinfer()
    from torch.distributed.tensor import DTensor  # type: ignore

    with torch.no_grad():
        for mod in model.modules():
            qm = getattr(mod, "quant_method", None)
            if not isinstance(qm, NVFP4QATQuantizeMethod):
                continue

            weight = getattr(mod, "weight", None)
            if weight is None:
                continue

            weight_local = weight.to_local() if isinstance(weight, DTensor) else weight  # type: ignore[arg-type]

            # Only the reduced scalar needs fp32; avoid a full fp32 copy.
            weight_absmax = (weight_local.detach().abs().nan_to_num().amax().to(dtype=torch.float32))
            weight_global_sf = (448 * 6) / weight_absmax
            fp4_w, fp4_s = _nvfp4_quantize(
                weight_local,
                weight_global_sf,
                sfLayout=SfLayout.layout_128x4,
                do_shuffle=False,
            )
            mod.register_buffer("_fp4_weight", fp4_w, persistent=False)
            mod.register_buffer("_fp4_weight_scale", fp4_s, persistent=False)
            mod.register_buffer(
                "_weight_global_sf",
                weight_global_sf.to(dtype=torch.bfloat16),
                persistent=False,
            )

            # Drop the dense weight as soon as the fp4 buffers are installed
            # so it cannot keep occupying GPU memory.
            removed_weight = mod._parameters.pop("weight", None)
            if removed_weight is not None:
                removed_weight.grad = None
            del removed_weight, weight, weight_local, weight_absmax
    gc.collect()
    torch.cuda.empty_cache()