[bugfix] Gate the FA4 NVFP4 path on the installed kernel, not the import (#2067)

Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com>

fastvideo/attention/backends/attn_qat_infer.py

@@ -134,15 +134,16 @@ _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)."""
+    """FP4 FA4 kernel 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) and requires the installed kernel
+    to accept the NVFP4 scale-factor kwargs: upstream flash-attn-4 imports but
+    rejects mSFQ/mSFK, so only the flash-attention-fp4 fork can serve this route."""
     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
+            from fastvideo.attention.utils.flash_attn_cute import fp4_forward_available
+            _fa4_fp4_import_ok = fp4_forward_available()
         except ImportError:
             _fa4_fp4_import_ok = False
     return _fa4_fp4_import_ok
@@ -193,7 +194,7 @@ def attn_qat_infer_receipt() -> str:
                 f"pv_mode={_current_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})"
+        return f"arch={arch} kernel=none (no NVFP4 flash_attn.cute forward -- {_FA4_INSTALL_HINT})"
     return f"arch={arch} kernel=none (supported: {supported})"
 
 

fastvideo/attention/backends/flash_attn.py

@@ -49,8 +49,11 @@ logger.info("Worker %s Using FlashAttention-%s backend",
 # The FP4 path uses a dedicated custom_op wrapper (flash_attn_fp4_func) so that
 # torch.compile treats the CuTeDSL kernel as an opaque boundary.
 try:
-    from fastvideo.attention.utils.flash_attn_cute import flash_attn_fp4_func
-    _FA4_FP4_AVAILABLE = True
+    from fastvideo.attention.utils.flash_attn_cute import flash_attn_fp4_func, fp4_forward_available
+    # flash_attn.cute importing is not enough: the FP4 forward passes mSFQ/mSFK, which
+    # only the flash-attention-fp4 fork's kernel accepts, not the upstream FA4 build the
+    # repo pins. Gate on the installed kernel.
+    _FA4_FP4_AVAILABLE = fp4_forward_available()
 except ImportError:
     flash_attn_fp4_func = None
     _FA4_FP4_AVAILABLE = False

fastvideo/attention/utils/flash_attn_cute.py

@@ -46,6 +46,18 @@ def _accepted(params: frozenset[str], **kwargs) -> dict:
     return {name: value for name, value in kwargs.items() if name in params}
 
 
+# The block-scaled NVFP4 forward (the mSFQ/mSFK scale-factor tensors) only exists in the
+# hao-ai-lab/flash-attention-fp4 fork; the upstream FA4 build the repo's flash-attn-4 pin
+# installs accepts neither kwarg. Probe the signature so callers gate the FP4 paths on the
+# kernel actually installed, not on flash_attn.cute merely importing.
+_FP4_FORWARD_KWARGS = ("mSFQ", "mSFK")
+
+
+def fp4_forward_available() -> bool:
+    """The installed flash_attn.cute accepts the NVFP4 block-scaled forward."""
+    return all(name in _FWD_PARAMS for name in _FP4_FORWARD_KWARGS)
+
+
 try:
     # FA2 serves the calls FA4 cute cannot on pre-sm90 GPUs (backward, GQA).
     # Optional so FA4-only installs can still import this module.

fastvideo/tests/attention/test_nvfp4_fused_qk_quant.py

@@ -6,6 +6,7 @@ scale-factor bytes in FlashAttention-4's layout, and therefore the same FA4 NVFP
 ``precondition_qk`` followed by ``_nvfp4_quantize_for_fa4``. GPU tests need Blackwell (sm_100a/sm_103a), flashinfer
 and the FA4 NVFP4 kernel; the selection tests run anywhere.
 """
+import inspect
 from unittest.mock import patch
 
 import pytest
@@ -210,3 +211,17 @@ def test_fused_path_env_opt_out(env_overrides, fa4_runtime_available) -> None:
 def test_fused_path_off_for_non_default_flashinfer_recipe(env_overrides) -> None:
     env_overrides.enter_context(envs.override_external("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH", "1"))
     assert fused.fused_qk_quant_available() is False
+
+
+def test_fp4_availability_tracks_the_installed_kernel() -> None:
+    """flash_attn.cute importing is not enough to run the FP4 path.
+
+    The repo's flash-attn-4 pin is upstream FA4, whose kernel rejects the mSFQ/mSFK scale-factor kwargs the FP4
+    forward passes; only the flash-attention-fp4 fork accepts them. A build with the upstream kernel must not
+    report the FP4 path available (and so must skip the GPU tests above), or its unit lane dies with TypeError.
+    """
+    flash_attn_cute = pytest.importorskip("fastvideo.attention.utils.flash_attn_cute",
+                                          reason="no usable flash-attention package installed")
+    kernel_accepts_fp4 = {"mSFQ", "mSFK"} <= set(inspect.signature(flash_attn_cute._flash_attn_fwd).parameters)
+    assert flash_attn_cute.fp4_forward_available() is kernel_accepts_fp4
+    assert flash_attn_module._FA4_FP4_AVAILABLE is kernel_accepts_fp4