[feat]: fp8 PV mode for the FA4-FP4 attention path

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)

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,18 +153,19 @@ 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()
     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")
+                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,18 +331,26 @@ 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]
 
         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]

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,14 +357,19 @@ 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]
 
         # Quantize Q/K to FP4 (internally pads to multiple of 128 for SF layout)
         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.

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)