[feat]: FA4-FP4 ATTN_QAT_INFER on sm_100/sm_103 + NVFP4 weight purge

Author: SolitaryThinker

examples/inference/basic/basic_ltx2_distilled_fast_profile.py

@@ -191,6 +191,21 @@ def main() -> None:
     print(f"Using refine upsampler: {refine_upsampler_path}")
 
     pipeline_config = PipelineConfig.from_pretrained(model_root)
+    # LTX-2 NVFP4 deploy contract (train==deploy surface):
+    #   * Linears: NVFP4 block-scaled GEMMs (per-16 E2M1 + E4M3 SFs) on every
+    #     arch, via flashinfer.
+    #   * ATTN_QAT_INFER attention differs per arch: sm_120a/sm_121a use the
+    #     fastvideo-kernel CUTLASS (SageAttention3-FP4) scheme that
+    #     ATTN_QAT_TRAIN simulates; sm_100a (GB200) / sm_103a (GB300) use the
+    #     FP4 FA4 kernel (flash-attention-fp4) with per-16 block-scaled NVFP4
+    #     Q/K and BF16 P/V -- a train-sim mismatch that is gated by MS-SSIM
+    #     measurement, not assumed equal. The selection receipt is logged at
+    #     backend resolution ("ATTN_QAT_INFER resolved: ...").
+    # Original-weight retention: default (None) purges the always-FP4 layers'
+    # bf16 originals after conversion and retains only refine-only layers,
+    # which this two-stage fast profile still runs dense during stage 1.
+    # Single-stage deploys (no refine stage) can purge everything with
+    # NVFP4Config(retain_original_weights=False).
     pipeline_config.dit_config.quant_config = NVFP4Config()
     OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
     torch_compile_kwargs = {

examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml

@@ -10,7 +10,16 @@
 # 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.
+# Validation-time ATTN_QAT_INFER is arch-aware:
+#   * sm_120a/sm_121a: fastvideo-kernel CUTLASS extension -- the exact
+#     quantization scheme ATTN_QAT_TRAIN simulates.
+#   * sm_100a (GB200) / sm_103a (GB300): FP4 FA4 kernel
+#     (github.com/hao-ai-lab/flash-attention-fp4, branch fp4; per-16
+#     block-scaled NVFP4 Q/K, BF16 P/V). This scheme DIFFERS from the
+#     CUTLASS one the training simulation matches, so sm_100/sm_103
+#     validation and deployment carry a train-sim mismatch -- gate quality
+#     by MS-SSIM measurement rather than assuming parity. The resolution
+#     receipt ("ATTN_QAT_INFER resolved: ...") records arch + scheme.
 #
 # Preprocess data first (same data as the bf16 overfit):
 #   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
@@ -19,8 +28,8 @@
 #   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:
+# On GB200 without flash-attention-fp4 installed (or any other arch with no
+# ATTN_QAT_INFER 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

fastvideo/attention/backends/attn_qat_infer.py

@@ -54,29 +54,108 @@ def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
 # kernel is compiled for (sm_120a / sm_121a -- see fastvideo-kernel/README.md).
 _SUPPORTED_DEVICE_CAPABILITIES = frozenset({(12, 0), (12, 1)})
 
-
-def _device_capability_supported() -> bool:
+# Datacenter-Blackwell capabilities served by the FP4 FA4 kernel
+# (flash-attention-fp4 @ fp4, sm_100a/sm_103a) through #1221's plumbing:
+# per-16 block-scaled NVFP4 Q/K (E4M3 scale factors), BF16 P/V. This is a
+# DIFFERENT quantization scheme from the sm_12x CUTLASS extension above --
+# ATTN_QAT_TRAIN simulates the CUTLASS scheme, so sm_100/sm_103 deployment
+# carries a train-sim mismatch that is measured (MS-SSIM gate), not assumed.
+_FA4_FP4_CAPABILITIES = frozenset({(10, 0), (10, 3)})
+
+_FA4_INSTALL_HINT = ("pip install flash-attention-fp4 from "
+                     "https://github.com/hao-ai-lab/flash-attention-fp4 (branch fp4); "
+                     "see docs/inference/optimizations.md")
+
+_fa4_fp4_import_ok: bool | None = None
+
+
+def _fa4_fp4_available() -> bool:
+    """flash_attn.cute (FA4) import probe, cached. Reuses #1221's guarded
+    import chain in fastvideo.attention.utils.flash_attn_cute (which maps
+    cutlass-dsl version skew to ImportError with a loud warning)."""
+    global _fa4_fp4_import_ok
+    if _fa4_fp4_import_ok is None:
+        try:
+            from fastvideo.attention.utils.flash_attn_cute import (  # noqa: F401
+                flash_attn_fp4_func, )
+            _fa4_fp4_import_ok = True
+        except ImportError:
+            _fa4_fp4_import_ok = False
+    return _fa4_fp4_import_ok
+
+
+def _active_capability() -> tuple[int, int] | None:
     if not torch.cuda.is_available():
-        return False
+        return None
     try:
-        return tuple(torch.cuda.get_device_capability()) in _SUPPORTED_DEVICE_CAPABILITIES
+        return tuple(torch.cuda.get_device_capability())
     except Exception:  # pragma: no cover - defensive: never break backend selection
-        return False
+        return None
+
+
+def _resolved_kernel() -> str | None:
+    """Which ATTN_QAT_INFER kernel serves the active device, or None.
+
+    Per-arch resolution (single source of truth -- extend the capability
+    sets above, do not add equality checks elsewhere):
+      * sm_12x consumer Blackwell -> fastvideo-kernel CUTLASS extension
+        (modified SageAttention3 FP4).
+      * sm_100a/sm_103a datacenter Blackwell -> FP4 FA4 (flash-attention-fp4)
+        via the merged #1221 plumbing.
+    """
+    cap = _active_capability()
+    if cap in _SUPPORTED_DEVICE_CAPABILITIES and _get_attn_qat_infer() is not None:
+        return "cutlass_sm12x"
+    if cap in _FA4_FP4_CAPABILITIES and _fa4_fp4_available():
+        return "fa4_fp4"
+    return None
+
+
+def attn_qat_infer_receipt() -> str:
+    """One-line receipt of the resolution decision (arch + kernel + quant
+    knobs), for the selection log and for tooling. The FA4 knobs are the
+    repo's tuned defaults passed through verbatim: qk_mode=nvfp4
+    (per-16 E4M3 SFs), pv_mode=bf16 -- see flash_attn/cute/README.md in
+    the kernel repo."""
+    cap = _active_capability()
+    arch = f"sm_{cap[0]}{cap[1]}" if cap is not None else "no-cuda"
+    kernel = _resolved_kernel()
+    if kernel == "cutlass_sm12x":
+        return f"arch={arch} kernel=fastvideo-kernel-cutlass scheme=sage3-fp4-sm120"
+    if kernel == "fa4_fp4":
+        return (f"arch={arch} kernel=flash-attention-fp4 qk_mode=nvfp4(per-16-e4m3-sf) "
+                f"pv_mode=bf16 train_sim_mismatch=measured")
+    supported = "sm_120a/sm_121a via fastvideo-kernel build.sh; sm_100a/sm_103a via flash-attention-fp4"
+    if cap is not None and cap in _FA4_FP4_CAPABILITIES:
+        return f"arch={arch} kernel=none (flash_attn.cute not importable -- {_FA4_INSTALL_HINT})"
+    return f"arch={arch} kernel=none (supported: {supported})"
+
+
+_receipt_logged = False
+
+
+def _log_receipt_once() -> None:
+    # One line per process, not per layer (the validation swap constructs
+    # one impl per attention layer).
+    global _receipt_logged
+    if not _receipt_logged:
+        _receipt_logged = True
+        logger.info("ATTN_QAT_INFER resolved: %s", attn_qat_infer_receipt())
 
 
 def is_attn_qat_infer_available() -> bool:
-    """True only when the extension imports AND the active device is a
-    consumer-Blackwell (sm_120/sm_121) GPU the kernel is compiled for.
+    """True only when the active device has a built ATTN_QAT_INFER kernel.
 
     The import check alone is not sufficient: CUDA 13 wheel builds can
-    carry the sm_120/sm_121 extension on any host (e.g. H100, GB200),
-    where the import succeeds, backend selection picks this backend, and
-    the first kernel call then fails with an unsupported-capability error
-    instead of ever reaching the documented FlashAttention fallback in
+    carry the sm_12x extension on any host (e.g. H100, GB200), where the
+    import succeeds, backend selection picks this backend, and the first
+    kernel call then fails with an unsupported-capability error instead of
+    ever reaching the documented FlashAttention fallback in
     fastvideo.platforms.cuda. Gating on the active device's capability
-    keeps that fallback working on every non-sm_120/121 GPU.
+    keeps that fallback working on every unsupported GPU, while
+    sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).
     """
-    return _device_capability_supported() and _get_attn_qat_infer() is not None
+    return _resolved_kernel() is not None
 
 
 class AttnQatInferBackend(AttentionBackend):
@@ -122,6 +201,12 @@ def __init__(
         if dropout_p > 0:
             raise NotImplementedError(f"attn_qat_infer does not support dropout (got dropout_p={dropout_p}). "
                                       "The QAT inference kernel applies no stochastic dropout.")
+        # Kernel resolution is per-forward, not per-construction: callers
+        # (the validation swap, backend selection) gate on
+        # is_attn_qat_infer_available() first, and constructing an impl on a
+        # host without the kernel must stay legal (pre-existing contract the
+        # validation-swap test pins).
+        _log_receipt_once()
 
     def forward(
         self,
@@ -130,10 +215,13 @@ def forward(
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
+        if _resolved_kernel() == "fa4_fp4":
+            return self._forward_fa4_fp4(query, key, value)
+
         attn_qat_infer = _get_attn_qat_infer()
         if attn_qat_infer is None:
-            raise ImportError("attn_qat_infer is not available. Please ensure the "
-                              "attn_qat_infer kernel package is installed.")
+            raise ImportError(f"attn_qat_infer is not available ({attn_qat_infer_receipt()}). "
+                              "Please ensure an ATTN_QAT_INFER kernel is installed for this device.")
 
         query = query.transpose(1, 2).contiguous()
         key = key.transpose(1, 2).contiguous()
@@ -148,3 +236,42 @@ def forward(
             sm_scale=self.softmax_scale,
         )
         return output.transpose(1, 2).contiguous()
+
+    def _forward_fa4_fp4(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        """sm_100a/sm_103a path: FP4 FA4 with the repo's tuned defaults
+        (NVFP4 per-16 block-scaled Q/K, BF16 V) -- mirrors
+        FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
+        (batch, seqlen, nheads, headdim); no transpose."""
+        from fastvideo.attention.backends.flash_attn import (
+            _nvfp4_quantize_for_fa4, )
+        from fastvideo.attention.utils.flash_attn_cute import (
+            flash_attn_fp4_func, )
+
+        orig_seqlen_q = query.shape[1]
+        orig_seqlen_k = key.shape[1]
+
+        q_fp4, q_sf = _nvfp4_quantize_for_fa4(query)
+        k_fp4, k_sf = _nvfp4_quantize_for_fa4(key)
+
+        # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
+        # original lengths so padding never biases the softmax.
+        q_fp4 = q_fp4[:, :orig_seqlen_q]
+        k_fp4 = k_fp4[:, :orig_seqlen_k]
+
+        output = flash_attn_fp4_func(
+            q_fp4,
+            k_fp4,
+            value,
+            q_sf,
+            k_sf,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        if isinstance(output, tuple):
+            output = output[0]
+        return output

fastvideo/layers/quantization/nvfp4_config.py

@@ -309,6 +309,11 @@ def __init__(self, layer_prefix: str = ""):
         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:
+        # None = auto (retain only refine-only layers, which can still take
+        # the stage-profile dense path); True = retain all; False = purge all
+        # (single-stage deploys that never run the base profile).
+        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):
@@ -349,16 +354,25 @@ def apply(
         | None = None,
     ) -> torch.Tensor:
         SfLayout, _, _ = _require_flashinfer()
-        out_dim = layer.weight.shape[0]
+        # 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":
-            out = (F.linear(x, layer.weight, bias) if torch.cuda.is_available() or bias is None else F.linear(
-                x, layer.weight, bias.to(x.dtype)))
+            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
@@ -416,11 +430,19 @@ class NVFP4Config(QuantizationConfig):
     instead of hardcoding it here.
     """
 
-    def __init__(self, layer_profile: str = "refine"):
+    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.
+        # None (auto): purge always-FP4 layers, retain refine-only layers --
+        # those still take the dense path when the streaming stage profile
+        # flips to ``base``. True: retain everything (pre-purge behavior).
+        # False: purge everything -- for single-stage deploys (e.g. the
+        # LTX-2 QAD distilled 8-step deploy) where no dense stage ever runs;
+        # a purged layer hitting the dense path raises with this flag named.
+        self.retain_original_weights = retain_original_weights
 
     def get_name(self):
         return "nvfp4"
@@ -438,22 +460,30 @@ def get_config_filenames():
 
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> NVFP4Config:
-        return cls(layer_profile=config.get("layer_profile", "refine"))
+        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):
-            return NVFP4QuantizeMethod(layer_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):
@@ -486,6 +516,29 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None:
                 persistent=False,
             )
 
+            retain_flag = getattr(qm, "_retain_original_weights", None)
+            retain = retain_flag if retain_flag is not None else qm._is_refine_only_layer
+            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",

fastvideo/platforms/cuda.py

@@ -142,11 +142,13 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea
                 logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
         elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
             from fastvideo.attention.backends.attn_qat_infer import (  # noqa: F401
-                AttnQatInferBackend, is_attn_qat_infer_available)
+                AttnQatInferBackend, attn_qat_infer_receipt, is_attn_qat_infer_available)
             if is_attn_qat_infer_available():
-                logger.info("Using Attn-QAT inference (modified SageAttention3 FP4) backend.")
+                logger.info("Using Attn-QAT inference backend (%s).", attn_qat_infer_receipt())
                 return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
-            logger.info("Attn-QAT inference kernel is not built. Fall back to Flash Attention.")
+            # Keep the trailing sentence stable: downstream receipts grep for it.
+            logger.info("Attn-QAT inference kernel is not built (%s). Fall back to Flash Attention.",
+                        attn_qat_infer_receipt())
         elif selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN:
             from fastvideo.attention.backends.attn_qat_train import (  # noqa: F401
                 AttnQatTrainBackend, is_attn_qat_train_available)