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() -> QuantizationMethodsName 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() -> intMinimum 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]) -> QuantizationConfigCreate 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 | NoneDetects 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]) -> AnyGet 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) -> AnyGet 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 | NoneGet the quantize method to use for the quantized layer.
Parameters
layer: torch.nn.ModuleThe layer for the quant method.
prefix: strThe 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) -> NoneFail 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 | Nonefastvideo.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.TensorApply 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.TensorGather 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) -> NoneProcess 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.TensorThe weight tensor.
weight_attrs: dict[str, Any] | NoneA 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 | NoneMain 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 Nonefastvideo.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)))</code></pre></details>
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</code></pre></details>
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</code></pre></details>
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</code></pre></details>
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())</code></pre></details>
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</code></pre></details>
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</code></pre></details>
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</code></pre></details>
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,
)</code></pre></details>
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)</code></pre></details>