diff --git a/fastvideo/attention/backends/attn_qat_infer.py b/fastvideo/attention/backends/attn_qat_infer.py index cb71c3764..f0030962a 100644 --- a/fastvideo/attention/backends/attn_qat_infer.py +++ b/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})" diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index eac762356..295ef6631 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/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 diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py index 7a7c09fbc..728709458 100644 --- a/fastvideo/attention/utils/flash_attn_cute.py +++ b/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. diff --git a/fastvideo/tests/attention/test_nvfp4_fused_qk_quant.py b/fastvideo/tests/attention/test_nvfp4_fused_qk_quant.py index 906906638..d4d0d6901 100644 --- a/fastvideo/tests/attention/test_nvfp4_fused_qk_quant.py +++ b/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