nvfp4_qat_config

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</code></pre></details>
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&quot;only allow bf16/fp16 inputs to fp4 linear, got &#123;x.dtype&#125;&quot;)
x = x.view(-1, x.shape&#91;-1&#93;)
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=&quot;cutlass&quot;,
)

if bias is not None:
    out = out + bias
out = out.view(*original_shape&#91;:-1&#93;, out_dim)
return out</code></pre></details>
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&#x27;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`).
&quot;&quot;&quot;

def __init__(self, target_layers: tuple&#91;str, ...&#93; | None = None) -&gt; None:
    super().__init__()
    self.target_layers = (tuple(target_layers) if target_layers else DEFAULT_FP4_LAYERS)

def get_name(self) -&gt; str:
    return &quot;nvfp4_qat&quot;

def get_supported_act_dtypes(self) -&gt; list&#91;torch.dtype&#93;:
    return &#91;torch.bfloat16, torch.float16&#93;

@classmethod
def get_min_capability(cls) -&gt; int:
    return 100

@staticmethod
def get_config_filenames() -&gt; list&#91;str&#93;:
    return &#91;&#93;

@classmethod
def from_config(cls, config: dict&#91;str, Any&#93;) -&gt; NVFP4QATConfig:
    target_layers = config.get(&quot;target_layers&quot;)
    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</code></pre></details>
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.
&quot;&quot;&quot;
SfLayout, _, _ = _require_flashinfer()
from torch.distributed.tensor import DTensor  # type: ignore

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

        weight = getattr(mod, &quot;weight&quot;, None)
        if weight is None:
            continue

        weight_local = weight.to_local() if isinstance(weight, DTensor) else weight  # type: ignore&#91;arg-type&#93;

        # 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(&quot;_fp4_weight&quot;, fp4_w, persistent=False)
        mod.register_buffer(&quot;_fp4_weight_scale&quot;, fp4_s, persistent=False)
        mod.register_buffer(
            &quot;_weight_global_sf&quot;,
            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(&quot;weight&quot;, 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()</code></pre></details>