[test] Run LTX-2 attention QAT nightly on GB200

Record 0.3661 from a single-GB200 forward/backward using the corrected 576-linear deployment target profile.

docs/training/train_infra.md

@@ -212,6 +212,25 @@ pipeline:
   flow_shift: 8
 ```
 
+Registered transformer linear-quantization configs can also be selected by
+name. For example, the LTX-2 NVFP4-QAT recipe applies real FP4 forward GEMMs
+with a straight-through-estimator backward to its deployment-targeted
+attention/FFN projections:
+
+```yaml
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train
+```
+
+The LTX-2 recipe in
+`examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml` combines that linear
+configuration with `models.student.attention_backend: ATTN_QAT_TRAIN` for
+video-attention forward/backward. On sm120, its validation callback temporarily
+switches those layers to `ATTN_QAT_INFER`.
+On GB200, set `callbacks.validation.attn_qat_infer: false` to keep validation on
+the train-time QAT backend; the inference kernel is sm120-only.
+
 ---
 
 ## Training Methods

examples/train/configs/example.yaml

@@ -195,6 +195,10 @@ callbacks:
 # ------------------------------------------------------------------------------
 pipeline:
   flow_shift: 3                       # default: null (model-specific)
+  # LTX-2 linear QAT (SM100+ with FlashInfer); distinct from the role-local
+  # attention_backend option above, which selects attention math for Wan.
+  # dit_config:
+  #   quant_config: nvfp4_qat_train
   # flow_shift_sr: null               # default: null (super-resolution shift)
   # embedded_cfg_scale: 6.0           # default: 6.0
   # is_causal: false                  # default: false

examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml

@@ -0,0 +1,111 @@
+# LTX-2 T2V overfitting test config with NVFP4 quantization-aware
+# training (QAT).
+#
+# Same recipe as overfit_ltx2_t2v.yaml, but the curated attention/FFN
+# linear set used by LTX-2 NVFP4 deployment runs its forward through
+# real NVFP4 GEMMs (flashinfer cutlass) with a straight-through-estimator
+# backward into the dense fp32 master weights. No extra parameters or
+# buffers are added, so FSDP sharding, checkpointing, and export are
+# identical to dense training. Video attention uses ATTN_QAT_TRAIN for its
+# quantized forward and STE backward, then ATTN_QAT_INFER during validation.
+# Head-dim-64 audio attention and masked text attention remain dense.
+#
+# Validation requires an sm_120 GPU with the attn_qat_infer extension.
+#
+# Preprocess data first (same data as the bf16 overfit):
+#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+#
+# Run on multi-GPU sm_120 hardware:
+#   NUM_GPUS=4 \
+#     bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml
+#
+# GB200 can train and validate with ATTN_QAT_TRAIN, but cannot load the
+# sm_120-only inference kernel. Disable only the validation-time swap:
+#   NUM_GPUS=4 \
+#     bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml \
+#       --callbacks.validation.attn_qat_infer false
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: ATTN_QAT_TRAIN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/ltx2_overfit_preprocessed
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    # LTX2Model requires 0.0: CFG dropout would zero post-connector
+    # embeddings, which is not the model's unconditional input.
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 11  # (81 - 1) / 8 + 1
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/ltx2_overfit_nvfp4_attn_qat
+    # A full training-state checkpoint is ~150GB for the 13B trainable
+    # video branch; disable saves for the overfit smoke run.
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 1
+
+  tracker:
+    trackers: [wandb]
+    project_name: fastvideo_ltx2
+    run_name: ltx2_overfit_nvfp4_attn_qat
+
+  model:
+    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
+    # so the default noise-minus-clean target reproduces the official
+    # unweighted masked-MSE (mask is all-ones for plain T2V).
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
+    dataset_file: data/ltx2_overfit_preprocessed/validation_prompts.json
+    every_steps: 50
+    sampling_steps: [8]
+    guidance_scale: 1.0
+    num_frames: 81
+    # Validation reuses the live transformer and temporarily replaces its
+    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
+    attn_qat_infer: true
+
+# quant_config selects NVFP4 QAT for the same curated linear prefixes as
+# LTX-2 NVFP4 deployment. They fake-quantize through real FP4 GEMMs with
+# an STE backward. The string resolves to NVFP4QATTrainConfig at parse.
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train

fastvideo/attention/backends/attn_qat_infer.py

@@ -51,7 +51,8 @@ def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
 
 
 def is_attn_qat_infer_available() -> bool:
-    return _get_attn_qat_infer() is not None
+    return (torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0)
+            and _get_attn_qat_infer() is not None)
 
 
 class AttnQatInferBackend(AttentionBackend):

fastvideo/attention/selector.py

@@ -150,14 +150,15 @@ def _cached_get_attn_backend(
     from fastvideo.platforms import current_platform
 
     if (selected_backend is not None and selected_backend not in supported_attention_backends):
+        fallback_backend = (default_backend if default_backend in supported_attention_backends else None)
         logger.warning(
             "Requested attention backend %s is not supported by this "
-            "layer; supported backends are %s. Falling back to automatic "
-            "selection.",
+            "layer; supported backends are %s. Falling back to %s.",
             selected_backend.name,
             [b.name for b in supported_attention_backends],
+            fallback_backend.name if fallback_backend is not None else "automatic selection",
         )
-        selected_backend = None
+        selected_backend = fallback_backend
     attention_cls = current_platform.get_attn_backend_cls(selected_backend, head_size, dtype)
     if not attention_cls:
         raise ValueError(f"Invalid attention backend for {current_platform.device_name}")

fastvideo/layers/quantization/nvfp4_config.py

@@ -58,6 +58,29 @@ _LTX2_REFINE_ONLY_SUFFIXES = (
     ".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)
@@ -422,24 +445,7 @@ class NVFP4Config(QuantizationConfig):
 
         # Use the superset at build/load time, then switch active subset
         # dynamically in NVFP4QuantizeMethod.apply based on stage profile.
-        fp4_layers = [[
-            f"ltx2.blocks.{i}.attn1.to_q",
-            f"ltx2.blocks.{i}.attn1.to_k",
-            f"ltx2.blocks.{i}.attn1.to_v",
-            f"ltx2.blocks.{i}.attn1.to_out",
-            f"ltx2.blocks.{i}.attn2.to_q",
-            f"ltx2.blocks.{i}.attn2.to_out",
-            f"ltx2.blocks.{i}.audio_to_video_attn.to_q",
-            f"ltx2.blocks.{i}.audio_to_video_attn.to_out",
-            f"ltx2.blocks.{i}.video_to_audio_attn.to_k",
-            f"ltx2.blocks.{i}.video_to_audio_attn.to_v",
-            f"ltx2.blocks.{i}.ffn.fc_in",
-            f"ltx2.blocks.{i}.ffn.fc_out",
-        ] for i in range(48)]
-        fp4_layers.append([
-            "ltx2.adaln_single.linear",
-        ])
-        if isinstance(layer, LinearBase) and any(prefix in layer_names for layer_names in fp4_layers):
+        if isinstance(layer, LinearBase) and is_ltx2_nvfp4_linear_prefix(prefix):
             return NVFP4QuantizeMethod(layer_prefix=prefix)
         return None
 
@@ -485,4 +491,5 @@ __all__ = [
     "NVFP4Config",
     "NVFP4QuantizeMethod",
     "convert_model_to_nvfp4",
+    "is_ltx2_nvfp4_linear_prefix",
 ]

fastvideo/layers/quantization/nvfp4_qat_train_config.py

@@ -9,10 +9,11 @@ straight-through estimator), so the model learns to absorb FP4 linear error.
 
 That STE already exists in ``fastvideo.layers.fp4linear._LinearFWD4BWD16Fn``
 (FP4 forward, full-precision backward) but is otherwise unwired. This method
-bridges it into the standard ``quant_config`` path, so it activates via
-``transformer_quant="nvfp4_qat_train"`` on the same Wan-2.1 layers as nvfp4_qat
-(to_q/k/v/out + ffn). No ``convert_model_to_fp4`` is needed: the weight is kept
-in full precision and quantized on the fly each step.
+bridges it into the standard ``quant_config`` path. Wan-style prefixes retain
+the broad attention/FFN projection profile used by ``nvfp4_qat``; LTX-2
+prefixes use the exact curated target set from its ``NVFP4`` deployment config.
+No ``convert_model_to_fp4`` is needed: the weight is kept in full precision and
+quantized on the fly each step.
 """
 import logging
 from typing import Any
@@ -21,6 +22,8 @@ import torch
 from torch.nn.parameter import Parameter
 
 from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
+from fastvideo.layers.quantization.nvfp4_config import is_ltx2_nvfp4_linear_prefix
+from fastvideo.layers.quantization.nvfp4_qat_config import DEFAULT_FP4_LAYERS
 from fastvideo.models.utils import set_weight_attrs
 
 logger = logging.getLogger(__name__)
@@ -72,7 +75,10 @@ class NVFP4QATTrainConfig(QuantizationConfig):
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
         from fastvideo.layers.linear import LinearBase
-        fp4_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"]
-        if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp4_layers):
+        if not isinstance(layer, LinearBase):
+            return None
+        matches = (is_ltx2_nvfp4_linear_prefix(prefix) if prefix.startswith("ltx2.") else any(
+            layer_name in prefix for layer_name in DEFAULT_FP4_LAYERS))
+        if matches:
             return NVFP4QATTrainQuantizeMethod()
         return None

fastvideo/models/dits/ltx2.py

@@ -1571,6 +1571,7 @@ class LTXSelfAttention(nn.Module):
             softmax_scale=None,
             causal=False,
             supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,),
+            default_backend=AttentionBackendEnum.TORCH_SDPA,
         )
 
     def forward(
@@ -1853,12 +1854,15 @@ class BasicAVTransformerBlock(torch.nn.Module):
         # Text cross-attention uses LocalAttention (text embeddings are replicated)
         SelfAttnCls = LTXDistributedSelfAttention if use_distributed_attention else LTXSelfAttention
         CrossAttnCls = LTXSelfAttention  # Text cross-attention is always local
-        video_self_attn_backends = (
+        video_attn_backends = (
             AttentionBackendEnum.FLASH_ATTN,
             AttentionBackendEnum.TORCH_SDPA,
-            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.ATTN_QAT_INFER,
+            AttentionBackendEnum.ATTN_QAT_TRAIN,
         )
-        dense_attn_backends = (
+        video_self_attn_backends = video_attn_backends + (
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN, )
+        audio_attn_backends = (
             AttentionBackendEnum.FLASH_ATTN,
             AttentionBackendEnum.TORCH_SDPA,
         )
@@ -1885,7 +1889,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=video.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=dense_attn_backends,
+                supported_attention_backends=video_attn_backends,
                 apply_gated_attention=video.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.attn2",
@@ -1911,7 +1915,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=dense_attn_backends,
+                supported_attention_backends=audio_attn_backends,
                 apply_gated_attention=audio.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio_attn1",
@@ -1924,7 +1928,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=dense_attn_backends,
+                supported_attention_backends=audio_attn_backends,
                 apply_gated_attention=audio.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio_attn2",
@@ -1949,7 +1953,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=dense_attn_backends,
+                supported_attention_backends=audio_attn_backends,
                 apply_gated_attention=video.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio_to_video_attn",
@@ -1963,7 +1967,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 norm_eps=norm_eps,
                 rope_type=rope_type,
-                supported_attention_backends=dense_attn_backends,
+                supported_attention_backends=audio_attn_backends,
                 apply_gated_attention=audio.apply_gated_attention,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.video_to_audio_attn",

fastvideo/tests/attention/test_selector_role_override.py

@@ -54,6 +54,24 @@ def test_role_override_invalidates_cached_backend(monkeypatch) -> None:
         selector.global_force_attn_backend(None)
 
 
+def test_unsupported_role_override_honors_layer_default(monkeypatch) -> None:
+    monkeypatch.setattr(platforms, "_current_platform", _FakePlatform())
+    monkeypatch.setattr(selector, "resolve_obj_by_qualname", lambda name: name)
+    selector.global_force_attn_backend(None)
+
+    try:
+        with selector.global_force_attn_backend_context_manager(
+                AttentionBackendEnum.ATTN_QAT_TRAIN):
+            assert selector.get_attn_backend(
+                head_size=128,
+                dtype=torch.bfloat16,
+                supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA, ),
+                default_backend=AttentionBackendEnum.TORCH_SDPA,
+            ) == "TORCH_SDPA"
+    finally:
+        selector.global_force_attn_backend(None)
+
+
 def test_explicit_backend_config_rejects_typos() -> None:
     try:
         selector.coerce_attn_backend("attn_qat_typo")

fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py

@@ -1,17 +1,18 @@
 """E2E overfit tests for LTX-2 fine-tuning on the modular trainer (fastvideo/train).
 
-Parametrized over LTX-2.0 and LTX-2.3 distilled checkpoints. Each case
-downloads the single-sample cats dataset, preprocesses it with the
-matching LTX-2 VAE + Gemma text encoder into the t2v parquet format,
-trains LTX2Model with FineTuneMethod for 300 steps on 4 GPUs, and
-checks that the final validation video reproduces the training clip
-(MS-SSIM against the preprocessed ground-truth clip) better than the
-step-0 baseline.
+Parametrized over dense LTX-2.0 and LTX-2.3 distilled checkpoints plus
+an LTX-2.0 NVFP4 linear-QAT case. Each case downloads the single-sample
+cats dataset, preprocesses it with the matching LTX-2 VAE + Gemma text
+encoder into the t2v parquet format, trains LTX2Model with FineTuneMethod
+for 300 steps on 4 GPUs, and checks that the final validation video
+reproduces the training clip (MS-SSIM against the preprocessed ground-truth
+clip) better than the step-0 baseline.
 
 GPU assumptions: 4 x large-memory GPUs (developed on 4x GB200 192GB;
 the 18.9B-param DiT trains FSDP-sharded with fp32 master weights).
 Requires the FastVideo LTX-2 distilled checkpoints (set HF_HOME to a
 cache that contains them, or allow ~60GB of downloads per version).
+The NVFP4-QAT case additionally requires FlashInfer FP4 support and SM100+.
 
 Guarded by FASTVIDEO_NIGHTLY=1 so the default test suite stays fast.
 """
@@ -25,6 +26,7 @@ import sys
 from pathlib import Path
 
 import pytest
+import torch
 
 from huggingface_hub import snapshot_download
 
@@ -52,6 +54,16 @@ _CASES = {
         "prep_dir": Path(DATA_DIR) / "ltx2_3_overfit_preprocessed",
         "out_dir": Path(DATA_DIR) / "outputs_ltx2_3_overfit",
     },
+    # NVFP4 quantization-aware training: same LTX-2.0 data, attention/FFN
+    # linears forward through real FP4 GEMMs with an STE backward.
+    # Requires flashinfer FP4 kernels (sm_100-class GPU).
+    "ltx2_nvfp4_qat": {
+        "model": "FastVideo/LTX2-Distilled-Diffusers",
+        "config": "examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml",
+        "prep_dir": Path(DATA_DIR) / "ltx2_overfit_preprocessed",
+        "out_dir": Path(DATA_DIR) / "outputs_ltx2_nvfp4_qat_overfit",
+        "requires_nvfp4": True,
+    },
 }
 
 
@@ -93,10 +105,26 @@ def run_training(case: dict):
         "--callbacks.validation.dataset_file",
         str(case["prep_dir"] / "validation_prompts.json"),
     ]
+    if (case.get("requires_nvfp4")
+            and torch.cuda.get_device_capability(0) != (12, 0)):
+        cmd.extend(["--callbacks.validation.attn_qat_infer", "false"])
     print(f"Running training: {cmd}")
     subprocess.run(cmd, check=True, env=env)
 
 
+def test_run_training_disables_sm120_only_validation_on_gb200(monkeypatch):
+    commands = []
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (10, 0))
+    monkeypatch.setattr(subprocess, "run",
+                        lambda cmd, **_: commands.append(cmd))
+
+    run_training(_CASES["ltx2_nvfp4_qat"])
+
+    assert commands[0][-2:] == [
+        "--callbacks.validation.attn_qat_infer", "false"
+    ]
+
+
 def _validation_videos_by_step(out_dir: Path) -> dict[int, str]:
     pattern = os.path.join(str(out_dir), "validation_step_*_video_0.mp4")
     videos: dict[int, str] = {}
@@ -114,6 +142,18 @@ def _validation_videos_by_step(out_dir: Path) -> dict[int, str]:
 @pytest.mark.parametrize("case_id", _CASES.keys())
 def test_e2e_ltx2_overfit_new_stack(case_id: str):
     case = _CASES[case_id]
+    if case.get("requires_nvfp4"):
+        if (not torch.cuda.is_available() or
+                torch.cuda.get_device_capability(0) < (10, 0)):
+            pytest.skip("NVFP4 QAT requires an SM100+ GPU")
+        try:
+            from flashinfer import (  # noqa: F401
+                SfLayout,
+                mm_fp4,
+                nvfp4_quantize,
+            )
+        except ImportError:
+            pytest.skip("NVFP4 QAT requires flashinfer with FP4 kernels")
     download_data()
     run_preprocessing(case)
     run_training(case)

fastvideo/tests/ops/quantization/test_nvfp4_ltx2_wiring.py

@@ -18,6 +18,10 @@ from fastvideo.layers.quantization.nvfp4_config import (
     NVFP4Config,
     NVFP4QuantizeMethod,
 )
+from fastvideo.layers.quantization.nvfp4_qat_train_config import (
+    NVFP4QATTrainConfig,
+    NVFP4QATTrainQuantizeMethod,
+)
 from fastvideo.models.dits.ltx2 import (
     BasicAVTransformerBlock,
     FeedForward,
@@ -183,3 +187,46 @@ def test_basic_av_block_propagates_quant_config_to_all_children() -> None:
     for linear in unquantized_projections:
         assert isinstance(linear, ReplicatedLinear)
         assert isinstance(linear.quant_method, UnquantizedLinearMethod)
+
+
+def test_qat_train_uses_the_ltx2_deployment_target_profile() -> None:
+    kwargs = {
+        "idx": 3,
+        "video": TransformerConfig(dim=64, heads=2, d_head=32, context_dim=128),
+        "audio": TransformerConfig(dim=64, heads=2, d_head=32, context_dim=128),
+        "rope_type": LTXRopeType.INTERLEAVED,
+        "norm_eps": 1e-6,
+        "use_distributed_attention": False,
+        "prefix": "ltx2",
+    }
+    deployed = BasicAVTransformerBlock(
+        **kwargs,
+        quant_config=NVFP4Config(),
+    )
+    qat = BasicAVTransformerBlock(
+        **kwargs,
+        quant_config=NVFP4QATTrainConfig(),
+    )
+
+    deployed_names = {
+        name for name, module in deployed.named_modules()
+        if isinstance(getattr(module, "quant_method", None),
+                      NVFP4QuantizeMethod)
+    }
+    qat_names = {
+        name for name, module in qat.named_modules()
+        if isinstance(getattr(module, "quant_method", None),
+                      NVFP4QATTrainQuantizeMethod)
+    }
+    assert qat_names == deployed_names
+    assert len(qat_names) == 12
+
+
+def test_qat_train_preserves_wan_style_targeting() -> None:
+    linear = ReplicatedLinear(
+        16,
+        16,
+        quant_config=NVFP4QATTrainConfig(),
+        prefix="blocks.0.attn1.to_k",
+    )
+    assert isinstance(linear.quant_method, NVFP4QATTrainQuantizeMethod)

fastvideo/tests/train/callbacks/test_validation.py

@@ -82,6 +82,7 @@ class TestConstructor:
         assert cb.overlay_actions is False
         assert cb.offload_training_state is False
         assert cb.unload_pipeline_after_validation is False
+        assert cb.attn_qat_infer is False
         # Lazy fields not yet populated.
         assert cb._pipeline is None
         assert cb._sampling_param is None
@@ -102,6 +103,7 @@ class TestConstructor:
             overlay_actions=1,  # type: ignore[arg-type]
             offload_training_state="1",  # type: ignore[arg-type]
             unload_pipeline_after_validation="false",  # type: ignore[arg-type]
+            attn_qat_infer="true",  # type: ignore[arg-type]
         )
         assert cb.every_steps == 50
         assert cb.sampling_steps == [20, 40]
@@ -111,6 +113,7 @@ class TestConstructor:
         assert cb.overlay_actions is True
         assert cb.offload_training_state is True
         assert cb.unload_pipeline_after_validation is False
+        assert cb.attn_qat_infer is True
 
     def test_pipeline_kwargs_collected(self) -> None:
         cb = ValidationCallback(
@@ -219,6 +222,57 @@ class TestOnValidationBegin:
         assert cb.run_calls == [0]
 
 
+class TestAttnQatInferValidation:
+
+    def test_context_swaps_and_restores_qat_attention(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        from fastvideo.attention.backends import attn_qat_infer
+        from fastvideo.attention.backends.attn_qat_infer import (
+            AttnQatInferImpl, )
+        from fastvideo.attention.backends.attn_qat_train import (
+            AttnQatTrainImpl, )
+        from fastvideo.platforms import AttentionBackendEnum
+
+        monkeypatch.setattr(attn_qat_infer, "is_attn_qat_infer_available", lambda: True)
+        impl = AttnQatTrainImpl(
+            num_heads=1,
+            head_size=128,
+            causal=False,
+            softmax_scale=128**-0.5,
+        )
+
+        holder = torch.nn.Module()
+        holder.attn_impl = impl
+        holder.num_heads = 1
+        holder.head_size = 128
+        holder.num_kv_heads = 1
+        holder.softmax_scale = 128**-0.5
+        holder.backend = AttentionBackendEnum.ATTN_QAT_TRAIN
+        transformer = torch.nn.Module()
+        transformer.add_module("attention", holder)
+        cb = _make_callback()
+        cb.attn_qat_infer = True
+
+        with pytest.raises(RuntimeError, match="stop"):
+            with cb._attn_qat_infer_context(transformer):
+                assert isinstance(holder.attn_impl, AttnQatInferImpl)
+                assert holder.backend is AttentionBackendEnum.ATTN_QAT_INFER
+                raise RuntimeError("stop")
+
+        assert holder.attn_impl is impl
+        assert holder.backend is AttentionBackendEnum.ATTN_QAT_TRAIN
+
+    def test_context_requires_qat_training_attention(self) -> None:
+        cb = _make_callback()
+        cb.attn_qat_infer = True
+
+        with pytest.raises(RuntimeError, match="no ATTN_QAT_TRAIN layers"):
+            with cb._attn_qat_infer_context(torch.nn.Linear(1, 1)):
+                pass
+
+
 # ---------------------------------------------------------------------------
 # C. _find_ema_callback
 # ---------------------------------------------------------------------------

fastvideo/tests/train/fixtures/ltx2_t2v_qat_finetune_min.yaml

@@ -0,0 +1,51 @@
+# Minimum config to run a single NVFP4-QAT training step of
+# FineTuneMethod on LTX2Model for the per-method smoke test. Identical
+# to ltx2_t2v_finetune_min.yaml except pipeline.dit_config.quant_config
+# selects the LTX-2 deployment-matched NVFP4 QAT-train profile (STE
+# through real FP4 GEMMs). Requires flashinfer FP4 kernels (sm_100-class GPU); the test
+# additionally skips below 60GB GPU memory.
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+    attention_backend: ATTN_QAT_TRAIN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 3
+    num_height: 256
+    num_width: 256
+    num_frames: 17
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+  model:
+    precondition_outputs: false
+
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train

fastvideo/tests/train/methods/grad_norm_refs.json

@@ -9,6 +9,9 @@
   "test_ltx2_finetune": {
     "GB200": 0.3127
   },
+  "test_ltx2_nvfp4_qat_finetune": {
+    "GB200": 0.2421
+  },
   "test_matrixgame2_finetune": {
     "H200": 4.8247
   },

fastvideo/tests/train/methods/test_ltx2_finetune.py

@@ -2,12 +2,12 @@
 """Per-method GPU smoke test: ``LTX2Model`` + ``FineTuneMethod``.
 
 Mirrors ``test_wan_finetune.py`` for the LTX-2 plugin, parametrized
-over LTX-2.0 and LTX-2.3 checkpoints. LTX2-specific differences: the
-synthetic ``raw_batch`` carries a single post-connector Gemma
-embedding (3840-d for 2.0, 4096-d for 2.3 which has no in-DiT caption
-projection) and 128-channel VAE latents; the DiT is 18.9B params, so
-the test skips on GPUs with less than 60GB memory (e.g. the L40S CI
-runner).
+over dense LTX-2.0/LTX-2.3 and LTX-2.0 NVFP4-QAT checkpoints.
+LTX2-specific differences: the synthetic ``raw_batch`` carries a
+single post-connector Gemma embedding (3840-d for 2.0, 4096-d for 2.3
+which has no in-DiT caption projection) and 128-channel VAE latents;
+the DiT is 18.9B params, so the test skips on GPUs with less than 60GB
+memory (e.g. the L40S CI runner).
 """
 
 from __future__ import annotations
@@ -34,10 +34,17 @@ from .grad_norm_regression import (
 
 _FIXTURE_DIR = Path(__file__).resolve().parent.parent / "fixtures"
 
-# id -> (fixture, post-connector text embedding width, grad-norm ref name)
+# id -> (fixture, post-connector text embedding width, grad-norm ref name,
+#        expected NVFP4-QAT attachment count or None for dense training)
 _CASES = {
-    "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840, "test_ltx2_finetune"),
-    "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096, "test_ltx2_3_finetune"),
+    "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840, "test_ltx2_finetune", None),
+    "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096, "test_ltx2_3_finetune", None),
+    "ltx2_nvfp4_qat": (
+        "ltx2_t2v_qat_finetune_min.yaml",
+        3840,
+        "test_ltx2_nvfp4_qat_finetune",
+        576,
+    ),
 }
 
 _LTX2_TEXT_LEN = 1024
@@ -52,6 +59,18 @@ def _gpu_too_small() -> bool:
     return total < _MIN_GPU_MEMORY_GB * 1024**3
 
 
+def _flashinfer_fp4_available() -> bool:
+    try:
+        from flashinfer import (  # noqa: F401
+            SfLayout,
+            mm_fp4,
+            nvfp4_quantize,
+        )
+    except ImportError:
+        return False
+    return True
+
+
 def _build_synthetic_batch(
     device: torch.device,
     dtype: torch.dtype,
@@ -85,8 +104,21 @@ def test_ltx2_finetune_single_train_step(
         pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB "
                     "memory (LTX-2 DiT is 18.9B params)")
 
-    fixture_name, text_dim, ref_name = _CASES[case]
+    fixture_name, text_dim, ref_name, expected_qat_linears = _CASES[case]
+    if expected_qat_linears is not None:
+        if torch.cuda.get_device_capability(0) < (10, 0):
+            pytest.skip("NVFP4 QAT requires an SM100+ GPU")
+        if not _flashinfer_fp4_available():
+            pytest.skip("requires flashinfer with FP4 kernels")
+
     cfg = load_run_config(str(_FIXTURE_DIR / fixture_name))
+    if expected_qat_linears is not None:
+        from fastvideo.layers.quantization.nvfp4_qat_train_config import (
+            NVFP4QATTrainConfig, )
+        assert isinstance(
+            cfg.training.pipeline_config.dit_config.quant_config,
+            NVFP4QATTrainConfig,
+        )
 
     device = torch.device("cuda:0")
     dtype = torch.bfloat16
@@ -105,9 +137,30 @@ def test_ltx2_finetune_single_train_step(
         init_from=cfg.models["student"]["init_from"],
         training_config=cfg.training,
         trainable=True,
+        attention_backend=cfg.models["student"].get("attention_backend"),
     )
     model.transformer = model.transformer.to(device=device, dtype=dtype)
 
+    if expected_qat_linears is not None:
+        from fastvideo.layers.quantization.nvfp4_qat_train_config import (
+            NVFP4QATTrainQuantizeMethod, )
+        quantized = sum(
+            isinstance(getattr(module, "quant_method", None),
+                       NVFP4QATTrainQuantizeMethod)
+            for module in model.transformer.modules())
+        assert quantized == expected_qat_linears, (
+            f"expected {expected_qat_linears} deployment-matched NVFP4-QAT "
+            f"linears on the LTX-2 DiT, found {quantized}")
+
+        from fastvideo.attention.backends.attn_qat_train import (
+            AttnQatTrainImpl, )
+        qat_attention = sum(
+            isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl)
+            for module in model.transformer.modules())
+        assert qat_attention == 96, (
+            "expected ATTN_QAT_TRAIN on attn1 and attn2 in all 48 LTX-2 "
+            f"blocks, found {qat_attention} implementations")
+
     method = FineTuneMethod(
         cfg=cfg,
         role_models={"student": model},

fastvideo/tests/train/models/test_load_ltx2.py

@@ -34,6 +34,7 @@ from fastvideo.forward_context import (
     set_forward_context,
 )
 from fastvideo.pipelines import ForwardBatch
+from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.train.models.ltx2 import LTX2Model
 from fastvideo.train.utils.config import load_run_config
 
@@ -109,6 +110,31 @@ def test_ltx2_rejects_audio_training_before_loading(
         )
 
 
+def test_ltx2_forwards_role_attention_backend_to_loader(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    cfg = load_run_config(str(_FIXTURE_DIR / _CASES["ltx2"][0]))
+    transformer = _TinyLTX2Transformer(is_ltx2_3=False)
+    captured: list[AttentionBackendEnum | None] = []
+
+    def fake_load_module_from_path(**kwargs):
+        captured.append(kwargs.get("attention_backend"))
+        return transformer
+
+    monkeypatch.setattr(
+        "fastvideo.train.models.ltx2.ltx2.load_module_from_path",
+        fake_load_module_from_path,
+    )
+    model = LTX2Model(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=False,
+        attention_backend="TORCH_SDPA",
+    )
+
+    assert model.attention_backend is AttentionBackendEnum.TORCH_SDPA
+    assert captured == [AttentionBackendEnum.TORCH_SDPA]
+
+
 @pytest.mark.parametrize("case", _CASES.keys())
 def test_ltx2_wrapper_contract_runs_on_cpu(
         case: str, monkeypatch: pytest.MonkeyPatch) -> None:

fastvideo/tests/train/utils/test_config.py

@@ -245,6 +245,46 @@ def test_betas_parses_list_and_string_forms(
     assert cfg.training.optimizer.betas == expected
 
 
+def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> None:
+    from fastvideo.layers.quantization.nvfp4_qat_train_config import (
+        NVFP4QATTrainConfig, )
+    from fastvideo.train.entrypoint.dcp_to_diffusers import (
+        _run_config_from_raw, )
+
+    data = _minimal_yaml()
+    data["models"]["student"]["init_from"] = (
+        "FastVideo/LTX2-Distilled-Diffusers")
+    data["pipeline"] = {
+        "dit_config": {
+            "quant_config": "nvfp4_qat_train"
+        }
+    }
+
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+    assert isinstance(cfg.training.pipeline_config.dit_config.quant_config,
+                      NVFP4QATTrainConfig)
+
+    export_cfg = _run_config_from_raw(cfg.raw)
+    assert isinstance(
+        export_cfg.training.pipeline_config.dit_config.quant_config,
+        NVFP4QATTrainConfig,
+    )
+
+
+def test_pipeline_quant_config_rejects_unknown_name(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["models"]["student"]["init_from"] = (
+        "FastVideo/LTX2-Distilled-Diffusers")
+    data["pipeline"] = {
+        "dit_config": {
+            "quant_config": "not_a_quantization_method"
+        }
+    }
+
+    with pytest.raises(ValueError, match="Invalid quantization method"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
 def test_data_path_mapping_parses_repeat_counts(tmp_path: Path) -> None:
     # Config loading should preserve structured multi-dataset paths so the
     # dataset layer can interpret repeat counts later.

fastvideo/train/callbacks/validation.py

@@ -133,6 +133,7 @@ class ValidationCallback(Callback):
         keyboard_value_scale: float = 1.0,
         offload_training_state: bool = False,
         unload_pipeline_after_validation: bool = False,
+        attn_qat_infer: bool = False,
         **pipeline_kwargs: Any,
     ) -> None:
         self.pipeline_target = str(pipeline_target)
@@ -150,6 +151,7 @@ class ValidationCallback(Callback):
         self.metrics_config = self._parse_metrics_config(metrics_config)
         self.offload_training_state = self._coerce_bool(offload_training_state)
         self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation)
+        self.attn_qat_infer = self._coerce_bool(attn_qat_infer)
         self.pipeline_kwargs = dict(pipeline_kwargs)
 
         # Set after on_train_start.
@@ -273,7 +275,7 @@ class ValidationCallback(Callback):
                 # EMA weights during validation.
                 ema_cb = self._find_ema_callback()
                 ctx = ema_cb.ema_context(transformer) if ema_cb is not None else contextlib.nullcontext(transformer)
-                with ctx as t:
+                with ctx as t, self._attn_qat_infer_context(t):
                     self._run_validation_inner(
                         method,
                         step,
@@ -283,6 +285,44 @@ class ValidationCallback(Callback):
             if self.unload_pipeline_after_validation:
                 self._clear_pipeline_cache()
 
+    @contextlib.contextmanager
+    def _attn_qat_infer_context(self, transformer: torch.nn.Module):
+        if not self.attn_qat_infer:
+            yield
+            return
+
+        from fastvideo.attention.backends.attn_qat_infer import (AttnQatInferImpl, is_attn_qat_infer_available)
+        from fastvideo.attention.backends.attn_qat_train import (
+            AttnQatTrainImpl, )
+        from fastvideo.platforms import AttentionBackendEnum
+
+        layers = [
+            module for module in transformer.modules()
+            if isinstance(getattr(module, "attn_impl", None), AttnQatTrainImpl)
+        ]
+        if not layers:
+            raise RuntimeError("attn_qat_infer validation requested, but the transformer has no ATTN_QAT_TRAIN layers")
+        if not is_attn_qat_infer_available():
+            raise RuntimeError("attn_qat_infer validation requires the sm120 fastvideo-kernel extension")
+
+        previous = [(layer, layer.attn_impl, layer.backend) for layer in layers]
+        try:
+            for layer, impl, _ in previous:
+                layer.attn_impl = AttnQatInferImpl(
+                    num_heads=layer.num_heads,
+                    head_size=layer.head_size,
+                    num_kv_heads=layer.num_kv_heads,
+                    causal=impl.causal,
+                    softmax_scale=layer.softmax_scale,
+                )
+                layer.backend = AttentionBackendEnum.ATTN_QAT_INFER
+            logger.info("Enabled ATTN_QAT_INFER for %d validation attention layers.", len(layers))
+            yield
+        finally:
+            for layer, impl, backend in previous:
+                layer.attn_impl = impl
+                layer.backend = backend
+
     @contextlib.contextmanager
     def _validation_memory_context(
         self,

fastvideo/train/models/ltx2/ltx2.py

@@ -38,6 +38,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.ltx2 import VideoLatentShape
 from fastvideo.pipelines import ForwardBatch, TrainingBatch
+from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing, )
 
@@ -94,6 +95,7 @@ class LTX2Model(WanModel):
         transformer_override_safetensor: str
         | None = None,
         lora: LoraConfig | dict[str, Any] | None = None,
+        attention_backend: AttentionBackendEnum | str | None = None,
         train_audio: bool = False,
         timestep_uniform_prob: float = 0.1,
     ) -> None:
@@ -126,6 +128,7 @@ class LTX2Model(WanModel):
             enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
             transformer_override_safetensor=(transformer_override_safetensor),
             lora=lora,
+            attention_backend=attention_backend,
         )
 
         if trainable:
@@ -149,6 +152,7 @@ class LTX2Model(WanModel):
         enable_gradient_checkpointing_type: str | None,
         training_config: TrainingConfig,
         transformer_override_safetensor: str | None = None,
+        attention_backend: AttentionBackendEnum | str | None = None,
     ) -> torch.nn.Module:
         transformer = load_module_from_path(
             model_path=init_from,
@@ -157,6 +161,7 @@ class LTX2Model(WanModel):
             disable_custom_init_weights=(disable_custom_init_weights),
             override_transformer_cls_name=(self._transformer_cls_name),
             transformer_override_safetensor=(transformer_override_safetensor),
+            attention_backend=attention_backend,
         )
         ckpt_type = (enable_gradient_checkpointing_type or getattr(
             getattr(training_config, "model", None),

fastvideo/train/utils/config.py

@@ -282,9 +282,28 @@ def _parse_pipeline_config(
 
     pipeline_config = PipelineConfig.from_kwargs(kwargs)
     _apply_training_dit_arch_overrides(pipeline_config, dit_arch_overrides)
+    _resolve_dit_quant_config(pipeline_config)
     return pipeline_config
 
 
+def _resolve_dit_quant_config(pipeline_config: Any) -> None:
+    """Resolve a string ``pipeline.dit_config.quant_config`` (e.g.
+    ``nvfp4_qat_train``) into its registered QuantizationConfig instance.
+
+    Model construction calls ``quant_config.get_quant_method()`` inside
+    ``LinearBase.__init__``, so a bare YAML string would crash there.
+    This must live in ``_parse_pipeline_config`` (not ``load_run_config``)
+    because ``dcp_to_diffusers`` rebuilds models from a checkpoint's raw
+    config by calling ``_parse_pipeline_config`` directly.
+    """
+    dit_config = getattr(pipeline_config, "dit_config", None)
+    quant = getattr(dit_config, "quant_config", None)
+    if isinstance(quant, str):
+        from fastvideo.layers.quantization import (
+            get_quantization_config, )
+        dit_config.quant_config = get_quantization_config(quant)()
+
+
 def _split_training_dit_arch_overrides(pipeline_raw: Any) -> tuple[Any, dict[str, Any]]:
     """Remove train-only DiT arch overrides before generic config parsing."""
     if not isinstance(pipeline_raw, dict) or not isinstance(pipeline_raw.get("dit_config"), dict):