Files

556 lines
20 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""LTX-2 NVFP4 quantization (FlashInfer-backed).
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).
Upstreamed from ``FastVideo-internal`` so consumers that load LTX-2
weights with NVFP4 quantization can drive the public package
end-to-end.
`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.
"""
from __future__ import annotations
import logging
from typing import Any
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from fastvideo.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
from fastvideo.models.utils import set_weight_attrs
logger = logging.getLogger(__name__)
def _require_flashinfer() -> tuple[Any, Any, Any]:
"""Lazy flashinfer import — raised at use time, not import time.
Returns the bound ``(SfLayout, mm_fp4, nvfp4_quantize)`` triple from
flashinfer. Raises ``ImportError`` with an actionable hint if the
package is not available.
"""
try:
from flashinfer import ( # type: ignore[import-not-found]
SfLayout, mm_fp4, nvfp4_quantize,
)
except ImportError as exc: # pragma: no cover - depends on host env
raise ImportError("NVFP4 quantization requires flashinfer. "
"Install with `pip install flashinfer-python`.") from exc
return SfLayout, mm_fp4, nvfp4_quantize
_LTX2_REFINE_ONLY_SUFFIXES = (
".audio_to_video_attn.to_q",
".video_to_audio_attn.to_k",
".video_to_audio_attn.to_v",
)
_LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES = (
"attn1.to_q",
"attn1.to_k",
"attn1.to_v",
"attn1.to_out",
"attn2.to_q",
"attn2.to_out",
"audio_to_video_attn.to_q",
"audio_to_video_attn.to_out",
"video_to_audio_attn.to_k",
"video_to_audio_attn.to_v",
"ffn.fc_in",
"ffn.fc_out",
)
_LTX2_NVFP4_LINEAR_PREFIXES = frozenset(f"ltx2.blocks.{block_idx}.{suffix}" for block_idx in range(48)
for suffix in _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES) | frozenset(
("ltx2.adaln_single.linear", ))
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
def _is_ltx2_refine_only_prefix(prefix: str) -> bool:
return any(prefix.endswith(suffix) for suffix in _LTX2_REFINE_ONLY_SUFFIXES)
def _get_ltx2_fp4_stage_profile(default: str = "refine") -> str:
"""Read the active stage profile from the forward context.
Streaming inference flips between ``base`` and ``refine`` between
segments; the FP4 layer set differs across the two. Falls back to
``default`` whenever the context is not available — this keeps the
op safe to run outside the streaming server (e.g. during eager
tests).
"""
try:
from fastvideo.forward_context import get_forward_context
forward_ctx = get_forward_context()
forward_batch = getattr(forward_ctx, "forward_batch", None)
if forward_batch is None:
return default
extra = getattr(forward_batch, "extra", None)
if not isinstance(extra, dict):
return default
profile = extra.get("ltx2_fp4_stage_profile", default)
if profile in ("base", "refine"):
return profile
return default
except Exception:
return default
_OPS_REGISTERED = False
def _register_ops_once() -> None:
"""Register the fastvideo_fp4 torch ops on first import that needs
them. Each op binds to flashinfer at call time; this just sets up
the dispatcher entries."""
global _OPS_REGISTERED
if _OPS_REGISTERED:
return
@torch.library.custom_op(
"fastvideo_fp4::nvfp4_quantize",
mutates_args=(),
device_types="cuda",
)
def _nvfp4_quantize_op(
x: torch.Tensor,
global_sf: torch.Tensor,
sf_layout: int,
do_shuffle: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
SfLayout, _, nvfp4_quantize = _require_flashinfer()
return nvfp4_quantize(x, global_sf, sfLayout=SfLayout(sf_layout), do_shuffle=do_shuffle)
@_nvfp4_quantize_op.register_fake
def _nvfp4_quantize_op_fake(
x: torch.Tensor,
global_sf: torch.Tensor,
sf_layout: int,
do_shuffle: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
del global_sf, sf_layout, do_shuffle
m, k = x.shape
quantized = torch.empty((m, (k + 1) // 2), device=x.device, dtype=torch.uint8)
scales = torch.empty((m, (k + 15) // 16), device=x.device, dtype=torch.uint8)
return quantized, scales
@torch.library.custom_op(
"fastvideo_fp4::mm_fp4",
mutates_args=(),
device_types="cuda",
)
def _mm_fp4_op(
a: torch.Tensor,
b: torch.Tensor,
a_scale: torch.Tensor,
b_scale: torch.Tensor,
alpha: torch.Tensor | None,
out_dtype: torch.dtype = torch.bfloat16,
out: torch.Tensor | None = None,
block_size: int = 16,
use_8x4_sf_layout: bool = False,
backend: str = "auto",
use_nvfp4: bool = True,
) -> torch.Tensor:
_, mm_fp4, _ = _require_flashinfer()
if a.dtype == torch.float4_e2m1fn_x2:
a = a.view(torch.uint8) if a.is_contiguous() else a.contiguous().view(torch.uint8)
if b.dtype == torch.float4_e2m1fn_x2:
b = b.view(torch.uint8) if b.is_contiguous() else b.contiguous().view(torch.uint8)
return mm_fp4(
a,
b,
a_scale,
b_scale,
alpha,
out_dtype,
out,
block_size=block_size,
use_8x4_sf_layout=use_8x4_sf_layout,
backend=backend,
use_nvfp4=use_nvfp4,
)
@_mm_fp4_op.register_fake
def _mm_fp4_op_fake(
a: torch.Tensor,
b: torch.Tensor,
a_scale: torch.Tensor,
b_scale: torch.Tensor,
alpha: torch.Tensor | None,
out_dtype: torch.dtype = torch.bfloat16,
out: torch.Tensor | None = None,
block_size: int = 16,
use_8x4_sf_layout: bool = False,
backend: str = "auto",
use_nvfp4: bool = True,
) -> torch.Tensor:
del a_scale, b_scale, alpha, block_size, use_8x4_sf_layout, backend
del use_nvfp4
if out is not None:
return out
out_shape = (*a.shape[:-1], b.shape[1])
return torch.empty(out_shape, device=a.device, dtype=out_dtype)
_OPS_REGISTERED = True
def _nvfp4_quantize(
x: torch.Tensor,
global_sf: Any,
*,
sfLayout: Any,
do_shuffle: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
_register_ops_once()
SfLayout, _, _ = _require_flashinfer()
if isinstance(sfLayout, SfLayout):
sf_layout = sfLayout.value
elif hasattr(sfLayout, "value"):
sf_layout = int(sfLayout.value)
else:
sf_layout = int(sfLayout)
if not torch.is_tensor(global_sf):
global_sf = torch.tensor(global_sf, device=x.device, dtype=torch.float32)
elif global_sf.device != x.device:
global_sf = global_sf.to(device=x.device)
if sf_layout == SfLayout.layout_linear.value:
x_for_quant = x
logical_rows = x.shape[0]
else:
# Sequence-parallel can feed either logical rows or row-padded
# rows. Normalize to the kernel tile shape for swizzled layouts
# so both paths share a stable quantization contract.
row_tile = 8 if sf_layout == SfLayout.layout_8x4.value else 128
logical_rows = x.shape[0]
pad_rows = (-logical_rows) % row_tile
x_for_quant = F.pad(x, (0, 0, 0, pad_rows))
quantized, scales = torch.ops.fastvideo_fp4.nvfp4_quantize(x_for_quant, global_sf, sf_layout, do_shuffle)
if sf_layout != SfLayout.layout_linear.value:
quantized = quantized.narrow(0, 0, logical_rows)
return quantized, scales
def _mm_fp4(
a: torch.Tensor,
b: torch.Tensor,
a_scale: torch.Tensor,
b_scale: torch.Tensor,
alpha: Any,
out_dtype: torch.dtype,
out: torch.Tensor | None,
**kwargs: Any,
) -> torch.Tensor:
_register_ops_once()
block_size = kwargs.pop("block_size", 16)
use_8x4_sf_layout = kwargs.pop("use_8x4_sf_layout", False)
backend = kwargs.pop("backend", "auto")
use_nvfp4 = kwargs.pop("use_nvfp4", True)
if kwargs:
raise TypeError(f"Unsupported kwargs for _mm_fp4: {sorted(kwargs)}")
if alpha is not None and not torch.is_tensor(alpha):
alpha = torch.tensor(alpha, device=a.device, dtype=torch.float32)
return torch.ops.fastvideo_fp4.mm_fp4(
a,
b,
a_scale,
b_scale,
alpha,
out_dtype,
out,
block_size,
use_8x4_sf_layout,
backend,
use_nvfp4,
)
def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor:
"""Coerce an activation to a dtype the FP4 linear accepts.
The pre-attention norm can emit fp32 (e.g. in eager mode, without the
torch.compile fusion that keeps it bf16). The FP4 linear emits bf16
regardless (see _mm_fp4 out dtype), so cast fp32 -> bf16 rather than
failing, matching the sibling fastvideo/layers/fp4linear.py. Non-floating
inputs (e.g. int/bool) are a genuine error and are rejected fast.
"""
if not x.is_floating_point():
raise TypeError(f"fp4 linear expects floating-point inputs, got {x.dtype}")
if x.dtype not in (torch.bfloat16, torch.float16):
x = x.to(torch.bfloat16)
return x
class NVFP4QuantizeMethod(QuantizeMethodBase):
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 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])
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='auto',
)
if bias is not None:
out = out + bias
out = out.view(*original_shape[:-1], out_dim)
return out
class NVFP4Config(QuantizationConfig):
"""LTX-2-specific NVFP4 quantization configuration.
NVFP4 is NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha,
``layout_128x4`` scale layout, group size 16). Today this class
hardcodes the LTX-2 layer paths it covers. When a second model
wants NVFP4, lift the layer-path list into a config field
instead of hardcoding it here.
"""
def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None):
super().__init__()
# ``base``: stage-1 set (no attn2.to_out, no cross-modal AV
# projections). ``refine``: full stage-2 set.
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"),
)
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
from fastvideo.layers.linear import LinearBase
# Use the superset at build/load time, then switch active subset
# dynamically in NVFP4QuantizeMethod.apply based on stage profile.
if isinstance(layer, LinearBase) and is_ltx2_nvfp4_linear_prefix(prefix):
method = NVFP4QuantizeMethod(layer_prefix=prefix)
method._retain_original_weights = self.retain_original_weights
return method
return None
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):
# ponytail: purging FSDP-sharded originals needs per-shard
# resharding bookkeeping; skip until a sharded deploy needs it.
retained += 1
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,
)
__all__ = [
"NVFP4Config",
"NVFP4QuantizeMethod",
"convert_model_to_nvfp4",
"is_ltx2_nvfp4_linear_prefix",
]