--- title: nvfp4_config ---

NVFP4 linear quantization for supported transformer layer sets.

NVFP4 is NVIDIA's block-scaled FP4 format (e2m1 mantissa, fp32 alpha, layout_128x4 scale layout, group size 16) — distinct from generic FP4 / OCP-FP4 / MX-FP4. We name the public surface NVFP4 explicitly so downstream callers don't conflate it with other FP4 variants that may land later (e.g. AMD's MX-FP4 or vendor-neutral e3m0).

The registered config targets the curated LTX-2 deployment set, the main MiniMax-H3 transformer-block FFN linears, and the packed MiniMax-H3 DiT export (layer_profile="h3_dit") covering attention plus FFN.

flashinfer is imported lazily inside the call paths that need it. This keeps import fastvideo cheap on hosts where flashinfer is not installed; only the actual NVFP4 quantize / matmul ops fail at use time, with a clear error.

fastvideo.layers.quantization.nvfp4_config.QuantizationConfig

QuantizationConfig()

Base class for quantization configs.

fastvideo.layers.quantization.nvfp4_config.QuantizationConfig.get_name

get_name() -> QuantizationMethods

Name of the quantization method.

fastvideo.layers.quantization.nvfp4_config.QuantizationConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

fastvideo.layers.quantization.nvfp4_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_config.QuantizationConfig.get_config_filenames

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

fastvideo.layers.quantization.nvfp4_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_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_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_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_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_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_config.QuantizationConfig.get_cache_scale

get_cache_scale(name: str) -> str | None

fastvideo.layers.quantization.nvfp4_config.QuantizeMethodBase

QuantizeMethodBase()

Base class for different quantized methods.

fastvideo.layers.quantization.nvfp4_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_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_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_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_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_config.minimax_h3_block_index

minimax_h3_block_index(prefix: str) -> int | None

Main transformer-block index of a MiniMax-H3 module prefix, or None outside the main blocks.

Source
def minimax_h3_block_index(prefix: str) -> int | None:
    """Main transformer-block index of a MiniMax-H3 module prefix, or None outside the main blocks."""
    match = _H3_MAIN_BLOCK.search(prefix)
    return int(match.group(1)) if match is not None else None

fastvideo.layers.quantization.nvfp4_config.read_minimax_h3_nvfp4_dit_skip_blocks

read_minimax_h3_nvfp4_dit_skip_blocks(path: str) -> tuple[int, ...]

The bf16 main blocks a packed NVFP4H3 export declares (empty for exports that quantize every block).

Source
def read_minimax_h3_nvfp4_dit_skip_blocks(path: str) -> tuple[int, ...]:
    """The bf16 main blocks a packed NVFP4H3 export declares (empty for exports that quantize every block)."""
    import json

    from safetensors import safe_open

    with safe_open(path, framework="pt", device="cpu") as reader:
        raw = (reader.metadata() or {}).get(H3_NVFP4_DIT_SKIP_BLOCKS_KEY)
    if raw is None:
        return ()
    blocks = json.loads(raw)
    if not isinstance(blocks, list) or not all(
            isinstance(b, int) and not isinstance(b, bool) and b >= 0 for b in blocks):
        raise ValueError(f"{path}: metadata {H3_NVFP4_DIT_SKIP_BLOCKS_KEY!r} must be a JSON list of block indices")
    return tuple(sorted(set(blocks)))

fastvideo.layers.quantization.nvfp4_config.is_ltx2_nvfp4_linear_prefix

is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool

Return whether prefix belongs to the LTX-2 NVFP4 deployment set.

Source
def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool:
    """Return whether *prefix* belongs to the LTX-2 NVFP4 deployment set."""
    return prefix in _LTX2_NVFP4_LINEAR_PREFIXES

fastvideo.layers.quantization.nvfp4_config.is_minimax_h3_nvfp4_linear_prefix

is_minimax_h3_nvfp4_linear_prefix(prefix: str) -> bool
Source
def is_minimax_h3_nvfp4_linear_prefix(prefix: str) -> bool:
    return _MINIMAX_H3_NVFP4_FF_PREFIX.search(prefix) is not None

fastvideo.layers.quantization.nvfp4_config.is_minimax_h3_nvfp4_dit_linear_prefix

is_minimax_h3_nvfp4_dit_linear_prefix(prefix: str) -> bool

Return whether prefix is a MiniMax-H3 DiT attention or FFN linear.

This is the packed NVFP4H3 export set: attn.to_{q,k,v,out} and ff.{fc_in,fc_out} in each main transformer block. Token-refiner, AdaLN, and embedding linears stay dense.

Source
def is_minimax_h3_nvfp4_dit_linear_prefix(prefix: str) -> bool:
    """Return whether *prefix* is a MiniMax-H3 DiT attention or FFN linear.

    This is the packed NVFP4H3 export set: ``attn.to_{q,k,v,out}`` and
    ``ff.{fc_in,fc_out}`` in each main transformer block. Token-refiner,
    AdaLN, and embedding linears stay dense.
    """
    return _MINIMAX_H3_NVFP4_DIT_PREFIX.search(prefix) is not None

fastvideo.layers.quantization.nvfp4_config.is_minimax_h3_nvfp4_dit_export_path

is_minimax_h3_nvfp4_dit_export_path(path: str) -> bool
Source
def is_minimax_h3_nvfp4_dit_export_path(path: str) -> bool:
    return os.path.basename(path) == H3_NVFP4_DIT_EXPORT_FILENAME

fastvideo.layers.quantization.nvfp4_config.find_minimax_h3_nvfp4_dit_export

find_minimax_h3_nvfp4_dit_export(weight_paths: list[str]) -> str | None
Source
def find_minimax_h3_nvfp4_dit_export(weight_paths: list[str]) -> str | None:
    seen: list[str] = []
    for path in weight_paths:
        if is_minimax_h3_nvfp4_dit_export_path(path) and os.path.isfile(path):
            return path
        directory = path if os.path.isdir(path) else os.path.dirname(path)
        if directory and directory not in seen:
            seen.append(directory)
    for directory in seen:
        candidate = os.path.join(directory, H3_NVFP4_DIT_EXPORT_FILENAME)
        if os.path.isfile(candidate):
            return candidate
    return None

fastvideo.layers.quantization.nvfp4_config.dense_transformer_safetensors

dense_transformer_safetensors(weight_paths: list[str]) -> list[str]

Drop the packed NVFP4 DiT export so it is not loaded as bf16 weights.

Source
def dense_transformer_safetensors(weight_paths: list[str]) -> list[str]:
    """Drop the packed NVFP4 DiT export so it is not loaded as bf16 weights."""
    return [path for path in weight_paths if not is_minimax_h3_nvfp4_dit_export_path(path)]

fastvideo.layers.quantization.nvfp4_config.nvfp4_quantize_fenced

nvfp4_quantize_fenced(x: torch.Tensor, global_sf: torch.Tensor, sf_layout: int, do_shuffle: bool = False) -> tuple[torch.Tensor, torch.Tensor]

FlashInfer NVFP4 quantization with the DGX Spark (GB10) ordering fence.

Every FastVideo NVFP4 quantization goes through here, so the GB10 workaround covers inference and the QAT straight-through linear alike.

Source
def nvfp4_quantize_fenced(
    x: torch.Tensor,
    global_sf: torch.Tensor,
    sf_layout: int,
    do_shuffle: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """FlashInfer NVFP4 quantization with the DGX Spark (GB10) ordering fence.

    Every FastVideo NVFP4 quantization goes through here, so the GB10
    workaround covers inference and the QAT straight-through linear alike.
    """
    device_index = x.device.index if x.device.index is not None else torch.cuda.current_device()
    spark = _is_dgx_spark(device_index)
    if spark and torch.cuda.is_current_stream_capturing():
        raise RuntimeError("NVFP4 activation quantization on DGX Spark requires a completion fence; "
                           "disable CUDA graph capture.")
    SfLayout, _, nvfp4_quantize = _require_flashinfer()
    if not spark:
        return nvfp4_quantize(x, global_sf, sfLayout=SfLayout(sf_layout), do_shuffle=do_shuffle)
    # FlashInfer's PDL kernel reads the global scale before its
    # dependency wait. Fresh dynamic scales require normal ordering.
    quantized, scales = nvfp4_quantize(x,
                                       global_sf,
                                       sfLayout=SfLayout(sf_layout),
                                       do_shuffle=do_shuffle,
                                       enable_pdl=False)
    # With FlashInfer 0.6.18 and 0.7.0 on GB10, queued activation
    # quantization plus GEMM can diverge: identical seeded H3 requests
    # produce different latents. Completing each quantization on the
    # host prevents it, at no measurable denoise cost.
    torch.cuda.current_stream(x.device).synchronize()
    return quantized, scales

fastvideo.layers.quantization.nvfp4_config.NVFP4QuantizeMethod

NVFP4QuantizeMethod(layer_prefix: str = '')
Source
class NVFP4QuantizeMethod(QuantizeMethodBase):

    # Lazily resolved by _static_activation_global_sf; class defaults also cover object.__new__ test doubles.
    _static_sf_checked: bool = False
    _static_sf: torch.Tensor | None = None

    def __init__(self, layer_prefix: str = ""):
        super().__init__()
        self.weight_fp4 = None
        self.weight_scale = None
        self.x_global_sf = torch.tensor(1.0, device="cuda", dtype=torch.float32)
        self.layer_prefix = layer_prefix
        self._is_refine_only_layer = _is_ltx2_refine_only_prefix(layer_prefix)
        # Set from NVFP4Config.retain_original_weights in get_quant_method:
        # True = retain every original bf16 weight; None/False (default) =
        # purge the purgeable set. Refine-only layers are always retained --
        # the base stage profile runs them dense by deployment contract.
        self._retain_original_weights: bool | None = None

    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):
        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 _static_activation_global_sf(self) -> torch.Tensor | None:
        """FASTVIDEO_NVFP4_ACT_AMAX: JSON of calibrated input amax per layer ("b<block>.<sub>" or full prefix)."""
        if getattr(self, "_static_sf_checked", False):
            return self._static_sf
        self._static_sf_checked, self._static_sf = True, None
        path = envs.FASTVIDEO_NVFP4_ACT_AMAX.get()
        if path:
            table = _load_amax_table(path)
            prefix = self.layer_prefix or ""
            match = re.search(r"transformer_blocks\.(\d+)\.(.+)$", prefix)
            keys = [prefix] + ([f"b{match.group(1)}.{match.group(2)}"] if match else [])
            amax = next((table[k] for k in keys if k in table), None)
            if amax is not None:
                self._static_sf = torch.tensor((448.0 * 6.0) / max(amax, 1e-12), dtype=torch.float32, device="cuda")
        return self._static_sf

    def _dynamic_activation_scale(self) -> bool:
        """FASTVIDEO_NVFP4_DYNAMIC_ACT: "all", or comma-separated layer-name suffixes (e.g. "ff.fc_out")."""
        cached = getattr(self, "_dynamic_act_cached", None)
        if cached is None:
            selected = envs.FASTVIDEO_NVFP4_DYNAMIC_ACT.get()
            suffixes = [part.strip() for part in selected.split(",") if part.strip()]
            prefix = self.layer_prefix or ""
            cached = "all" in suffixes or any(prefix.endswith(suffix) for suffix in suffixes)
            self._dynamic_act_cached = cached
        return cached

    def uses_unit_activation_scale(self, layer: torch.nn.Module) -> bool:
        """Whether ``apply`` quantizes this layer's input with the unit global scale.

        False when a calibrated scale (env table or the export's ``_nvfp4_input_global_sf``) or a dynamic
        per-call scale applies; such inputs cannot share one pre-quantized copy across layers.
        """
        return (self._static_activation_global_sf() is None and getattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, None) is None
                and not self._dynamic_activation_scale())

    def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        SfLayout, _, _ = _require_flashinfer()
        x = _coerce_fp4_input_dtype(x)
        x_2d = x.view(-1, x.shape[-1])
        x_fp4, x_scale = _nvfp4_quantize(
            x_2d,
            self.x_global_sf,
            sfLayout=SfLayout.layout_128x4,
            do_shuffle=False,
        )
        return x_fp4, x_scale, self.x_global_sf

    def wants_prequantized_input(self) -> bool:
        if not self._is_refine_only_layer:
            return True
        stage_profile = _get_ltx2_fp4_stage_profile(default="refine")
        return stage_profile != "base"

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
        pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        | None = None,
    ) -> torch.Tensor:
        SfLayout, _, _ = _require_flashinfer()
        # The original bf16 weight may have been purged after FP4 conversion
        # (see convert_model_to_nvfp4); the packed FP4 weight keeps the
        # output dim as its first dimension (only K is packed 2-per-byte).
        weight = getattr(layer, "weight", None)
        out_dim = weight.shape[0] if weight is not None else layer._nvfp4_weight.shape[0]
        original_shape = x.shape

        # Stage-aware profile: keep refine-only FP4 layers in dense mode
        # during stage-1 denoising so the base path doesn't pay the
        # quantize/dequantize tax for layers it never touches.
        stage_profile = _get_ltx2_fp4_stage_profile(default="refine")
        if self._is_refine_only_layer and stage_profile == "base":
            if weight is None:
                raise RuntimeError(f"NVFP4 layer {self.layer_prefix!r} hit the stage-profile dense path, "
                                   "but its original weights were purged "
                                   "(NVFP4Config(retain_original_weights=False)). Streaming/two-stage "
                                   "deploys must load with retain_original_weights left unset (auto) or True.")
            out = (F.linear(x, weight, bias) if torch.cuda.is_available() or bias is None else F.linear(
                x, weight, bias.to(x.dtype)))
            return out.view(*original_shape[:-1], out_dim)
        if pre_quantized is not None:
            x_fp4, x_scale, x_global_sf = pre_quantized
            # FlashInfer fused norm+quant APIs may return 3D tensors for
            # 3D inputs. mm_fp4 only accepts 2D tensors, so flatten
            # batch/sequence dims here.
            if x_fp4.dim() > 2:
                x_fp4 = x_fp4.view(-1, x_fp4.shape[-1])
            if x_scale.dim() > 2:
                x_scale = x_scale.view(-1, x_scale.shape[-1])
        else:
            x = _coerce_fp4_input_dtype(x)
            x = x.view(-1, x.shape[-1])
            static_sf = self._static_activation_global_sf()
            if static_sf is None:
                static_sf = getattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, None)
            if static_sf is not None:
                x_global_sf = static_sf
            elif self._dynamic_activation_scale():
                # A unit global scale caps FP8 block scales at |x| = 6 * 448; inputs such as H3's ff.fc_out
                # (post-SwiGLU) exceed that, so derive the global scale from this call's amax.
                x_global_sf = (448.0 * 6.0) / x.abs().amax().float().clamp(min=1e-12)
            else:
                x_global_sf = self.x_global_sf
            x_fp4, x_scale = _nvfp4_quantize(
                x,
                x_global_sf,
                sfLayout=SfLayout.layout_128x4,
                do_shuffle=False,
            )

        weight_fp4 = layer._nvfp4_weight
        weight_scale = layer._nvfp4_weight_scale
        weight_global_sf = layer._weight_global_sf

        if hasattr(layer, "_nvfp4_alpha"):
            alpha = layer._nvfp4_alpha / x_global_sf
        else:
            alpha = 1.0 / (x_global_sf * weight_global_sf)

        out = _mm_fp4(
            x_fp4,
            weight_fp4.T,
            x_scale,
            weight_scale.T,
            alpha,
            torch.bfloat16,
            None,
            backend=_mm_fp4_backend(),
        )

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

fastvideo.layers.quantization.nvfp4_config.NVFP4QuantizeMethod.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)
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):
    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_config.NVFP4QuantizeMethod.uses_unit_activation_scale

uses_unit_activation_scale(layer: torch.nn.Module) -> bool

Whether apply quantizes this layer's input with the unit global scale.

False when a calibrated scale (env table or the export's _nvfp4_input_global_sf) or a dynamic per-call scale applies; such inputs cannot share one pre-quantized copy across layers.

Source
def uses_unit_activation_scale(self, layer: torch.nn.Module) -> bool:
    """Whether ``apply`` quantizes this layer's input with the unit global scale.

    False when a calibrated scale (env table or the export's ``_nvfp4_input_global_sf``) or a dynamic
    per-call scale applies; such inputs cannot share one pre-quantized copy across layers.
    """
    return (self._static_activation_global_sf() is None and getattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, None) is None
            and not self._dynamic_activation_scale())

fastvideo.layers.quantization.nvfp4_config.NVFP4QuantizeMethod.quantize_input

quantize_input(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
Source
def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    SfLayout, _, _ = _require_flashinfer()
    x = _coerce_fp4_input_dtype(x)
    x_2d = x.view(-1, x.shape[-1])
    x_fp4, x_scale = _nvfp4_quantize(
        x_2d,
        self.x_global_sf,
        sfLayout=SfLayout.layout_128x4,
        do_shuffle=False,
    )
    return x_fp4, x_scale, self.x_global_sf

fastvideo.layers.quantization.nvfp4_config.NVFP4QuantizeMethod.wants_prequantized_input

wants_prequantized_input() -> bool
Source
def wants_prequantized_input(self) -> bool:
    if not self._is_refine_only_layer:
        return True
    stage_profile = _get_ltx2_fp4_stage_profile(default="refine")
    return stage_profile != "base"

fastvideo.layers.quantization.nvfp4_config.NVFP4QuantizeMethod.apply

apply(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor
Source
def apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: torch.Tensor | None = None,
    pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
    | None = None,
) -> torch.Tensor:
    SfLayout, _, _ = _require_flashinfer()
    # The original bf16 weight may have been purged after FP4 conversion
    # (see convert_model_to_nvfp4); the packed FP4 weight keeps the
    # output dim as its first dimension (only K is packed 2-per-byte).
    weight = getattr(layer, "weight", None)
    out_dim = weight.shape[0] if weight is not None else layer._nvfp4_weight.shape[0]
    original_shape = x.shape

    # Stage-aware profile: keep refine-only FP4 layers in dense mode
    # during stage-1 denoising so the base path doesn't pay the
    # quantize/dequantize tax for layers it never touches.
    stage_profile = _get_ltx2_fp4_stage_profile(default="refine")
    if self._is_refine_only_layer and stage_profile == "base":
        if weight is None:
            raise RuntimeError(f"NVFP4 layer {self.layer_prefix!r} hit the stage-profile dense path, "
                               "but its original weights were purged "
                               "(NVFP4Config(retain_original_weights=False)). Streaming/two-stage "
                               "deploys must load with retain_original_weights left unset (auto) or True.")
        out = (F.linear(x, weight, bias) if torch.cuda.is_available() or bias is None else F.linear(
            x, weight, bias.to(x.dtype)))
        return out.view(*original_shape[:-1], out_dim)
    if pre_quantized is not None:
        x_fp4, x_scale, x_global_sf = pre_quantized
        # FlashInfer fused norm+quant APIs may return 3D tensors for
        # 3D inputs. mm_fp4 only accepts 2D tensors, so flatten
        # batch/sequence dims here.
        if x_fp4.dim() > 2:
            x_fp4 = x_fp4.view(-1, x_fp4.shape[-1])
        if x_scale.dim() > 2:
            x_scale = x_scale.view(-1, x_scale.shape[-1])
    else:
        x = _coerce_fp4_input_dtype(x)
        x = x.view(-1, x.shape[-1])
        static_sf = self._static_activation_global_sf()
        if static_sf is None:
            static_sf = getattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, None)
        if static_sf is not None:
            x_global_sf = static_sf
        elif self._dynamic_activation_scale():
            # A unit global scale caps FP8 block scales at |x| = 6 * 448; inputs such as H3's ff.fc_out
            # (post-SwiGLU) exceed that, so derive the global scale from this call's amax.
            x_global_sf = (448.0 * 6.0) / x.abs().amax().float().clamp(min=1e-12)
        else:
            x_global_sf = self.x_global_sf
        x_fp4, x_scale = _nvfp4_quantize(
            x,
            x_global_sf,
            sfLayout=SfLayout.layout_128x4,
            do_shuffle=False,
        )

    weight_fp4 = layer._nvfp4_weight
    weight_scale = layer._nvfp4_weight_scale
    weight_global_sf = layer._weight_global_sf

    if hasattr(layer, "_nvfp4_alpha"):
        alpha = layer._nvfp4_alpha / x_global_sf
    else:
        alpha = 1.0 / (x_global_sf * weight_global_sf)

    out = _mm_fp4(
        x_fp4,
        weight_fp4.T,
        x_scale,
        weight_scale.T,
        alpha,
        torch.bfloat16,
        None,
        backend=_mm_fp4_backend(),
    )

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

fastvideo.layers.quantization.nvfp4_config.NVFP4Config

NVFP4Config(layer_profile: str = 'refine', retain_original_weights: bool | None = None, skip_blocks: tuple[int, ...] | list[int] = ())

Select NVFP4 for the supported LTX-2 and MiniMax-H3 linear sets.

NVFP4 is NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha, layout_128x4 scale layout, group size 16). LTX-2 uses its curated attention and FFN deployment set. MiniMax-H3's default profile uses only fc_in and fc_out in each main transformer-block FFN. layer_profile="h3_dit" expands that to the packed NVFP4H3 DiT set (attention to_{q,k,v,out} plus those FFN linears). layer_profile="h3_dit_ffn" loads a packed export of the FFN linears only, keeping attention projections dense (e.g. calibrated FFN-only checkpoints such as FastH3 V2 NVFP4). layer_profile="h3_dit_vsa" is h3_dit plus each block's VSA compression gate attn.to_gate_compress (VSA-distilled students). skip_blocks keeps every linear of those MiniMax-H3 main blocks in bf16 (mixed exports such as FastH3 OmniRef QAD, blocks 42-49). A packed export records its own list, which the loader adopts when this is empty.

Source
class NVFP4Config(QuantizationConfig):
    """Select NVFP4 for the supported LTX-2 and MiniMax-H3 linear sets.

    NVFP4 is NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha,
    ``layout_128x4`` scale layout, group size 16). LTX-2 uses its curated
    attention and FFN deployment set. MiniMax-H3's default profile uses only
    ``fc_in`` and ``fc_out`` in each main transformer-block FFN.
    ``layer_profile="h3_dit"`` expands that to the packed NVFP4H3 DiT set
    (attention ``to_{q,k,v,out}`` plus those FFN linears).
    ``layer_profile="h3_dit_ffn"`` loads a packed export of the FFN linears
    only, keeping attention projections dense (e.g. calibrated FFN-only
    checkpoints such as FastH3 V2 NVFP4).
    ``layer_profile="h3_dit_vsa"`` is ``h3_dit`` plus each block's VSA
    compression gate ``attn.to_gate_compress`` (VSA-distilled students).
    ``skip_blocks`` keeps every linear of those MiniMax-H3 main blocks in
    bf16 (mixed exports such as FastH3 OmniRef QAD, blocks 42-49). A packed
    export records its own list, which the loader adopts when this is empty.
    """

    def __init__(self,
                 layer_profile: str = "refine",
                 retain_original_weights: bool | None = None,
                 skip_blocks: tuple[int, ...] | list[int] = ()):
        super().__init__()
        if any(isinstance(b, bool) or not isinstance(b, int) or b < 0 for b in skip_blocks):
            raise ValueError(f"NVFP4Config.skip_blocks must be non-negative block indices, got {skip_blocks!r}")
        self.skip_blocks: tuple[int, ...] = tuple(sorted(set(skip_blocks)))
        if layer_profile not in ("base", "refine", "h3_dit", "h3_dit_ffn", "h3_dit_vsa"):
            raise ValueError("NVFP4Config.layer_profile must be one of 'base', 'refine', 'h3_dit', "
                             f"'h3_dit_ffn', or 'h3_dit_vsa', got {layer_profile!r}")
        self.layer_profile = layer_profile
        # Original bf16 ``layer.weight`` retention after FP4 conversion.
        # Default (None/False): purge the purgeable originals -- every
        # always-FP4 layer. Refine-only layers (the cross-modal AV
        # projections) are ALWAYS retained: the ``base`` stage profile runs
        # them dense by deployment contract, including the distilled
        # single-stage deploy. True: retain everything (debugging /
        # pre-purge behavior).
        self.retain_original_weights = retain_original_weights

    def get_name(self):
        return "nvfp4"

    def get_supported_act_dtypes(self):
        return [torch.bfloat16, torch.float16]

    @classmethod
    def get_min_capability(cls):
        return 100

    @staticmethod
    def get_config_filenames():
        return []

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> NVFP4Config:
        return cls(
            layer_profile=config.get("layer_profile", "refine"),
            retain_original_weights=config.get("retain_original_weights"),
            skip_blocks=tuple(config.get("skip_blocks", ())),
        )

    def with_skip_blocks(self, skip_blocks: tuple[int, ...]) -> NVFP4Config:
        """A copy of this config with ``skip_blocks`` replaced."""
        return NVFP4Config(layer_profile=self.layer_profile,
                           retain_original_weights=self.retain_original_weights,
                           skip_blocks=skip_blocks)

    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        from fastvideo.layers.linear import LinearBase

        if not isinstance(layer, LinearBase):
            return None
        if self.skip_blocks and minimax_h3_block_index(prefix) in self.skip_blocks:
            return None
        if self.layer_profile == "h3_dit":
            tagged = is_minimax_h3_nvfp4_dit_linear_prefix(prefix)
        elif self.layer_profile == "h3_dit_vsa":
            tagged = (is_minimax_h3_nvfp4_dit_linear_prefix(prefix)
                      or _MINIMAX_H3_NVFP4_VSA_GATE_PREFIX.search(prefix) is not None)
        elif self.layer_profile == "h3_dit_ffn":
            tagged = is_minimax_h3_nvfp4_linear_prefix(prefix)
        else:
            tagged = is_ltx2_nvfp4_linear_prefix(prefix) or is_minimax_h3_nvfp4_linear_prefix(prefix)
        if tagged:
            method = NVFP4QuantizeMethod(layer_prefix=prefix)
            method._retain_original_weights = self.retain_original_weights
            return method
        if (self.layer_profile == "h3_dit_ffn" and envs.FASTVIDEO_H3_FP8_ATTENTION.get()
                and _H3_BLOCK_ATTN_PROJ.search(prefix) is not None):
            # Mixed precision: NVFP4 MLPs, FP8 (per-tensor weight, dynamic per-tensor activation) attention.
            from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod
            return FP8QuantizeMethod(granularity=envs.FASTVIDEO_H3_FP8_GRANULARITY.get())
        return None

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.get_name

get_name()
Source
def get_name(self):
    return "nvfp4"

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.get_supported_act_dtypes

get_supported_act_dtypes()
Source
def get_supported_act_dtypes(self):
    return [torch.bfloat16, torch.float16]

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.get_min_capability

get_min_capability()
Source
@classmethod
def get_min_capability(cls):
    return 100

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.get_config_filenames

get_config_filenames()
Source
@staticmethod
def get_config_filenames():
    return []

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.from_config

from_config(config: dict[str, Any]) -> NVFP4Config
Source
@classmethod
def from_config(cls, config: dict[str, Any]) -> NVFP4Config:
    return cls(
        layer_profile=config.get("layer_profile", "refine"),
        retain_original_weights=config.get("retain_original_weights"),
        skip_blocks=tuple(config.get("skip_blocks", ())),
    )

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.with_skip_blocks

with_skip_blocks(skip_blocks: tuple[int, ...]) -> NVFP4Config

A copy of this config with skip_blocks replaced.

Source
def with_skip_blocks(self, skip_blocks: tuple[int, ...]) -> NVFP4Config:
    """A copy of this config with ``skip_blocks`` replaced."""
    return NVFP4Config(layer_profile=self.layer_profile,
                       retain_original_weights=self.retain_original_weights,
                       skip_blocks=skip_blocks)

fastvideo.layers.quantization.nvfp4_config.NVFP4Config.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 not isinstance(layer, LinearBase):
        return None
    if self.skip_blocks and minimax_h3_block_index(prefix) in self.skip_blocks:
        return None
    if self.layer_profile == "h3_dit":
        tagged = is_minimax_h3_nvfp4_dit_linear_prefix(prefix)
    elif self.layer_profile == "h3_dit_vsa":
        tagged = (is_minimax_h3_nvfp4_dit_linear_prefix(prefix)
                  or _MINIMAX_H3_NVFP4_VSA_GATE_PREFIX.search(prefix) is not None)
    elif self.layer_profile == "h3_dit_ffn":
        tagged = is_minimax_h3_nvfp4_linear_prefix(prefix)
    else:
        tagged = is_ltx2_nvfp4_linear_prefix(prefix) or is_minimax_h3_nvfp4_linear_prefix(prefix)
    if tagged:
        method = NVFP4QuantizeMethod(layer_prefix=prefix)
        method._retain_original_weights = self.retain_original_weights
        return method
    if (self.layer_profile == "h3_dit_ffn" and envs.FASTVIDEO_H3_FP8_ATTENTION.get()
            and _H3_BLOCK_ATTN_PROJ.search(prefix) is not None):
        # Mixed precision: NVFP4 MLPs, FP8 (per-tensor weight, dynamic per-tensor activation) attention.
        from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod
        return FP8QuantizeMethod(granularity=envs.FASTVIDEO_H3_FP8_GRANULARITY.get())
    return None

fastvideo.layers.quantization.nvfp4_config.convert_model_to_nvfp4

convert_model_to_nvfp4(model: torch.nn.Module) -> None
Source
def convert_model_to_nvfp4(model: torch.nn.Module) -> None:
    SfLayout, _, _ = _require_flashinfer()
    from torch.distributed.tensor import DTensor  # type: ignore

    purged = 0
    retained = 0
    purged_bytes = 0
    for mod in model.modules():
        qm = getattr(mod, "quant_method", None)
        if isinstance(qm, NVFP4QuantizeMethod):
            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]
            weight_global_sf = (448 * 6) / weight_local.float().abs().nan_to_num().max()
            fp4_w, fp4_s = _nvfp4_quantize(
                weight_local,
                weight_global_sf,
                sfLayout=SfLayout.layout_128x4,
                do_shuffle=False,
            )
            weight_global_sf_t = torch.as_tensor(
                weight_global_sf,
                device=weight_local.device,
                dtype=torch.float32,
            )
            mod.register_buffer("_nvfp4_weight", fp4_w, persistent=False)
            mod.register_buffer("_nvfp4_weight_scale", fp4_s, persistent=False)
            mod.register_buffer(
                "_weight_global_sf",
                weight_global_sf_t.to(dtype=torch.bfloat16),
                persistent=False,
            )
            mod.register_buffer(
                "_nvfp4_alpha",
                (1.0 / weight_global_sf_t).to(dtype=torch.float32),
                persistent=False,
            )

            retain_flag = getattr(qm, "_retain_original_weights", None)
            # Refine-only layers are NEVER purgeable: the "base" stage profile
            # runs them dense by deployment contract (the distilled
            # single-stage deploy included — its forward context is the base
            # profile, so e.g. audio_to_video_attn routes dense every step).
            # retain_original_weights therefore only widens retention
            # (True = keep everything); it cannot narrow it below the
            # dense-capable set.
            retain = qm._is_refine_only_layer or retain_flag is True
            if retain:
                retained += 1
            elif isinstance(weight, DTensor):
                raise RuntimeError("NVFP4 cannot purge FSDP-sharded bf16 weights. Use a packed NVFP4 "
                                   "export, or convert without FSDP sharding.")
            else:
                purged_bytes += weight.numel() * weight.element_size()
                purged += 1
                mod.register_parameter("weight", None)

    if purged or retained:
        logger.info(
            "NVFP4 weight purge receipt: purged %d original bf16 weight tensors "
            "(%.2f GiB freed); retained %d (refine-only dense fallback or "
            "retain_original_weights).",
            purged,
            purged_bytes / (1 << 30),
            retained,
        )

fastvideo.layers.quantization.nvfp4_config.nvfp4_linear_weight_param_names

nvfp4_linear_weight_param_names(model: torch.nn.Module) -> set[str]

State-dict names of weight on layers tagged with NVFP4QuantizeMethod.

Source
def nvfp4_linear_weight_param_names(model: torch.nn.Module) -> set[str]:
    """State-dict names of ``weight`` on layers tagged with ``NVFP4QuantizeMethod``."""
    names: set[str] = set()
    for module_name, module in model.named_modules():
        if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod):
            names.add(f"{module_name}.weight" if module_name else "weight")
    return names

fastvideo.layers.quantization.nvfp4_config.load_minimax_h3_nvfp4_dit_export

load_minimax_h3_nvfp4_dit_export(model: torch.nn.Module, path: str, device: torch.device | str) -> int

Load a packed NVFP4H3 DiT export onto already-tagged NVFP4 linears.

Keys are <module>::<buffer> with the four buffers convert_model_to_nvfp4 registers, plus an optional calibrated _nvfp4_input_global_sf. Every export prefix must match an NVFP4 linear, and every NVFP4 linear must appear in the export.

Source
def load_minimax_h3_nvfp4_dit_export(
    model: torch.nn.Module,
    path: str,
    device: torch.device | str,
) -> int:
    """Load a packed NVFP4H3 DiT export onto already-tagged NVFP4 linears.

    Keys are ``<module>::<buffer>`` with the four buffers
    ``convert_model_to_nvfp4`` registers, plus an optional calibrated
    ``_nvfp4_input_global_sf``. Every export prefix must match an
    NVFP4 linear, and every NVFP4 linear must appear in the export.
    """
    from safetensors import safe_open

    modules = dict(model.named_modules())
    tagged = {
        name
        for name, module in modules.items() if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
    }
    groups: dict[str, dict[str, str]] = {}
    with safe_open(path, framework="pt", device="cpu") as reader:
        for key in reader.keys():  # noqa: SIM118
            if H3_NVFP4_DIT_KEY_SEP not in key:
                raise ValueError(f"MiniMax-H3 NVFP4 DiT export key {key!r} is missing {H3_NVFP4_DIT_KEY_SEP!r}")
            prefix, buffer_name = key.split(H3_NVFP4_DIT_KEY_SEP, 1)
            groups.setdefault(prefix, {})[buffer_name] = key

        loaded_names: set[str] = set()
        for prefix, buffers in groups.items():
            missing_buffers = [name for name in H3_NVFP4_DIT_BUFFER_NAMES if name not in buffers]
            if missing_buffers:
                raise ValueError(f"MiniMax-H3 NVFP4 DiT export layer {prefix!r} is missing {missing_buffers}")
            extra_buffers = sorted(set(buffers) - set(H3_NVFP4_DIT_BUFFER_NAMES) - {H3_NVFP4_DIT_INPUT_SF_NAME})
            if extra_buffers:
                raise ValueError(f"MiniMax-H3 NVFP4 DiT export layer {prefix!r} has unknown buffers {extra_buffers}")
            module = _module_by_nvfp4_export_prefix(modules, prefix)
            if module is None:
                raise ValueError(f"MiniMax-H3 NVFP4 DiT export layer {prefix!r} is not in the model")
            if not isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod):
                raise RuntimeError("MiniMax-H3 NVFP4 DiT export layer "
                                   f"{prefix!r} is not an NVFP4 linear; set NVFP4Config(layer_profile='h3_dit')")
            for buffer_name in H3_NVFP4_DIT_BUFFER_NAMES + (H3_NVFP4_DIT_INPUT_SF_NAME, ):
                if buffer_name not in buffers:
                    continue
                tensor = reader.get_tensor(buffers[buffer_name]).to(device=device)
                module.register_buffer(buffer_name, tensor, persistent=False)
            module.register_parameter("weight", None)
            loaded_names.add(next(name for name, candidate in modules.items() if candidate is module))

    missing_layers = tagged - loaded_names
    extra_layers = loaded_names - tagged
    if missing_layers or extra_layers:
        raise RuntimeError("MiniMax-H3 NVFP4 DiT export does not cover the tagged linear set; "
                           f"missing={sorted(missing_layers)[:8]} extra={sorted(extra_layers)[:8]}")
    logger.info("Loaded MiniMax-H3 NVFP4 DiT export: %d linears from %s", len(loaded_names), path)
    return len(loaded_names)