diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index de315e6209..ef2bfc6dce 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -156,6 +156,18 @@ gen = VideoGenerator.from_pretrained( gen.generate_video(prompt="A raccoon in sunflowers", save_video=True) ``` +#### fp8 PV mode **[target]** + +The FA4-FP4 path keeps V in BF16 by default. An `fa4_pv_mode="fp8"` knob (an +attention-impl config field accepted by both `FLASH_ATTN` and `ATTN_QAT_INFER` +on this path; allowed values `"bf16"`/`"fp8"`) casts V to fp8 e4m3 before the +kernel — an unscaled cast, per the kernel's plain-fp8 PV contract. Kernel-level +benchmarks on datacenter Blackwell show it is faster at large shapes, but it +stays opt-in **[target]** pending an end-to-end compiled benchmark and a +quality gate before any default consideration. The `ATTN_QAT_INFER` receipt +line reports the configured `pv_mode`, and the V dtype actually fed to the +kernel is logged once on the first FA4 forward. + #### Known Limitations - `use_fsdp_inference=True` is incompatible with the FP4 path (FSDP shards invalidate tensor pointers) diff --git a/fastvideo/attention/backends/attn_qat_infer.py b/fastvideo/attention/backends/attn_qat_infer.py index 29d50fbb52..8fe5eb74bd 100644 --- a/fastvideo/attention/backends/attn_qat_infer.py +++ b/fastvideo/attention/backends/attn_qat_infer.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import importlib +import os import sys from collections.abc import Callable from pathlib import Path @@ -75,6 +76,36 @@ def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None: "flashinfer-python==0.6.8 and FASTVIDEO_FA4=1; " "see docs/inference/optimizations.md") +# PV-mode knob for the FA4-FP4 path (the "fa4_pv_mode" extra_impl_args key, +# consumed here and by FlashAttentionImpl). "bf16" keeps V in BF16 (default, +# byte-identical to the pre-knob behavior); "fp8" casts V to e4m3 before the +# kernel -- the fork's plain-fp8 PV contract: unscaled cast, no mSFV scale +# factors and no v_descale, BF16 output. +_FA4_PV_MODES = ("bf16", "fp8") + + +# The last configured pv mode; the receipt derives its pv_mode field from +# this instead of declaring a literal. Impl construction records it. +def _default_fa4_pv_mode() -> str: + """Env bridge, mirroring the sibling nvfp4_fa4 pattern: kwargs win, the + FASTVIDEO_FA4_PV_MODE env var is the user-reachable fallback (model code + constructs attention with fixed literals, so without this bridge the knob + has no user path). Reading it here also makes the resolution-time receipt + correct for env-driven runs before any impl is constructed.""" + return os.environ.get("FASTVIDEO_FA4_PV_MODE", "bf16") + + +_configured_fa4_pv_mode = _default_fa4_pv_mode() + + +def validate_fa4_pv_mode(mode: str) -> str: + # Fail fast at impl construction so a typo never survives to the first + # forward on a Blackwell box. + if mode not in _FA4_PV_MODES: + raise ValueError(f"fa4_pv_mode must be one of {_FA4_PV_MODES}, got {mode!r}") + return mode + + _fa4_fp4_import_ok: bool | None = None @@ -122,10 +153,11 @@ def _resolved_kernel() -> str | 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.""" + knobs), for the selection log and for tooling. qk_mode=nvfp4 (per-16 E4M3 + SFs) is the repo's tuned default passed through verbatim; pv_mode is + derived from the configured fa4_pv_mode knob, and the dtype actually fed + to the kernel is logged once on the first FA4 forward -- 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() @@ -133,7 +165,7 @@ def attn_qat_infer_receipt() -> str: 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") + f"pv_mode={_configured_fa4_pv_mode} 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})" @@ -145,12 +177,14 @@ def attn_qat_infer_receipt() -> str: 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.""" + resolves the FA4 quantize/V-cast helpers and kernel entry point.""" from fastvideo.attention.backends.flash_attn import ( - _nvfp4_quantize_for_fa4, ) + _fa4_v_to_fp8, + _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) + return (_nvfp4_quantize_for_fa4, _fa4_v_to_fp8, flash_attn_fp4_func) def _resolve_fa4_route_ops() -> tuple: @@ -174,6 +208,19 @@ def _log_receipt_once() -> None: logger.info("ATTN_QAT_INFER resolved: %s", attn_qat_infer_receipt()) +_pv_dtype_logged = False + + +def _log_pv_dtype_once(dtype: torch.dtype) -> None: + # Derived-from-runtime companion to the receipt: the dtype actually fed + # to the FA4 kernel as V on the first forward, once per process. + global _pv_dtype_logged + if not _pv_dtype_logged: + _pv_dtype_logged = True + logger.info("ATTN_QAT_INFER FA4 first forward: observed V dtype=%s (configured pv_mode=%s)", dtype, + _configured_fa4_pv_mode) + + def is_attn_qat_infer_available() -> bool: """True only when the active device has a built ATTN_QAT_INFER kernel. @@ -232,6 +279,11 @@ 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.") + self.fa4_pv_mode = validate_fa4_pv_mode(extra_impl_args.get("fa4_pv_mode") or _default_fa4_pv_mode()) + # Record the configured mode before the once-log so the receipt line + # (whose pv_mode field is derived from this) carries it. + global _configured_fa4_pv_mode + _configured_fa4_pv_mode = self.fa4_pv_mode # 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 @@ -279,11 +331,11 @@ def _forward_fa4_fp4( 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 + """sm_100a/sm_103a path: FP4 FA4 (NVFP4 per-16 block-scaled Q/K; V in + BF16 by default or fp8 e4m3 per the fa4_pv_mode knob) -- 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() + _nvfp4_quantize_for_fa4, _fa4_v_to_fp8, flash_attn_fp4_func = _resolve_fa4_route_ops() orig_seqlen_q = query.shape[1] orig_seqlen_k = key.shape[1] @@ -291,6 +343,14 @@ def _forward_fa4_fp4( q_fp4, q_sf = _nvfp4_quantize_for_fa4(query) k_fp4, k_sf = _nvfp4_quantize_for_fa4(key) + # fp8 PV: unscaled e4m3 cast (no mSFV/v_descale); output stays BF16. + if self.fa4_pv_mode == "fp8": + value = _fa4_v_to_fp8(value) + # Keep the once-log (and its global flag) out of compiled traces, + # matching the FLASH_ATTN backend's logging convention. + if not torch.compiler.is_compiling(): + _log_pv_dtype_once(value.dtype) + # 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] diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index d53f0d9b09..8ed7a55983 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -16,6 +16,7 @@ AttentionMetadata, AttentionMetadataBuilder, ) +from fastvideo.attention.backends.attn_qat_infer import (_default_fa4_pv_mode, validate_fa4_pv_mode) from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -141,6 +142,36 @@ def _nvfp4_quantize_for_fa4(tensor_4d: torch.Tensor) -> tuple[torch.Tensor, torc return torch.ops.fastvideo.nvfp4_quantize_fa4(tensor_4d) +# fp8 PV mode (fa4_pv_mode="fp8"): the FA4 kernel's plain-fp8 V contract is a +# bare elementwise e4m3 cast in the usual (batch, seqlen, nheads, headdim) +# headdim-contiguous layout -- no mSFV scale-factor tensor and no v_descale +# (those belong to the block-scaled fp4/mxfp8 PV modes; absent v_descale means +# an implicit dequant scale of 1.0). The kernel's output stays BF16 whenever +# block-scaled Q/K are enabled. A naive in-forward `.to(torch.float8_e4m3fn)` +# graph-breaks under torch.compile, so the cast gets the same custom-op +# boundary treatment as the quantize step above. +@torch.library.custom_op( + "fastvideo::fa4_v_to_fp8", + mutates_args=(), + device_types="cuda", +) +def _fa4_v_to_fp8_op(value: torch.Tensor) -> torch.Tensor: + return value.to(torch.float8_e4m3fn) + + +@torch.library.register_fake("fastvideo::fa4_v_to_fp8") +def _fa4_v_to_fp8_fake(value: torch.Tensor) -> torch.Tensor: + # `.to(dtype)` uses preserve_format, so the impl keeps the input's strides + # for dense tensors (and falls back to contiguous otherwise); empty_like's + # default preserve_format reproduces exactly that layout rule. + return torch.empty_like(value, dtype=torch.float8_e4m3fn) + + +def _fa4_v_to_fp8(value: torch.Tensor) -> torch.Tensor: + """Cast V to fp8 e4m3 for FA4's fp8 PV mode via the custom-op boundary.""" + return torch.ops.fastvideo.fa4_v_to_fp8(value) + + class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @@ -220,12 +251,18 @@ def __init__( self.causal = causal self.softmax_scale = softmax_scale self.nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4", False) or os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1" + # PV-mode knob for the FA4-FP4 path. Validated unconditionally so a + # typo fails at construction, not at the first forward on a Blackwell + # box. The default leaves behavior identical to before the knob. + # kwargs win; the FASTVIDEO_FA4_PV_MODE env bridge is the user-reachable + # fallback (same pattern as nvfp4_fa4 above). + self.fa4_pv_mode = validate_fa4_pv_mode(extra_impl_args.get("fa4_pv_mode") or _default_fa4_pv_mode()) if self.nvfp4_fa4: cap = torch.cuda.get_device_capability() assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}") assert _FA4_FP4_AVAILABLE, ("NVFP4 FA4 requires flash-attention-fp4 (flash_attn.cute). " "Install via instructions in docs/inference/optimizations.md") - logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk only)") + logger.info("NVFP4 FA4 enabled for FlashAttentionImpl (quant_qk, pv_mode=%s)", self.fa4_pv_mode) def forward( self, @@ -320,7 +357,8 @@ def _forward_impl( return output def _forward_nvfp4(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: - """FP4 flash attention with quantized Q and K, BF16 V.""" + """FP4 flash attention with quantized Q and K; V in BF16 (default) or + fp8 e4m3 per the fa4_pv_mode knob.""" orig_seqlen_q = query.shape[1] orig_seqlen_k = key.shape[1] @@ -328,6 +366,10 @@ def _forward_nvfp4(self, query: torch.Tensor, key: torch.Tensor, value: torch.Te q_fp4, q_sf = _nvfp4_quantize_for_fa4(query) k_fp4, k_sf = _nvfp4_quantize_for_fa4(key) + # fp8 PV: unscaled e4m3 cast (no mSFV/v_descale); output stays BF16. + if self.fa4_pv_mode == "fp8": + value = _fa4_v_to_fp8(value) + # Pass original seqlen to FA4 — the kernel handles non-multiple-of-128 # via boundary masking. FP4/SF data is padded to 128-multiple but FA4 # only attends to orig_seqlen positions, avoiding softmax bias on padding. diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py index ca38539924..8b0125a86b 100644 --- a/fastvideo/attention/utils/flash_attn_cute.py +++ b/fastvideo/attention/utils/flash_attn_cute.py @@ -359,10 +359,12 @@ def _flash_attn_cute_fp4_forward_fake( causal: bool, ) -> torch.Tensor: del k, sfq, sfk, softmax_scale, causal - # q is FP4 packed: shape (batch, seqlen, nheads, headdim/2). Output is in - # V's dtype with full headdim. + # q is FP4 packed: shape (batch, seqlen, nheads, headdim/2). The + # block-scaled kernel always writes a BF16 output with full headdim -- + # including when V is fp8 e4m3 (pv_mode=fp8), so the fake must not follow + # V's dtype. batch, seqlen_q, nheads = q.shape[:3] - return v.new_empty(batch, seqlen_q, nheads, v.shape[-1]) + return v.new_empty(batch, seqlen_q, nheads, v.shape[-1], dtype=torch.bfloat16) def flash_attn_fp4_func( @@ -374,7 +376,9 @@ def flash_attn_fp4_func( softmax_scale: float | None = None, causal: bool = False, ) -> torch.Tensor: - """FP4 (NVFP4 block-scaled) flash attention. q/k are FP4-packed; v is BF16.""" + """FP4 (NVFP4 block-scaled) flash attention. q/k are FP4-packed; v is + BF16, or fp8 e4m3 as a plain unscaled cast (no mSFV scale factors and no + v_descale -- the kernel's plain-fp8 PV contract). Output is BF16.""" return torch.ops.fastvideo._flash_attn_cute_fp4_forward(q, k, v, sfq, sfk, softmax_scale, causal) diff --git a/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py b/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py index c13c9a55de..b6885e5b37 100644 --- a/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py +++ b/fastvideo/tests/attention/test_attn_qat_infer_arch_gate.py @@ -43,6 +43,7 @@ def test_arch_resolution(monkeypatch, cap, cutlass, fa4, expected_kernel) -> Non def test_receipt_records_fa4_quant_knobs(monkeypatch) -> None: _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True) + monkeypatch.setattr(aqi, "_configured_fa4_pv_mode", "bf16") receipt = aqi.attn_qat_infer_receipt() assert "arch=sm_100" in receipt assert "qk_mode=nvfp4(per-16-e4m3-sf)" in receipt @@ -50,6 +51,32 @@ def test_receipt_records_fa4_quant_knobs(monkeypatch) -> None: assert "train_sim_mismatch=measured" in receipt +def test_receipt_pv_mode_is_derived_from_configured_knob(monkeypatch) -> None: + """The receipt's pv_mode field derives from the fa4_pv_mode knob recorded + at impl construction, not a literal.""" + _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True) + monkeypatch.setattr(aqi, "_configured_fa4_pv_mode", "bf16") + monkeypatch.setattr(aqi, "_receipt_logged", False) + + aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5, fa4_pv_mode="fp8") + assert "pv_mode=fp8" in aqi.attn_qat_infer_receipt() + + +@pytest.mark.parametrize("impl_module", ["attn_qat_infer", "flash_attn"]) +def test_fa4_pv_mode_typo_fails_at_construction(monkeypatch, impl_module) -> None: + """Both knob consumers validate fa4_pv_mode eagerly: a bad value raises at + impl construction (CPU-only), never surviving to a GPU forward.""" + monkeypatch.setattr(aqi, "_configured_fa4_pv_mode", aqi._configured_fa4_pv_mode) + if impl_module == "attn_qat_infer": + impl_cls = aqi.AttnQatInferImpl + else: + from fastvideo.attention.backends.flash_attn import FlashAttentionImpl + impl_cls = FlashAttentionImpl + + with pytest.raises(ValueError, match="fa4_pv_mode"): + impl_cls(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5, fa4_pv_mode="fp16") + + def test_receipt_records_cutlass_scheme(monkeypatch) -> None: _patch(monkeypatch, cap=(12, 0), cutlass=True, fa4=False) receipt = aqi.attn_qat_infer_receipt() @@ -102,10 +129,14 @@ def test_fa4_route_resolution_runs_once_across_forwards(monkeypatch) -> None: def fake_quant(t): return t, torch.zeros(1) + def fake_cast(t): + return t + 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, "_import_fa4_route_ops", + lambda: (resolves.append(1) or (fake_quant, fake_cast, 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) @@ -182,3 +213,88 @@ def test_fa4_quantize_op_fake_matches_real() -> None: (x,), test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), ) + + +def _register_fa4_v_to_fp8_cpu_kernel(): + """Register a CPU kernel for fastvideo::fa4_v_to_fp8. Unlike the quantize + op this needs no mock: the production impl is a plain dtype cast, which + CPU torch executes natively, so the CPU kernel IS the production body.""" + import torch + + def _cpu_kernel(value: torch.Tensor) -> torch.Tensor: + return value.to(torch.float8_e4m3fn) + + try: + torch.library.register_kernel("fastvideo::fa4_v_to_fp8", "cpu")(_cpu_kernel) + except RuntimeError: + pass # already registered by a previous test/run + + +def test_fa4_v_to_fp8_op_fake_matches_real() -> None: + """torch.library.opcheck for the fp8 V-cast op: the fake must reproduce + the impl's shape, dtype, AND strides (`.to` preserves the input layout for + dense tensors; empty_like's preserve_format mirrors that). Forward-only op + (no autograd registration), so restrict to the non-autograd suites.""" + import torch + + import fastvideo.attention.backends.flash_attn # noqa: F401 registers the op + fake + + _register_fa4_v_to_fp8_cpu_kernel() + + for v in ( + torch.randn(2, 200, 2, 128, dtype=torch.bfloat16), + # Non-contiguous dense input: strides must carry through the cast. + torch.randn(2, 2, 200, 128, dtype=torch.bfloat16).transpose(1, 2), + ): + torch.library.opcheck( + torch.ops.fastvideo.fa4_v_to_fp8, + (v,), + test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), + ) + + +def test_fa4_fp8_pv_path_is_fullgraph_traceable(monkeypatch) -> None: + """With fa4_pv_mode="fp8" the production _forward_fa4_fp4 body (quantize + op + V-cast op + kernel call) must compile fullgraph with no graph break: + the naive in-forward `.to(torch.float8_e4m3fn)` cast breaks the graph, + which is exactly why the cast lives behind a custom op.""" + import torch + + from fastvideo.attention.backends import flash_attn as fa + + _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True) + monkeypatch.setattr(aqi, "_configured_fa4_pv_mode", aqi._configured_fa4_pv_mode) + _register_fa4_quantize_cpu_kernel() + _register_fa4_v_to_fp8_cpu_kernel() + + def fake_kernel(q, k, v, sfq, sfk, softmax_scale=None, causal=False): + # The real kernel emits BF16 with full headdim regardless of V dtype. + return v.to(torch.bfloat16) + + # Pre-resolve the route with the real op-backed helpers and a mocked + # kernel entry point (importing the real one needs a CUDA install). + monkeypatch.setattr(aqi, "_FA4_ROUTE_OPS", (fa._nvfp4_quantize_for_fa4, fa._fa4_v_to_fp8, fake_kernel)) + + impl = aqi.AttnQatInferImpl(num_heads=2, head_size=128, causal=False, softmax_scale=128**-0.5, fa4_pv_mode="fp8") + compiled = torch.compile(impl._forward_fa4_fp4, fullgraph=True, backend="eager") + x = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) + out = compiled(x, x, x) + assert out.shape == x.shape + assert out.dtype == torch.bfloat16 + + +def test_fa4_pv_mode_env_bridge(monkeypatch) -> None: + """The FASTVIDEO_FA4_PV_MODE env bridge is the user-reachable path to the + knob (model code constructs attention with fixed literals); explicit + kwargs still win over the environment.""" + import torch # noqa: F401 + + _patch(monkeypatch, cap=(10, 0), cutlass=False, fa4=True) + monkeypatch.setenv("FASTVIDEO_FA4_PV_MODE", "fp8") + + impl = aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5) + assert impl.fa4_pv_mode == "fp8" + + explicit = aqi.AttnQatInferImpl(num_heads=1, head_size=128, causal=False, softmax_scale=128**-0.5, + fa4_pv_mode="bf16") + assert explicit.fa4_pv_mode == "bf16"