feat(ltx2): wire FP4 inference through fastvideo.layers.quantization

The public LTX-2 path was previously running full bf16 even when
callers set ``pipeline_config.dit_config.quant_config = FP4Config()``,
because:

* the DiT used plain ``torch.nn.Linear`` for FP4-eligible projections
  (``attn1/attn2.to_q/k/v/out``, ``audio_to_video_attn``,
  ``video_to_audio_attn``, FFN ``fc_in``/``fc_out``), so
  ``LinearBase.quant_method.apply`` was never reached;
* the loader did not call ``convert_model_to_fp4`` after weights
  loaded, so even matched layers had no ``_fp4_weight*`` buffers;
* ``QuantizationMethods`` did not register ``"FP4"``, so the typed
  ``engine.quantization.transformer_quant`` surface had no way to
  select it.

This commit completes the inference wire-up through the existing
``fastvideo.layers.quantization`` registry — no parallel pathway:

* swap ``nn.Linear`` → ``ReplicatedLinear`` for the LTX-2 attention
  ``to_q``/``to_k``/``to_v``/``to_out``/``to_gate_compress`` and the
  FFN ``fc_in``/``fc_out`` (``GELUApprox.proj`` and ``FeedForward``
  ``project_out``). Other linears (timestep MLP, caption proj,
  patchify, AdaLN scale/shift) stay ``nn.Linear`` to match internal.
* port ``_supports_prequantized_input`` and
  ``_linear_project_with_optional_prequant`` so the attention forward
  quantizes input once and reuses the ``(x_fp4, x_scale, x_global_sf)``
  tuple across q/k/v projections — matches internal exactly.
* plumb ``quant_config`` and ``prefix`` through
  ``BasicAVTransformerBlock`` → ``_init_transformer_blocks`` →
  ``LTXModel`` → ``LTX2Transformer3DModel`` so each
  ``ReplicatedLinear`` gets the correct
  ``ltx2.blocks.<i>.<sub>.<proj>`` prefix that
  ``FP4Config.get_quant_method`` matches against.
* register ``"FP4"`` in ``QuantizationMethods`` literal +
  ``get_quantization_config`` registry so typed
  ``engine.quantization.transformer_quant: "FP4"`` resolves to a
  concrete ``FP4Config()`` instance.
* add ``transformer_quant`` carrier on ``FastVideoArgs`` and
  ``__post_init__._apply_transformer_quant`` to pin the resolved
  quant config onto ``dit_config.quant_config`` (without
  overwriting an explicit setter on the dit_config).
* add ``_maybe_convert_model_to_fp4`` in ``fsdp_load.py`` that
  walks the loaded model once and registers
  ``_fp4_weight*``/``_fp4_alpha``/``_weight_global_sf`` buffers on
  layers whose ``quant_method`` is ``FP4QuantizeMethod``. flashinfer
  is imported lazily inside ``convert_model_to_fp4`` so this is a
  no-op on hosts without the FP4 kernels.
* restore ``LinearBase`` fallback to ``UnquantizedLinearMethod`` when
  ``quant_config.get_quant_method`` returns ``None`` — ``FP4Config``
  only tags a curated subset of layers, and the previous
  ``assert quant_method is not None`` would crash any non-tagged
  layer that received a quant_config.

Non-FP4 callers are unaffected: ``UnquantizedLinearMethod.apply``
runs the standard ``F.linear`` path, and parameter names
(``to_q.weight``, ``to_out.0.weight`` from ``nn.ModuleList``) are
identical to the previous ``nn.Sequential`` layout.

Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com>
Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com>
Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com>
Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>

docs/design/inference_schema_parity_inventory.yaml

@@ -37,6 +37,7 @@ surfaces:
       torch_compile_kwargs_text_encoder: generator.engine.compile.text_encoder_kwargs
       torch_compile_kwargs_vae: generator.engine.compile.vae_kwargs
       torch_compile_kwargs_audio_vae: generator.engine.compile.audio_vae_kwargs
+      transformer_quant: generator.engine.quantization.transformer_quant
       disable_autocast: generator.engine.disable_autocast
       enable_stage_verification: generator.engine.enable_stage_verification
       prompt_txt: request.inputs.prompt_path

fastvideo/api/compat.py

@@ -268,7 +268,15 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
     if quantization is not None and quantization.text_encoder_quant is not None:
         kwargs["override_text_encoder_quant"] = quantization.text_encoder_quant
     if quantization is not None and quantization.transformer_quant is not None:
-        kwargs["transformer_quant"] = quantization.transformer_quant
+        # Resolve the typed quant name to a concrete ``QuantizationConfig``
+        # instance and pin it on ``dit_config.quant_config``. The legacy
+        # path expected callers to do this themselves via
+        # ``pipeline_config.dit_config.quant_config = FP4Config()``; the
+        # typed surface accepts a string and does the wiring here so
+        # downstream code can rely on a single source of truth.
+        from fastvideo.layers.quantization import get_quantization_config
+        _resolved_quant_cls = get_quantization_config(quantization.transformer_quant)
+        kwargs["transformer_quant"] = _resolved_quant_cls()
 
     components = normalized.pipeline.components
     if components.pipeline_config_path is not None:

fastvideo/fastvideo_args.py

@@ -217,6 +217,15 @@ class FastVideoArgs:
 
     override_text_encoder_safetensors: str | None = None  # path to safetensors file for text encoder override
     override_text_encoder_quant: QuantizationMethods = None
+    # Typed transformer quantization carrier. The typed inference API
+    # accepts ``engine.quantization.transformer_quant: "FP4"`` and the
+    # compat layer resolves the name to a concrete ``QuantizationConfig``
+    # instance (e.g. ``FP4Config()``); ``__post_init__`` then pins it on
+    # ``pipeline_config.dit_config.quant_config`` so the loader can detect
+    # FP4 layers via the standard ``get_quant_method`` path. ``None``
+    # leaves whatever value the caller already set on ``dit_config``
+    # untouched.
+    transformer_quant: Any | None = None
 
     override_transformer_cls_name: str | None = None
     init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading
@@ -245,8 +254,29 @@ class FastVideoArgs:
                 raise
         self._apply_ltx2_vae_overrides()
         self._resolve_refine_args()
+        self._apply_transformer_quant()
         self.check_fastvideo_args()
 
+    def _apply_transformer_quant(self) -> None:
+        """Pin the typed ``transformer_quant`` instance onto ``dit_config``.
+
+        ``transformer_quant`` is populated by the typed compat layer when
+        a caller writes ``engine.quantization.transformer_quant: "FP4"``
+        in their config. We pin it here rather than at request time so
+        the model loader sees the quant_config when constructing the
+        DiT (linear layers attach their quant_method during ``__init__``).
+        """
+        if self.transformer_quant is None or self.pipeline_config is None:
+            return
+        dit_config = getattr(self.pipeline_config, "dit_config", None)
+        if dit_config is None:
+            return
+        # Don't overwrite if the caller already set it explicitly on
+        # dit_config (e.g. via ``pipeline_config.dit_config.quant_config = FP4Config()``);
+        # the explicit setter wins.
+        if getattr(dit_config, "quant_config", None) is None:
+            dit_config.quant_config = self.transformer_quant
+
     def _resolve_refine_args(self) -> None:
         """Map generic refine_* args to LTX-2-specific refine fields."""
         if self.refine_enabled is not None:

fastvideo/layers/linear.py

@@ -191,7 +191,15 @@ class LinearBase(torch.nn.Module):
         if quant_config is None:
             self.quant_method: QuantizeMethodBase | None = (UnquantizedLinearMethod())
         else:
+            # ``get_quant_method`` returns ``None`` for layers the config
+            # has decided not to quantize (e.g. ``FP4Config`` only tags
+            # a curated subset of LTX-2 attention/FFN layers). Fall back
+            # to ``UnquantizedLinearMethod`` so untagged layers behave
+            # like a plain ``nn.Linear`` instead of breaking subclass
+            # asserts.
             self.quant_method = quant_config.get_quant_method(self, prefix=prefix)
+            if self.quant_method is None:
+                self.quant_method = UnquantizedLinearMethod()
 
     def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
         raise NotImplementedError
@@ -230,8 +238,14 @@ class ReplicatedLinear(LinearBase):
             prefix=prefix,
         )
 
-        # All the linear layer supports quant method.
-        assert self.quant_method is not None
+        # ``QuantizationConfig.get_quant_method`` may return ``None`` for
+        # layers it doesn't intend to quantize (e.g. ``FP4Config`` only
+        # tags a specific subset of LTX-2 attention/FFN layers). Fall
+        # back to ``UnquantizedLinearMethod`` so non-matched layers
+        # behave like a plain ``nn.Linear``.
+        if self.quant_method is None:
+            self.quant_method = UnquantizedLinearMethod()
+
         self.quant_method.create_weights(
             self,
             self.input_size,

fastvideo/layers/quantization/__init__.py

@@ -2,7 +2,7 @@ from typing import Literal, get_args
 
 from fastvideo.layers.quantization.base_config import QuantizationConfig
 
-QuantizationMethods = Literal[None, "AbsMaxFP8"]
+QuantizationMethods = Literal[None, "AbsMaxFP8", "FP4"]
 
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
@@ -51,9 +51,11 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
 
     # lazy import to avoid triggering `torch.compile` too early
     from .absmax_fp8 import AbsMaxFP8Config
+    from .fp4_config import FP4Config
 
     method_to_config: dict[str, type[QuantizationConfig]] = {
         "AbsMaxFP8": AbsMaxFP8Config,
+        "FP4": FP4Config,
     }
     # Update the `method_to_config` with customized quantization methods.
     method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)

fastvideo/models/dits/ltx2.py

@@ -28,6 +28,8 @@ from fastvideo.distributed.communication_op import (
 )
 from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
 from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.quantization.base_config import QuantizationConfig
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum
@@ -35,6 +37,50 @@ from fastvideo.platforms import AttentionBackendEnum
 logger = init_logger(__name__)
 
 
+def _supports_prequantized_input(linear: ReplicatedLinear) -> bool:
+    """Whether ``linear``'s quant method is willing to accept a
+    pre-quantized ``(x_fp4, x_scale, x_global_sf)`` input tuple.
+
+    Used by the LTX-2 attention forward path so that a single input
+    tensor can be quantized once and reused across q/k/v projections
+    when they share the same source (self-attention).
+    """
+    quant_method = getattr(linear, "quant_method", None)
+    if not callable(getattr(quant_method, "quantize_input", None)):
+        return False
+    wants_prequant = getattr(quant_method, "wants_prequantized_input", None)
+    if callable(wants_prequant):
+        try:
+            return bool(wants_prequant())
+        except Exception:
+            return False
+    return True
+
+
+def _linear_project_with_optional_prequant(
+    linear: ReplicatedLinear,
+    x: torch.Tensor,
+    pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None,
+) -> torch.Tensor:
+    """Project ``x`` through ``linear``, optionally bypassing the
+    in-method quantize step when ``pre_quantized`` is supplied.
+
+    When the linear's quant method does not support pre-quantized
+    inputs (e.g. ``UnquantizedLinearMethod``), falls back to the
+    standard ``linear(x)`` call and discards the bias-pass-through
+    tuple element.
+    """
+    if pre_quantized is None or not _supports_prequantized_input(linear):
+        return linear(x)[0]
+    bias = linear.bias if not linear.skip_bias_add else None
+    return linear.quant_method.apply(  # type: ignore[union-attr]
+        linear,
+        x,
+        bias=bias,
+        pre_quantized=pre_quantized,
+    )
+
+
 def get_timestep_embedding(
     timesteps: torch.Tensor,
     embedding_dim: int,
@@ -172,25 +218,57 @@ class PixArtAlphaTextProjection(torch.nn.Module):
 
 class GELUApprox(nn.Module):
     """Linear + tanh-approximate GELU used by LTX-2 FFN."""
-    def __init__(self, in_features: int, out_features: int):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.proj = nn.Linear(in_features, out_features)
+        self.proj = ReplicatedLinear(
+            in_features,
+            out_features,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc_in",
+        )
         self.act = nn.GELU(approximate="tanh")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.act(self.proj(x))
+        return self.act(self.proj(x)[0])
 
 
 class FeedForward(nn.Module):
     """LTX-2 FFN: GELUApprox -> Identity -> Linear."""
-    def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None:
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int,
+        mult: int = 4,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
         super().__init__()
         inner_dim = int(dim * mult)
-        project_in = GELUApprox(dim, inner_dim)
-        self.net = nn.Sequential(project_in, nn.Identity(), nn.Linear(inner_dim, dim_out))
+        project_in = GELUApprox(
+            dim,
+            inner_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ffn",
+        )
+        project_out = ReplicatedLinear(
+            inner_dim,
+            dim_out,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ffn.fc_out",
+        )
+        self.net = nn.ModuleList([project_in, nn.Identity(), project_out])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.net(x)
+        x = self.net[0](x)
+        x = self.net[1](x)
+        x = self.net[2](x)[0]
+        return x
 
 
 class VideoLatentShape(tuple):
@@ -1246,6 +1324,8 @@ class LTXSelfAttention(nn.Module):
         norm_eps: float,
         rope_type: LTXRopeType,
         supported_attention_backends: tuple[AttentionBackendEnum, ...],
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         inner_dim = dim_head * heads
@@ -1257,10 +1337,37 @@ class LTXSelfAttention(nn.Module):
 
         self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
         self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=True)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity())
+        self.to_q = ReplicatedLinear(
+            query_dim,
+            inner_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q",
+        )
+        self.to_k = ReplicatedLinear(
+            context_dim,
+            inner_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k",
+        )
+        self.to_v = ReplicatedLinear(
+            context_dim,
+            inner_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v",
+        )
+        self.to_out = nn.ModuleList([
+            ReplicatedLinear(
+                inner_dim,
+                query_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_out",
+            ),
+            nn.Identity(),
+        ])
 
         self.attn = LTXLocalAttention(
             num_heads=heads,
@@ -1271,9 +1378,15 @@ class LTXSelfAttention(nn.Module):
             causal=False,
             supported_attention_backends=supported_attention_backends,
         )
-        self.to_gate_compress: nn.Linear | None = None
+        self.to_gate_compress: ReplicatedLinear | None = None
         if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
-            self.to_gate_compress = nn.Linear(context_dim, inner_dim, bias=True)
+            self.to_gate_compress = ReplicatedLinear(
+                context_dim,
+                inner_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_gate_compress",
+            )
         self.attn_masked = LTXLocalAttention(
             num_heads=heads,
             head_size=dim_head,
@@ -1292,11 +1405,24 @@ class LTXSelfAttention(nn.Module):
         pe: tuple[torch.Tensor, torch.Tensor] | None = None,
         k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
-        q = self.to_q(x)
         context = x if context is None else context
-        k = self.to_k(context)
-        v = self.to_v(context)
-        gate_compress = (self.to_gate_compress(context)
+
+        q_pre_quantized = (
+            self.to_q.quant_method.quantize_input(x)  # type: ignore[union-attr]
+            if _supports_prequantized_input(self.to_q) else None)
+        kv_pre_quantized = None
+        if (_supports_prequantized_input(self.to_k)
+                and _supports_prequantized_input(self.to_v)):
+            if context is x and q_pre_quantized is not None:
+                kv_pre_quantized = q_pre_quantized
+            else:
+                kv_pre_quantized = self.to_k.quant_method.quantize_input(  # type: ignore[union-attr]
+                    context)
+
+        q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized)
+        k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized)
+        v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized)
+        gate_compress = (self.to_gate_compress(context)[0]
                          if self.to_gate_compress is not None else None)
 
         q = self.q_norm(q)
@@ -1346,7 +1472,8 @@ class LTXSelfAttention(nn.Module):
                             ltx_freqs_cis=pe,
                             ltx_k_freqs_cis=k_pe)
         out = out.reshape(b, q_len, -1)
-        return self.to_out(out)
+        out = self.to_out[0](out)[0]
+        return self.to_out[1](out)
 
 
 class LTXDistributedSelfAttention(nn.Module):
@@ -1361,6 +1488,7 @@ class LTXDistributedSelfAttention(nn.Module):
         norm_eps: float,
         rope_type: LTXRopeType,
         supported_attention_backends: tuple[AttentionBackendEnum, ...],
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -1373,10 +1501,37 @@ class LTXDistributedSelfAttention(nn.Module):
 
         self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
         self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=True)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity())
+        self.to_q = ReplicatedLinear(
+            query_dim,
+            inner_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q",
+        )
+        self.to_k = ReplicatedLinear(
+            context_dim,
+            inner_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k",
+        )
+        self.to_v = ReplicatedLinear(
+            context_dim,
+            inner_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v",
+        )
+        self.to_out = nn.ModuleList([
+            ReplicatedLinear(
+                inner_dim,
+                query_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_out",
+            ),
+            nn.Identity(),
+        ])
 
         self.attn = LTXDistributedAttention(
             num_heads=heads,
@@ -1386,9 +1541,15 @@ class LTXDistributedSelfAttention(nn.Module):
             supported_attention_backends=supported_attention_backends,
             prefix=f"{prefix}.attn",
         )
-        self.to_gate_compress: nn.Linear | None = None
+        self.to_gate_compress: ReplicatedLinear | None = None
         if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
-            self.to_gate_compress = nn.Linear(context_dim, inner_dim, bias=True)
+            self.to_gate_compress = ReplicatedLinear(
+                context_dim,
+                inner_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_gate_compress",
+            )
 
     def forward(
         self,
@@ -1408,11 +1569,24 @@ class LTXDistributedSelfAttention(nn.Module):
             k_pe: Rotary position embeddings for K (cos, sin), or None to use pe
             original_seq_len: Original (unpadded) full sequence length
         """
-        q = self.to_q(x)
         context = x if context is None else context
-        k = self.to_k(context)
-        v = self.to_v(context)
-        gate_compress = (self.to_gate_compress(context)
+
+        q_pre_quantized = (
+            self.to_q.quant_method.quantize_input(x)  # type: ignore[union-attr]
+            if _supports_prequantized_input(self.to_q) else None)
+        kv_pre_quantized = None
+        if (_supports_prequantized_input(self.to_k)
+                and _supports_prequantized_input(self.to_v)):
+            if context is x and q_pre_quantized is not None:
+                kv_pre_quantized = q_pre_quantized
+            else:
+                kv_pre_quantized = self.to_k.quant_method.quantize_input(  # type: ignore[union-attr]
+                    context)
+
+        q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized)
+        k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized)
+        v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized)
+        gate_compress = (self.to_gate_compress(context)[0]
                          if self.to_gate_compress is not None else None)
 
         q = self.q_norm(q)
@@ -1443,7 +1617,8 @@ class LTXDistributedSelfAttention(nn.Module):
         )
 
         out = out.reshape(b, q_len, -1)
-        return self.to_out(out)
+        out = self.to_out[0](out)[0]
+        return self.to_out[1](out)
 
 
 class BasicAVTransformerBlock(torch.nn.Module):
@@ -1457,6 +1632,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
         norm_eps: float = 1e-6,
         use_distributed_attention: bool = False,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -1488,15 +1664,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 norm_eps=norm_eps,
                 rope_type=rope_type,
                 supported_attention_backends=video_self_attn_backends,
-                prefix=f"{prefix}.blocks.{idx}.attn1" if use_distributed_attention else "",
-            ) if use_distributed_attention else LTXSelfAttention(
-                query_dim=video.dim,
-                context_dim=None,
-                heads=video.heads,
-                dim_head=video.d_head,
-                norm_eps=norm_eps,
-                rope_type=rope_type,
-                supported_attention_backends=video_self_attn_backends,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.attn1",
             )
             # Text cross-attention - always local (text is replicated)
             self.attn2 = CrossAttnCls(
@@ -1507,8 +1676,15 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 norm_eps=norm_eps,
                 rope_type=rope_type,
                 supported_attention_backends=dense_attn_backends,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.attn2",
+            )
+            self.ff = FeedForward(
+                video.dim,
+                dim_out=video.dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}",
             )
-            self.ff = FeedForward(video.dim, dim_out=video.dim)
             self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))
 
         if audio is not None:
@@ -1521,15 +1697,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 norm_eps=norm_eps,
                 rope_type=rope_type,
                 supported_attention_backends=dense_attn_backends,
-                prefix=f"{prefix}.blocks.{idx}.audio_attn1" if use_distributed_attention else "",
-            ) if use_distributed_attention else LTXSelfAttention(
-                query_dim=audio.dim,
-                context_dim=None,
-                heads=audio.heads,
-                dim_head=audio.d_head,
-                norm_eps=norm_eps,
-                rope_type=rope_type,
-                supported_attention_backends=dense_attn_backends,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.audio_attn1",
             )
             # Text cross-attention - always local (text is replicated)
             self.audio_attn2 = CrossAttnCls(
@@ -1540,8 +1709,15 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 norm_eps=norm_eps,
                 rope_type=rope_type,
                 supported_attention_backends=dense_attn_backends,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.audio_attn2",
+            )
+            self.audio_ff = FeedForward(
+                audio.dim,
+                dim_out=audio.dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.audio",
             )
-            self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
             self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))
 
         if audio is not None and video is not None:
@@ -1555,6 +1731,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 norm_eps=norm_eps,
                 rope_type=rope_type,
                 supported_attention_backends=dense_attn_backends,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.audio_to_video_attn",
             )
             # Video-to-audio cross-attention
             # Uses local attention - context is gathered from all SP ranks in forward()
@@ -1566,6 +1744,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 norm_eps=norm_eps,
                 rope_type=rope_type,
                 supported_attention_backends=dense_attn_backends,
+                quant_config=quant_config,
+                prefix=f"{prefix}.blocks.{idx}.video_to_audio_attn",
             )
             self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
             self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))
@@ -1889,6 +2069,7 @@ class LTXModel(torch.nn.Module):
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
         double_precision_rope: bool = False,
         use_distributed_attention: bool = False,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
@@ -1942,6 +2123,7 @@ class LTXModel(torch.nn.Module):
             audio_cross_attention_dim=audio_cross_attention_dim,
             norm_eps=norm_eps,
             use_distributed_attention=use_distributed_attention,
+            quant_config=quant_config,
             prefix=prefix,
         )
 
@@ -2073,6 +2255,7 @@ class LTXModel(torch.nn.Module):
         audio_cross_attention_dim: int,
         norm_eps: float,
         use_distributed_attention: bool = False,
+        quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
         video_config = (
@@ -2105,6 +2288,7 @@ class LTXModel(torch.nn.Module):
                     rope_type=self.rope_type,
                     norm_eps=norm_eps,
                     use_distributed_attention=use_distributed_attention,
+                    quant_config=quant_config,
                     prefix=prefix,
                 )
                 for idx in range(num_layers)
@@ -2293,6 +2477,7 @@ class LTX2Transformer3DModel(BaseDiT):
             audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos,
             av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier,
             use_distributed_attention=use_distributed_attention,
+            quant_config=config.quant_config,
             prefix=config.prefix,
         )
 

fastvideo/models/loader/fsdp_load.py

@@ -28,6 +28,35 @@ from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
 logger = init_logger(__name__)
 
 
+def _maybe_convert_model_to_fp4(model: nn.Module) -> None:
+    """Quantize FP4-tagged linear layers in-place after weights are loaded.
+
+    Walks the module tree once, looking for layers whose ``quant_method``
+    is an :class:`FP4QuantizeMethod` (attached at construction time by
+    :meth:`FP4Config.get_quant_method`). When at least one such layer
+    exists, calls :func:`convert_model_to_fp4` to register the
+    ``_fp4_weight*`` / ``_fp4_alpha`` / ``_weight_global_sf`` buffers
+    on each targeted layer.
+
+    The walk is short-circuited on the first non-FP4 layer in the
+    ``UnquantizedLinearMethod`` case so non-FP4 callers pay only a
+    single ``isinstance`` check. flashinfer is imported lazily inside
+    :func:`convert_model_to_fp4` so this helper is a no-op on hosts
+    without the FP4 backend.
+    """
+    # Defer the import: fp4_config imports heavy diffusers/torch.distributed
+    # symbols at module-load time, and unconditional import would penalize
+    # every loader call regardless of whether FP4 is wired.
+    from fastvideo.layers.quantization.fp4_config import (FP4QuantizeMethod,
+                                                          convert_model_to_fp4)
+
+    for mod in model.modules():
+        if isinstance(getattr(mod, "quant_method", None), FP4QuantizeMethod):
+            logger.info("Converting loaded model weights for FP4 linear layers")
+            convert_model_to_fp4(model)
+            return
+
+
 # TODO(PY): move this to utils elsewhere
 @contextlib.contextmanager
 def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
@@ -158,6 +187,15 @@ def maybe_load_fsdp_model(
         if isinstance(p, torch.nn.Parameter):
             p.requires_grad = False
 
+    # FP4 weight prequantization. We detect by the registered
+    # ``quant_method`` on linear layers rather than by a separate flag —
+    # construction-time ``FP4Config.get_quant_method`` already attached
+    # ``FP4QuantizeMethod`` to every targeted layer, so the loader's
+    # responsibility is just to materialize the per-layer fp4 weight /
+    # scale buffers from the freshly-loaded bf16 weights. No-op when
+    # ``flashinfer`` is not installed (lazy import inside the helper).
+    _maybe_convert_model_to_fp4(model)
+
     compile_in_loader = enable_torch_compile and training_mode
     if compile_in_loader:
         compile_kwargs = torch_compile_kwargs or {}