From 8e5d6090571d30b13bc34e690da48ab667756fe1 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 26 Jul 2026 05:35:59 -0700 Subject: [PATCH] [feat]: FA4-FP4 ATTN_QAT_INFER on sm_100/sm_103 + NVFP4 weight purge --- docs/inference/optimizations.md | 16 ++ .../basic_ltx2_distilled_fast_profile.py | 16 ++ .../inference/ltx2_3/optimized_nvfp4_t2v.py | 176 ++++++++++++++++ .../configs/overfit_ltx2_t2v_nvfp4_qat.yaml | 18 +- .../attention/backends/attn_qat_infer.py | 192 ++++++++++++++++-- fastvideo/attention/backends/flash_attn.py | 63 +++++- fastvideo/layers/quantization/nvfp4_config.py | 72 ++++++- fastvideo/platforms/cuda.py | 8 +- .../test_attn_qat_infer_capability_gate.py | 24 ++- .../test_attn_qat_infer_arch_gate.py | 184 +++++++++++++++++ .../ops/quantization/test_nvfp4_purge.py | 127 ++++++++++++ fastvideo/train/callbacks/validation.py | 6 +- tests/local_tests/test_nvfp4_fa4.py | 54 +++++ 13 files changed, 923 insertions(+), 33 deletions(-) create mode 100644 examples/inference/ltx2_3/optimized_nvfp4_t2v.py create mode 100644 fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_purge.py diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 4f383827d2..de315e6209 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -119,6 +119,22 @@ pip install "nvidia-cutlass-dsl>=4.5.2" apache-tvm-ffi flashinfer-python The `--no-deps` flag prevents upgrading torch/torchvision. Use the supported PyTorch 2.12.0 and CUDA 13 environment for this kernel. +Branch-to-`nvidia-cutlass-dsl` compatibility (the fork tracks the CuTe DSL API +surface closely): + +| fork branch | cutlass-dsl | notes | +|---|---|---| +| `fp4` | `==4.4.2` (+ `nvidia-cutlass-dsl-libs-base==4.4.2`) | validated set on GB200: `quack-kernels==0.4.1`, `flashinfer-python==0.6.8`, `CUTE_DSL_ENABLE_TVM_FFI=1`, `FASTVIDEO_FA4=1` | +| `fix/cutlass-dsl-4.5` | `>=4.5.2` | carries the `cute.core.ThrMma` -> `cute.ThrMma` fix | +| any | 4.6-era | unsupported: `cute.make_fragment` was removed at module level; fails at CuTe JIT trace | + +`FASTVIDEO_FA4=1` is required alongside the fork: it ships no compiled +FlashAttention-2, so dense attention paths raise ImportError without the FA4 +opt-in. The same kernel also serves `ATTN_QAT_INFER` on sm_100a/sm_103a +(datacenter Blackwell) — the selection log's receipt line +(`ATTN_QAT_INFER resolved: ...`) records the arch, kernel, and quantization +mode that actually bound. + #### Usage Enable FP4 attention via the `--nvfp4_fa4` flag: diff --git a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py index e5f2847464..dae08f5d40 100644 --- a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py +++ b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py @@ -191,6 +191,22 @@ 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: the default purges the always-FP4 layers' + # bf16 originals after conversion. Refine-only layers (the cross-modal + # AV projections) always keep theirs: the base stage profile runs them + # dense by deployment contract -- in the two-stage fast profile AND the + # distilled single-stage deploy. retain_original_weights=True keeps + # everything (debugging). pipeline_config.dit_config.quant_config = NVFP4Config() OUTPUT_DIR.mkdir(parents=True, exist_ok=True) torch_compile_kwargs = { diff --git a/examples/inference/ltx2_3/optimized_nvfp4_t2v.py b/examples/inference/ltx2_3/optimized_nvfp4_t2v.py new file mode 100644 index 0000000000..5e326a9eb7 --- /dev/null +++ b/examples/inference/ltx2_3/optimized_nvfp4_t2v.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LTX-2.3 distilled text-to-video with the optimized NVFP4 inference stack. + +Runs `FastVideo/LTX-2.3-Distilled-Diffusers` on a single GPU with the full +validated optimization stack: + +* NVFP4 block-scaled linear layers (per-16 E2M1 weights + E4M3 scale factors), +* ATTN_QAT_INFER FP4 attention (arch-resolved kernel, receipt logged), +* torch.compile (fullgraph) over the DiT, text encoder, and VAE, +* single-stage 8-step distilled sampling at guidance 1.0. + +Quick start +----------- + # On GB200-class ARM hosts, unset LD_LIBRARY_PATH (see Hardware notes): + env -u LD_LIBRARY_PATH python examples/inference/ltx2_3/optimized_nvfp4_t2v.py + + # Optional overrides: + # export LTX23_MODEL_PATH=/path/to/local/snapshot + # export LTX23_T2V_PROMPT="a red fox running through fresh snow" + # export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_nvfp4_t2v + +Hardware notes +-------------- +- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a + system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM (some ARM + container images ship an LD_LIBRARY_PATH that breaks torch.compile's + toolchain discovery). The `_inductor.shape_padding = False` line below + also avoids a pad_mm landmine on the same generation of cards. +""" +from __future__ import annotations + +import os +import time +from pathlib import Path + +import torch._inductor.config as _inductor + +from fastvideo import VideoGenerator +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.layers.quantization.nvfp4_config import NVFP4Config +from fastvideo.utils import maybe_download_model + +# ATTN_QAT_INFER is the FP4 attention half of the NVFP4 deploy contract. It +# resolves per arch (CUTLASS SageAttention3-FP4 on sm_120a/sm_121a, FP4 FA4 +# on sm_100a/sm_103a) and logs a one-line receipt of what actually bound. +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER") +os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1") + +# Inductor knobs. The first one (shape_padding=False) is mandatory on +# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm. The rest +# are the same matmul-friendliness flags the sibling LTX-2 examples use. +_inductor.shape_padding = False +_inductor.conv_1x1_as_mm = True # treat 1x1 convolutions as matrix muls +_inductor.coordinate_descent_tuning = True +_inductor.coordinate_descent_check_all_directions = True +_inductor.epilogue_fusion = False # do not fuse pointwise ops into matmuls + +MODEL_ID = os.path.expandvars( + os.path.expanduser( + os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers") + ) +) +OUTPUT_DIR = Path(os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_nvfp4_t2v")) +DEFAULT_PROMPT = ( + "A fashion model takes a slow step forward and shifts her weight, " + "the soft fabric of her clothing swaying and rippling with the " + "motion, her hair shifting gently, soft even studio lighting on a " + "clean light background, elegant slow-motion runway feel." +) +PROMPT = os.getenv("LTX23_T2V_PROMPT", DEFAULT_PROMPT) + + +def main() -> None: + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + model_root = maybe_download_model(MODEL_ID) + print(f"Model: {model_root}") + print(f"Output dir: {OUTPUT_DIR.resolve()}") + + # Loading the pipeline config *with model_path* binds model-specific + # tuning (notably VAE precision/decoder defaults) into the config. + pipeline_config = PipelineConfig.from_pretrained(model_root) + + # NVFP4 linear layers on the DiT. The default purges the original BF16 + # weights of the always-FP4 linears right after conversion — a large + # peak-memory reduction. Refine-only layers (the cross-modal AV + # projections) always keep theirs: the base stage profile runs them + # dense by deployment contract. retain_original_weights=True keeps + # everything (debugging). + pipeline_config.dit_config.quant_config = NVFP4Config() + + # fullgraph=True is supported on the NVFP4 path: the FP4 quantize step + # is a registered custom op (fastvideo::nvfp4_quantize_fa4), so dynamo + # traces through it without graph breaks. mode="default" — the two + # CUDAGraph modes ("reduce-overhead" / "max-autotune") are a separate + # opt-in, not part of this validated preset. + torch_compile_kwargs = { + "backend": "inductor", + "fullgraph": True, + "mode": "default", + "dynamic": False, + } + + generator = VideoGenerator.from_pretrained( + model_root, + num_gpus=1, + pipeline_config=pipeline_config, + # Compile the DiT, text encoder, and VAE — all three stages benefit, + # and the VAE's codec submodules compile cleanly under fullgraph. + enable_torch_compile=True, + enable_torch_compile_text_encoder=True, + enable_torch_compile_vae=True, + torch_compile_kwargs=torch_compile_kwargs, + torch_compile_kwargs_vae=torch_compile_kwargs, + # Keep everything resident — no CPU offload for serving-style runs. + dit_cpu_offload=False, + text_encoder_cpu_offload=False, + vae_cpu_offload=False, + ltx2_vae_tiling=False, + ) + + common_kwargs = dict( + prompt=PROMPT, + negative_prompt="", # distilled is CFG-free; no negative needed + guidance_scale=1.0, # CFG=1 for distilled + height=1280, width=832, # portrait runway aspect + num_frames=121, fps=24, # ~5s clip + # Single-stage 8-step distilled sampling — the validated preset for + # this checkpoint (no two-stage refine; the NVFP4 deploy contract + # runs the distilled single-stage recipe). + num_inference_steps=8, + save_video=True, + ) + + try: + # Warmup: pays cold compile + first-shape guard work, untimed. + print("\n[warmup] compiling + generating…") + generator.generate_video( + output_path=str(OUTPUT_DIR / "_warmup.mp4"), + seed=7, + **common_kwargs, + ) + (OUTPUT_DIR / "_warmup.mp4").unlink(missing_ok=True) + + # Measured run. + out_path = OUTPUT_DIR / "output_ltx2_3_nvfp4_t2v.mp4" + print(f"\n[measured] generating: {out_path}") + t0 = time.perf_counter() + result = generator.generate_video( + output_path=str(out_path), + seed=2002, + **common_kwargs, + ) + wall = time.perf_counter() - t0 + e2e = (result.get("e2e_latency") if isinstance(result, dict) else None) or wall + print(f"[measured] e2e={e2e:.2f}s wall={wall:.2f}s") + finally: + generator.shutdown() + + # Expected receipts — verify these two lines in your own run's log: + # + # 1. ATTN_QAT_INFER routing receipt (logged at backend resolution; on a + # GB200-class part it reads): + # + # ATTN_QAT_INFER resolved: arch=sm_100 kernel=flash-attention-fp4 \ + # qk_mode=nvfp4(per-16-e4m3-sf) pv_mode=bf16 train_sim_mismatch=measured + # + # 2. NVFP4 weight purge receipt (logged after model conversion; N/M/X + # depend on the checkpoint): + # + # NVFP4 weight purge receipt: purged N original bf16 weight tensors \ + # (X.XX GiB freed); retained M (refine-only dense fallback or \ + # retain_original_weights). + + +if __name__ == "__main__": + main() diff --git a/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml b/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml index 46db30c6f9..242ceb2ed6 100644 --- a/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml +++ b/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml @@ -10,7 +10,19 @@ # 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; validated install set: +# nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1, +# flashinfer-python==0.6.8, FASTVIDEO_FA4=1 — see +# docs/inference/optimizations.md). 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 +31,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 diff --git a/fastvideo/attention/backends/attn_qat_infer.py b/fastvideo/attention/backends/attn_qat_infer.py index 0749b2f548..29d50fbb52 100644 --- a/fastvideo/attention/backends/attn_qat_infer.py +++ b/fastvideo/attention/backends/attn_qat_infer.py @@ -54,29 +54,139 @@ 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)}) + +# The fork is written against the cutlass-dsl 4.4 API surface; the validated +# install set (GB200-proven) is nvidia-cutlass-dsl==4.4.2 + +# nvidia-cutlass-dsl-libs-base==4.4.2 + quack-kernels==0.4.1 + +# flashinfer-python==0.6.8, with the fork on PYTHONPATH, +# CUTE_DSL_ENABLE_TVM_FFI=1, and FASTVIDEO_FA4=1 (the fork ships no compiled +# FA2, so dense attention paths need the FA4 opt-in). dsl 4.6-era installs +# fail at CuTe JIT trace (cute.make_fragment was removed at module level). +_FA4_INSTALL_HINT = ("install flash-attention-fp4 (branch fp4) from " + "https://github.com/hao-ai-lab/flash-attention-fp4 with " + "nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1, " + "flashinfer-python==0.6.8 and FASTVIDEO_FA4=1; " + "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})" + + +_FA4_ROUTE_OPS: tuple | None = None + + +def _import_fa4_route_ops() -> tuple: + """Slow path (own function so tests pin it runs once per process): + resolves the FA4 quantize helper and kernel entry point.""" + from fastvideo.attention.backends.flash_attn import ( + _nvfp4_quantize_for_fa4, ) + from fastvideo.attention.utils.flash_attn_cute import ( + flash_attn_fp4_func, ) + return (_nvfp4_quantize_for_fa4, flash_attn_fp4_func) + + +def _resolve_fa4_route_ops() -> tuple: + # Lazy but memoized: per-forward resolution graph-breaks dynamo every + # step and blocks fullgraph compilation of the NVFP4 path. + global _FA4_ROUTE_OPS + if _FA4_ROUTE_OPS is None: + _FA4_ROUTE_OPS = _import_fa4_route_ops() + return _FA4_ROUTE_OPS + + +_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 +232,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 +246,18 @@ def forward( value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: + # Dispatch on the single per-arch resolution: importability of the + # bundled sm_12x extension is NOT sufficient (CUDA 13 wheels carry it + # on unsupported hosts, where calling it is the wrong binary). + kernel = _resolved_kernel() + if kernel == "fa4_fp4": + return self._forward_fa4_fp4(query, key, value) + if kernel is None: + 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.") + 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.") + assert attn_qat_infer is not None # kernel == "cutlass_sm12x" implies the import succeeded query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() @@ -148,3 +272,39 @@ 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.""" + _nvfp4_quantize_for_fa4, flash_attn_fp4_func = _resolve_fa4_route_ops() + + 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 diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 3962173ca0..d53f0d9b09 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -32,8 +32,30 @@ flash_attn_fp4_func = None _FA4_FP4_AVAILABLE = False +_FA4_QUANT_OPS: tuple | None = None -def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + +def _import_fa4_quant_ops() -> tuple: + """The slow path: resolves flashinfer's FP4 quantization entry points + (imports + JIT-module lookup, which probes the CUDA toolchain via a + subprocess on first use). Kept as its own function so tests can pin + that it runs at most once per process.""" + from flashinfer.quantization import SfLayout, nvfp4_quantize + return (nvfp4_quantize, SfLayout) + + +def _resolve_fa4_quant_ops() -> tuple: + # Lazy (construct-anywhere/fail-at-forward stays intact) but memoized: + # re-resolving per forward graph-breaks dynamo every step (making + # fullgraph compilation impossible for the NVFP4 path) and keeps eager + # dispatch overhead on the hot path. + global _FA4_QUANT_OPS + if _FA4_QUANT_OPS is None: + _FA4_QUANT_OPS = _import_fa4_quant_ops() + return _FA4_QUANT_OPS + + +def _nvfp4_quantize_for_fa4_impl(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Quantize a (batch, seqlen, nheads, headdim) BF16 tensor to FP4. Returns: @@ -42,7 +64,7 @@ def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, to Caller should slice [:, :orig_seqlen] before passing to FA4. sf_tensor: torch.uint8, shape (32, 4, rest_m, 4, rest_k, nheads, batch) with stride[3]=1 """ - from flashinfer.quantization import nvfp4_quantize, SfLayout + nvfp4_quantize, SfLayout = _resolve_fa4_quant_ops() batch, seqlen, nheads, headdim = tensor_4d.shape sf_vec_size = 16 @@ -82,6 +104,43 @@ def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor, ) -> tuple[torch.Tensor, to return fp4_tensor, sf_mma +# Dynamo boundary for the FP4 quantize path (same pattern as the masked +# flash-attention entry points in flash_attn_no_pad.py): tracing a python +# body that resolves flashinfer's JIT module descends into its toolchain +# probe (a subprocess) regardless of any runtime memoization -- a cache hit +# is invisible at trace time. Registering the whole quantize step as a +# custom op makes it one opaque graph node, unlocking +# torch.compile(fullgraph=True) for the NVFP4 attention path. +@torch.library.custom_op( + "fastvideo::nvfp4_quantize_fa4", + mutates_args=(), + device_types="cuda", +) +def _nvfp4_quantize_fa4_op(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return _nvfp4_quantize_for_fa4_impl(tensor_4d) + + +@torch.library.register_fake("fastvideo::nvfp4_quantize_fa4") +def _nvfp4_quantize_fa4_fake(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch, seqlen, nheads, headdim = tensor_4d.shape + seqlen_padded = (seqlen + 127) // 128 * 128 + fp4 = tensor_4d.new_empty((batch, seqlen_padded, nheads, headdim // 2), dtype=torch.float4_e2m1fn_x2) + rest_m = seqlen_padded // 128 + rest_k = (headdim // 16) // 4 + # Must reproduce the impl's output STRIDES, not just its shape: the real + # sf is a permuted view of a contiguous (batch, nheads, rest_m, rest_k, + # 32, 4, 4) buffer, and torch.compile bakes the fake's strides into the + # generated code (a contiguous fake here asserts at runtime). + sf = tensor_4d.new_empty((batch, nheads, rest_m, rest_k, 32, 4, 4), dtype=torch.uint8).permute(4, 5, 2, 6, 3, 1, 0) + return fp4, sf + + +def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize (batch, seqlen, nheads, headdim) BF16 to FP4 via the custom + op boundary; see _nvfp4_quantize_for_fa4_impl for the layout contract.""" + return torch.ops.fastvideo.nvfp4_quantize_fa4(tensor_4d) + + class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index f4adb02e8b..ba90fa4a73 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/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: + # True = retain every original bf16 weight; None/False (default) = + # purge the purgeable set. Refine-only layers are always retained -- + # the base stage profile runs them dense by deployment contract. + self._retain_original_weights: bool | None = None def create_weights(self, layer: torch.nn.Module, input_size_per_partition: int, output_partition_sizes: list[int], input_size: int, output_size: int, params_dtype: torch.dtype, **extra_weight_attrs): @@ -349,7 +354,11 @@ 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 @@ -357,8 +366,13 @@ def apply( # 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. + # Default (None/False): purge the purgeable originals -- every + # always-FP4 layer. Refine-only layers (the cross-modal AV + # projections) are ALWAYS retained: the ``base`` stage profile runs + # them dense by deployment contract, including the distilled + # single-stage deploy. True: retain everything (debugging / + # pre-purge behavior). + self.retain_original_weights = retain_original_weights def get_name(self): return "nvfp4" @@ -438,7 +460,10 @@ 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 @@ -446,7 +471,9 @@ def get_quant_method(self, layer: torch.nn.Module, prefix: str): # 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 @@ -454,6 +481,9 @@ 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,36 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None: persistent=False, ) + retain_flag = getattr(qm, "_retain_original_weights", None) + # Refine-only layers are NEVER purgeable: the "base" stage profile + # runs them dense by deployment contract (the distilled + # single-stage deploy included — its forward context is the base + # profile, so e.g. audio_to_video_attn routes dense every step). + # retain_original_weights therefore only widens retention + # (True = keep everything); it cannot narrow it below the + # dense-capable set. + retain = qm._is_refine_only_layer or retain_flag is True + if retain: + retained += 1 + elif isinstance(weight, DTensor): + # ponytail: purging FSDP-sharded originals needs per-shard + # resharding bookkeeping; skip until a sharded deploy needs it. + retained += 1 + else: + purged_bytes += weight.numel() * weight.element_size() + purged += 1 + mod.register_parameter("weight", None) + + if purged or retained: + logger.info( + "NVFP4 weight purge receipt: purged %d original bf16 weight tensors " + "(%.2f GiB freed); retained %d (refine-only dense fallback or " + "retain_original_weights).", + purged, + purged_bytes / (1 << 30), + retained, + ) + __all__ = [ "NVFP4Config", diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index 4807201f4f..ec257d6d1c 100644 --- a/fastvideo/platforms/cuda.py +++ b/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) diff --git a/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py b/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py index ff157169af..4adc874ee0 100644 --- a/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py +++ b/fastvideo/tests/api/test_attn_qat_infer_capability_gate.py @@ -44,7 +44,11 @@ } -def _fake_gpu(monkeypatch, *, capability: tuple[int, int], extension_imports: bool) -> None: +def _fake_gpu(monkeypatch, + *, + capability: tuple[int, int], + extension_imports: bool, + fa4_imports: bool = False) -> None: monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability) monkeypatch.setattr( @@ -52,6 +56,10 @@ def _fake_gpu(monkeypatch, *, capability: tuple[int, int], extension_imports: bo "_get_attn_qat_infer", lambda: (lambda *a, **k: None) if extension_imports else None, ) + # The FP4 FA4 probe (flash-attention-fp4, the sm_100a/sm_103a route) is + # a physical fact of the running environment; fake it like the others so + # these tests are deterministic on hosts with/without flash_attn.cute. + monkeypatch.setattr(attn_qat_infer_module, "_fa4_fp4_available", lambda: fa4_imports) def _resolve() -> str: @@ -72,12 +80,24 @@ def test_sm90_host_with_bundled_extension_falls_back(monkeypatch): def test_sm100_host_with_bundled_extension_falls_back(monkeypatch): - _fake_gpu(monkeypatch, capability=(10, 0), extension_imports=True) + """sm_100 with only the (unrunnable) bundled sm_12x extension and no + FP4 FA4 kernel still falls back -- the original reviewed failure class.""" + _fake_gpu(monkeypatch, capability=(10, 0), extension_imports=True, fa4_imports=False) assert not is_attn_qat_infer_available() assert _resolve() in FALLBACK_CLASSES +@pytest.mark.parametrize("capability", [(10, 0), (10, 3)]) +def test_datacenter_blackwell_with_fa4_selects_backend(monkeypatch, capability): + """sm_100a/sm_103a resolve ATTN_QAT_INFER through the FP4 FA4 route + when flash-attention-fp4 is installed (even without the sm_12x ext).""" + _fake_gpu(monkeypatch, capability=capability, extension_imports=False, fa4_imports=True) + + assert is_attn_qat_infer_available() + assert _resolve() == ATTN_QAT_INFER_CLS + + @pytest.mark.parametrize("capability", [(12, 0), (12, 1)]) def test_consumer_blackwell_with_extension_selects_backend(monkeypatch, capability): _fake_gpu(monkeypatch, capability=capability, extension_imports=True) diff --git a/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py b/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py new file mode 100644 index 0000000000..c13c9a55de --- /dev/null +++ b/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the per-arch ATTN_QAT_INFER kernel resolution. + +The capability sets route: + * sm_120a/sm_121a -> fastvideo-kernel CUTLASS extension, + * sm_100a/sm_103a -> FP4 FA4 (flash-attention-fp4, #1221 plumbing), + * anything else -> unavailable (honest FlashAttention fallback upstream). + +All device/import probes are monkeypatched; no GPU or kernel install needed. +""" +from __future__ import annotations + +import pytest + +import fastvideo.attention.backends.attn_qat_infer as aqi + + +def _patch(monkeypatch, *, cap, cutlass, fa4) -> None: + monkeypatch.setattr(aqi, "_active_capability", lambda: cap) + monkeypatch.setattr(aqi, "_get_attn_qat_infer", lambda: (lambda *a, **k: None) if cutlass else None) + monkeypatch.setattr(aqi, "_fa4_fp4_available", lambda: fa4) + + +@pytest.mark.parametrize( + "cap,cutlass,fa4,expected_kernel", + [ + ((12, 0), True, False, "cutlass_sm12x"), + ((12, 1), True, False, "cutlass_sm12x"), + ((12, 0), False, False, None), # ext not built + ((10, 0), False, True, "fa4_fp4"), # GB200 + ((10, 3), False, True, "fa4_fp4"), # GB300 + ((10, 0), False, False, None), # fork not installed + ((9, 0), False, True, None), # H100: FA4-FP4 set excludes it + ((8, 9), True, True, None), # Ada: neither set + (None, True, True, None), # no CUDA + ], +) +def test_arch_resolution(monkeypatch, cap, cutlass, fa4, expected_kernel) -> None: + _patch(monkeypatch, cap=cap, cutlass=cutlass, fa4=fa4) + assert aqi._resolved_kernel() == expected_kernel + assert aqi.is_attn_qat_infer_available() == (expected_kernel is not None) + + +def test_receipt_records_fa4_quant_knobs(monkeypatch) -> None: + _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True) + receipt = aqi.attn_qat_infer_receipt() + assert "arch=sm_100" in receipt + assert "qk_mode=nvfp4(per-16-e4m3-sf)" in receipt + assert "pv_mode=bf16" in receipt + assert "train_sim_mismatch=measured" in receipt + + +def test_receipt_records_cutlass_scheme(monkeypatch) -> None: + _patch(monkeypatch, cap=(12, 0), cutlass=True, fa4=False) + receipt = aqi.attn_qat_infer_receipt() + assert "arch=sm_120" in receipt + assert "cutlass" in receipt + + +def test_receipt_names_install_hint_on_uninstalled_fa4_arch(monkeypatch) -> None: + _patch(monkeypatch, cap=(10, 3), cutlass=False, fa4=False) + receipt = aqi.attn_qat_infer_receipt() + assert "arch=sm_103" in receipt + assert "flash-attention-fp4" in receipt + + +def test_unsupported_arch_receipt_lists_support_matrix(monkeypatch) -> None: + _patch(monkeypatch, cap=(9, 0), cutlass=True, fa4=True) + receipt = aqi.attn_qat_infer_receipt() + assert "kernel=none" in receipt + assert "sm_100a/sm_103a" in receipt + + +def test_unsupported_arch_forward_never_calls_bundled_extension(monkeypatch) -> None: + """An unsupported GPU (e.g. sm_90) carrying an importable sm_12x wheel + must fail cleanly at forward — never dispatch into the wrong binary.""" + calls = [] + monkeypatch.setattr(aqi, "_active_capability", lambda: (9, 0)) + monkeypatch.setattr(aqi, "_get_attn_qat_infer", lambda: (lambda *a, **k: calls.append(1))) + monkeypatch.setattr(aqi, "_fa4_fp4_available", lambda: False) + + import torch + + impl = aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5) + q = torch.zeros(1, 4, 1, 128) + with pytest.raises(ImportError, match="not available"): + impl.forward(q, q, q, attn_metadata=None) + assert not calls, "unsupported arch dispatched into the bundled CUTLASS extension" + + +def test_fa4_route_resolution_runs_once_across_forwards(monkeypatch) -> None: + """The FA4 kernel/quantizer resolution is lazy but memoized: across N + forward calls the slow resolution path (imports + toolchain probe) + executes exactly once — per-forward re-resolution graph-breaks dynamo + every step and blocks fullgraph compilation.""" + import torch + + _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True) + + resolves = [] + + def fake_quant(t): + return t, torch.zeros(1) + + def fake_kernel(q, k, v, sfq, sfk, softmax_scale=None, causal=False): + return v + + monkeypatch.setattr(aqi, "_import_fa4_route_ops", lambda: (resolves.append(1) or (fake_quant, fake_kernel))) + monkeypatch.setattr(aqi, "_FA4_ROUTE_OPS", None) + + impl = aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5) + q = torch.zeros(1, 4, 1, 128) + for _ in range(3): + out = impl.forward(q, q, q, attn_metadata=None) + assert out.shape == q.shape + assert len(resolves) == 1, f"resolution ran {len(resolves)}x across 3 forwards" + + +def _register_fa4_quantize_cpu_kernel(): + """Register a CPU kernel for fastvideo::nvfp4_quantize_fa4 that honors the + production output contract — including STRIDES: the real sf output is a + permuted view of a contiguous (batch, nheads, rest_m, rest_k, 32, 4, 4) + buffer, and torch.compile bakes output strides into generated code.""" + import torch + + def _cpu_kernel(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch, seqlen, nheads, headdim = tensor_4d.shape + seqlen_padded = (seqlen + 127) // 128 * 128 + fp4 = torch.zeros(batch, seqlen_padded, nheads, headdim // 2, + dtype=torch.int8).view(torch.float4_e2m1fn_x2) + rest_m, rest_k = seqlen_padded // 128, (headdim // 16) // 4 + sf = torch.zeros(batch, nheads, rest_m, rest_k, 32, 4, 4, + dtype=torch.uint8).permute(4, 5, 2, 6, 3, 1, 0) + return fp4, sf + + try: + torch.library.register_kernel("fastvideo::nvfp4_quantize_fa4", "cpu")(_cpu_kernel) + except RuntimeError: + pass # already registered by a previous test/run + + +def test_fa4_quantize_path_is_fullgraph_traceable() -> None: + """The FP4 quantize step must be a single opaque graph node: tracing its + python body descends into flashinfer's toolchain probe (a subprocess), + which torch.compile(fullgraph=True) rejects — a runtime memo cannot fix + that because cache hits are invisible at trace time. fullgraph=True over + the op-backed path must compile and run without a graph break.""" + import torch + + from fastvideo.attention.backends import flash_attn as fa + + _register_fa4_quantize_cpu_kernel() + + def path(x: torch.Tensor) -> torch.Tensor: + fp4, sf = fa._nvfp4_quantize_for_fa4(x) + return fp4[:, :x.shape[1]].view(torch.int8).float() + float(sf.shape[0]) + + compiled = torch.compile(path, fullgraph=True, backend="eager") + out = compiled(torch.randn(1, 64, 2, 128, dtype=torch.bfloat16)) + assert out.shape == (1, 64, 2, 64) + + +def test_fa4_quantize_op_fake_matches_real() -> None: + """torch.library.opcheck cross-checks the registered fake against a real + kernel run — shapes, dtypes, AND strides — so a fake whose output layout + drifts from the impl (e.g. contiguous fake vs permuted-view real sf, which + torch.compile turns into a runtime stride assertion) fails here on CPU. + + The op is forward-only (no autograd registration), so restrict opcheck to + the non-autograd suites.""" + import torch + + import fastvideo.attention.backends.flash_attn # noqa: F401 registers the op + fake + + _register_fa4_quantize_cpu_kernel() + + # batch>1 and seqlen>128 so every sf dim (incl. batch, rest_m) is + # non-trivial and its stride actually participates in the check. + x = torch.randn(2, 200, 2, 128, dtype=torch.bfloat16) + torch.library.opcheck( + torch.ops.fastvideo.nvfp4_quantize_fa4, + (x,), + test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), + ) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py new file mode 100644 index 0000000000..5209108021 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU construction-level tests for the NVFP4 original-weight purge. + +``convert_model_to_nvfp4`` keeps the packed FP4 buffers and, by default +(auto), purges the bf16 ``layer.weight`` of always-FP4 layers while +retaining refine-only layers (which still take the stage-profile dense +path). ``retain_original_weights`` on ``NVFP4Config`` overrides in both +directions. flashinfer and the FP4 GEMM are monkeypatched, so this runs +on any host. +""" +from __future__ import annotations + +import logging +import types + +import pytest +import torch +import torch.nn as nn + +import fastvideo.layers.quantization.nvfp4_config as nv + +_ALWAYS_FP4_PREFIX = "ltx2.blocks.0.attn1.to_q" +_REFINE_ONLY_PREFIX = "ltx2.blocks.0.audio_to_video_attn.to_q" + + +def _method(prefix: str, retain: bool | None) -> nv.NVFP4QuantizeMethod: + # NVFP4QuantizeMethod.__init__ allocates x_global_sf on cuda; build the + # object without it so the purge tests run on CPU-only hosts. + m = object.__new__(nv.NVFP4QuantizeMethod) + m.weight_fp4 = None + m.weight_scale = None + m.x_global_sf = torch.tensor(1.0, dtype=torch.float32) + m.layer_prefix = prefix + m._is_refine_only_layer = nv._is_ltx2_refine_only_prefix(prefix) + m._retain_original_weights = retain + return m + + +class _FakeLinear(nn.Module): + + def __init__(self, prefix: str, retain: bool | None, out_dim: int = 8, in_dim: int = 16) -> None: + super().__init__() + self.weight = nn.Parameter(torch.randn(out_dim, in_dim, dtype=torch.bfloat16), requires_grad=False) + self.quant_method = _method(prefix, retain) + + +def _fake_quantize(weight, global_sf, sfLayout=None, do_shuffle=False): + out_dim, in_dim = weight.shape[0], weight.shape[-1] + return (torch.zeros(out_dim, in_dim // 2, dtype=torch.int8), torch.zeros(out_dim, in_dim // 16, + dtype=torch.uint8)) + + +@pytest.fixture(autouse=True) +def _patch_flashinfer(monkeypatch): + fake_sf_layout = types.SimpleNamespace(layout_128x4=None) + monkeypatch.setattr(nv, "_require_flashinfer", lambda: (fake_sf_layout, None, None)) + monkeypatch.setattr(nv, "_nvfp4_quantize", _fake_quantize) + + +def _model(retain: bool | None) -> nn.Module: + root = nn.Module() + root.always_fp4 = _FakeLinear(_ALWAYS_FP4_PREFIX, retain) + root.refine_only = _FakeLinear(_REFINE_ONLY_PREFIX, retain) + return root + + +def test_prefix_classification_sanity() -> None: + assert not nv._is_ltx2_refine_only_prefix(_ALWAYS_FP4_PREFIX) + assert nv._is_ltx2_refine_only_prefix(_REFINE_ONLY_PREFIX) + + +def test_auto_purges_always_fp4_and_retains_refine_only(caplog) -> None: + model = _model(retain=None) + with caplog.at_level(logging.INFO): + nv.convert_model_to_nvfp4(model) + assert model.always_fp4.weight is None + assert model.refine_only.weight is not None + assert model.always_fp4._nvfp4_weight is not None + receipt = [r.message for r in caplog.records if "NVFP4 weight purge receipt" in r.message] + assert receipt and "purged 1" in receipt[0] and "retained 1" in receipt[0] + + +def test_retain_true_keeps_everything() -> None: + model = _model(retain=True) + nv.convert_model_to_nvfp4(model) + assert model.always_fp4.weight is not None + assert model.refine_only.weight is not None + + +def test_retain_false_still_retains_dense_capable_layers() -> None: + """Refine-only layers run dense under the base stage profile by + deployment contract (single-stage deploys included), so no flag value + may purge them.""" + model = _model(retain=False) + nv.convert_model_to_nvfp4(model) + assert model.always_fp4.weight is None + assert model.refine_only.weight is not None + + +def test_apply_out_dim_survives_purge(monkeypatch) -> None: + model = _model(retain=False) + nv.convert_model_to_nvfp4(model) + layer = model.always_fp4 + captured = {} + + def fake_mm_fp4(x_fp4, w_t, x_scale, w_scale_t, alpha, out_dtype, out, backend): + captured["out_dim"] = w_t.shape[-1] if w_t.dim() == 2 else None + return torch.zeros(x_fp4.shape[0], w_t.shape[-1], dtype=torch.bfloat16) + + monkeypatch.setattr(nv, "_mm_fp4", fake_mm_fp4) + monkeypatch.setattr(nv, "_get_ltx2_fp4_stage_profile", lambda default="refine": "refine") + layer._weight_global_sf = torch.tensor(1.0, dtype=torch.bfloat16) + x = torch.randn(2, 3, 16, dtype=torch.bfloat16) + out = layer.quant_method.apply(layer, x) + assert out.shape == (2, 3, 8) + + +def test_dense_path_after_purge_raises_with_flag_named(monkeypatch) -> None: + """Defensive guard: convert never purges dense-capable layers anymore, + but a hand-purged module hitting the dense path must fail loudly.""" + model = _model(retain=False) + nv.convert_model_to_nvfp4(model) + model.refine_only.register_parameter("weight", None) # simulate misuse + monkeypatch.setattr(nv, "_get_ltx2_fp4_stage_profile", lambda default="refine": "base") + x = torch.randn(2, 3, 16, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="retain_original_weights"): + model.refine_only.quant_method.apply(model.refine_only, x) diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index bc84443a86..800a681c38 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -303,7 +303,11 @@ def _attn_qat_infer_context(self, transformer: torch.nn.Module): 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") + from fastvideo.attention.backends.attn_qat_infer import ( + attn_qat_infer_receipt, ) + raise RuntimeError("attn_qat_infer validation requested but no ATTN_QAT_INFER kernel serves " + f"this device ({attn_qat_infer_receipt()}). Set " + "callbacks.validation.attn_qat_infer=false to validate with ATTN_QAT_TRAIN.") previous = [(layer, layer.attn_impl, layer.backend) for layer in layers] try: diff --git a/tests/local_tests/test_nvfp4_fa4.py b/tests/local_tests/test_nvfp4_fa4.py index d6d837c6dc..132c6590c5 100644 --- a/tests/local_tests/test_nvfp4_fa4.py +++ b/tests/local_tests/test_nvfp4_fa4.py @@ -149,3 +149,57 @@ def test_attention_speedup(self): if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)], + reason="Requires Blackwell GPU (sm100a or sm103a)" +) +class TestAttnQatInferFA4Route: + """ATTN_QAT_INFER resolves to the FP4 FA4 kernel on sm_100/sm_103. + + Hardware-validated on both GB200 (sm_100a) and GB300 (sm_103a): route + resolution, SDPA parity, and cross-attention lengths pass on real + silicon for each capability in the FA4-FP4 set. + """ + + def test_resolves_to_fa4(self): + from fastvideo.attention.backends.attn_qat_infer import ( + _resolved_kernel, attn_qat_infer_receipt, is_attn_qat_infer_available) + assert is_attn_qat_infer_available() + assert _resolved_kernel() == "fa4_fp4" + receipt = attn_qat_infer_receipt() + assert "qk_mode=nvfp4(per-16-e4m3-sf)" in receipt + assert "pv_mode=bf16" in receipt + + def test_forward_parity_vs_sdpa(self): + from fastvideo.attention.backends.attn_qat_infer import AttnQatInferImpl + torch.manual_seed(0) + b, s, h, d = 1, 4096, 12, 128 + q = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) + k = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) + v = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) + + impl = AttnQatInferImpl(num_heads=h, head_size=d, causal=False, softmax_scale=d**-0.5) + out = impl.forward(q, k, v, attn_metadata=None) + + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), + v.transpose(1, 2).float()).transpose(1, 2) + + cos = torch.nn.functional.cosine_similarity(out.float().flatten(), ref.flatten(), dim=0).item() + # Same bound the sm_120 CUTLASS kernel test uses; FA4 NVFP4 QK + # measures cos_sim ~0.99 vs BF16 (kernel repo README precision table). + assert cos >= 0.97, f"cos_sim={cos:.4f} < 0.97" + + def test_cross_attention_lengths(self): + from fastvideo.attention.backends.attn_qat_infer import AttnQatInferImpl + torch.manual_seed(0) + b, h, d = 1, 12, 128 + q = torch.randn(b, 384, h, d, device="cuda", dtype=torch.bfloat16) + k = torch.randn(b, 512, h, d, device="cuda", dtype=torch.bfloat16) + v = torch.randn(b, 512, h, d, device="cuda", dtype=torch.bfloat16) + impl = AttnQatInferImpl(num_heads=h, head_size=d, causal=False, softmax_scale=d**-0.5) + out = impl.forward(q, k, v, attn_metadata=None) + assert out.shape == (b, 384, h, d)