@@ -54,29 +54,108 @@ def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
5454# kernel is compiled for (sm_120a / sm_121a -- see fastvideo-kernel/README.md).
5555_SUPPORTED_DEVICE_CAPABILITIES = frozenset ({(12 , 0 ), (12 , 1 )})
5656
57-
58- def _device_capability_supported () -> bool :
57+ # Datacenter-Blackwell capabilities served by the FP4 FA4 kernel
58+ # (flash-attention-fp4 @ fp4, sm_100a/sm_103a) through #1221's plumbing:
59+ # per-16 block-scaled NVFP4 Q/K (E4M3 scale factors), BF16 P/V. This is a
60+ # DIFFERENT quantization scheme from the sm_12x CUTLASS extension above --
61+ # ATTN_QAT_TRAIN simulates the CUTLASS scheme, so sm_100/sm_103 deployment
62+ # carries a train-sim mismatch that is measured (MS-SSIM gate), not assumed.
63+ _FA4_FP4_CAPABILITIES = frozenset ({(10 , 0 ), (10 , 3 )})
64+
65+ _FA4_INSTALL_HINT = ("pip install flash-attention-fp4 from "
66+ "https://github.com/hao-ai-lab/flash-attention-fp4 (branch fp4); "
67+ "see docs/inference/optimizations.md" )
68+
69+ _fa4_fp4_import_ok : bool | None = None
70+
71+
72+ def _fa4_fp4_available () -> bool :
73+ """flash_attn.cute (FA4) import probe, cached. Reuses #1221's guarded
74+ import chain in fastvideo.attention.utils.flash_attn_cute (which maps
75+ cutlass-dsl version skew to ImportError with a loud warning)."""
76+ global _fa4_fp4_import_ok
77+ if _fa4_fp4_import_ok is None :
78+ try :
79+ from fastvideo .attention .utils .flash_attn_cute import ( # noqa: F401
80+ flash_attn_fp4_func , )
81+ _fa4_fp4_import_ok = True
82+ except ImportError :
83+ _fa4_fp4_import_ok = False
84+ return _fa4_fp4_import_ok
85+
86+
87+ def _active_capability () -> tuple [int , int ] | None :
5988 if not torch .cuda .is_available ():
60- return False
89+ return None
6190 try :
62- return tuple (torch .cuda .get_device_capability ()) in _SUPPORTED_DEVICE_CAPABILITIES
91+ return tuple (torch .cuda .get_device_capability ())
6392 except Exception : # pragma: no cover - defensive: never break backend selection
64- return False
93+ return None
94+
95+
96+ def _resolved_kernel () -> str | None :
97+ """Which ATTN_QAT_INFER kernel serves the active device, or None.
98+
99+ Per-arch resolution (single source of truth -- extend the capability
100+ sets above, do not add equality checks elsewhere):
101+ * sm_12x consumer Blackwell -> fastvideo-kernel CUTLASS extension
102+ (modified SageAttention3 FP4).
103+ * sm_100a/sm_103a datacenter Blackwell -> FP4 FA4 (flash-attention-fp4)
104+ via the merged #1221 plumbing.
105+ """
106+ cap = _active_capability ()
107+ if cap in _SUPPORTED_DEVICE_CAPABILITIES and _get_attn_qat_infer () is not None :
108+ return "cutlass_sm12x"
109+ if cap in _FA4_FP4_CAPABILITIES and _fa4_fp4_available ():
110+ return "fa4_fp4"
111+ return None
112+
113+
114+ def attn_qat_infer_receipt () -> str :
115+ """One-line receipt of the resolution decision (arch + kernel + quant
116+ knobs), for the selection log and for tooling. The FA4 knobs are the
117+ repo's tuned defaults passed through verbatim: qk_mode=nvfp4
118+ (per-16 E4M3 SFs), pv_mode=bf16 -- see flash_attn/cute/README.md in
119+ the kernel repo."""
120+ cap = _active_capability ()
121+ arch = f"sm_{ cap [0 ]} { cap [1 ]} if cap is not None else "no-cuda"
122+ kernel = _resolved_kernel ()
123+ if kernel == "cutlass_sm12x" :
124+ return f"arch={ arch }
125+ if kernel == "fa4_fp4" :
126+ return (f"arch={ arch }
127+ f"pv_mode=bf16 train_sim_mismatch=measured" )
128+ supported = "sm_120a/sm_121a via fastvideo-kernel build.sh; sm_100a/sm_103a via flash-attention-fp4"
129+ if cap is not None and cap in _FA4_FP4_CAPABILITIES :
130+ return f"arch={ arch } { _FA4_INSTALL_HINT }
131+ return f"arch={ arch } { supported }
132+
133+
134+ _receipt_logged = False
135+
136+
137+ def _log_receipt_once () -> None :
138+ # One line per process, not per layer (the validation swap constructs
139+ # one impl per attention layer).
140+ global _receipt_logged
141+ if not _receipt_logged :
142+ _receipt_logged = True
143+ logger .info ("ATTN_QAT_INFER resolved: %s" , attn_qat_infer_receipt ())
65144
66145
67146def is_attn_qat_infer_available () -> bool :
68- """True only when the extension imports AND the active device is a
69- consumer-Blackwell (sm_120/sm_121) GPU the kernel is compiled for.
147+ """True only when the active device has a built ATTN_QAT_INFER kernel.
70148
71149 The import check alone is not sufficient: CUDA 13 wheel builds can
72- carry the sm_120/sm_121 extension on any host (e.g. H100, GB200),
73- where the import succeeds, backend selection picks this backend, and
74- the first kernel call then fails with an unsupported-capability error
75- instead of ever reaching the documented FlashAttention fallback in
150+ carry the sm_12x extension on any host (e.g. H100, GB200), where the
151+ import succeeds, backend selection picks this backend, and the first
152+ kernel call then fails with an unsupported-capability error instead of
153+ ever reaching the documented FlashAttention fallback in
76154 fastvideo.platforms.cuda. Gating on the active device's capability
77- keeps that fallback working on every non-sm_120/121 GPU.
155+ keeps that fallback working on every unsupported GPU, while
156+ sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).
78157 """
79- return _device_capability_supported () and _get_attn_qat_infer () is not None
158+ return _resolved_kernel () is not None
80159
81160
82161class AttnQatInferBackend (AttentionBackend ):
@@ -122,6 +201,10 @@ def __init__(
122201 if dropout_p > 0 :
123202 raise NotImplementedError (f"attn_qat_infer does not support dropout (got dropout_p={ dropout_p }
124203 "The QAT inference kernel applies no stochastic dropout." )
204+ self ._kernel = _resolved_kernel ()
205+ if self ._kernel is None :
206+ raise ImportError (f"attn_qat_infer is not available ({ attn_qat_infer_receipt ()} )
207+ _log_receipt_once ()
125208
126209 def forward (
127210 self ,
@@ -130,6 +213,9 @@ def forward(
130213 value : torch .Tensor ,
131214 attn_metadata : AttentionMetadata ,
132215 ) -> torch .Tensor :
216+ if self ._kernel == "fa4_fp4" :
217+ return self ._forward_fa4_fp4 (query , key , value )
218+
133219 attn_qat_infer = _get_attn_qat_infer ()
134220 if attn_qat_infer is None :
135221 raise ImportError ("attn_qat_infer is not available. Please ensure the "
@@ -148,3 +234,42 @@ def forward(
148234 sm_scale = self .softmax_scale ,
149235 )
150236 return output .transpose (1 , 2 ).contiguous ()
237+
238+ def _forward_fa4_fp4 (
239+ self ,
240+ query : torch .Tensor ,
241+ key : torch .Tensor ,
242+ value : torch .Tensor ,
243+ ) -> torch .Tensor :
244+ """sm_100a/sm_103a path: FP4 FA4 with the repo's tuned defaults
245+ (NVFP4 per-16 block-scaled Q/K, BF16 V) -- mirrors
246+ FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
247+ (batch, seqlen, nheads, headdim); no transpose."""
248+ from fastvideo .attention .backends .flash_attn import (
249+ _nvfp4_quantize_for_fa4 , )
250+ from fastvideo .attention .utils .flash_attn_cute import (
251+ flash_attn_fp4_func , )
252+
253+ orig_seqlen_q = query .shape [1 ]
254+ orig_seqlen_k = key .shape [1 ]
255+
256+ q_fp4 , q_sf = _nvfp4_quantize_for_fa4 (query )
257+ k_fp4 , k_sf = _nvfp4_quantize_for_fa4 (key )
258+
259+ # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
260+ # original lengths so padding never biases the softmax.
261+ q_fp4 = q_fp4 [:, :orig_seqlen_q ]
262+ k_fp4 = k_fp4 [:, :orig_seqlen_k ]
263+
264+ output = flash_attn_fp4_func (
265+ q_fp4 ,
266+ k_fp4 ,
267+ value ,
268+ q_sf ,
269+ k_sf ,
270+ softmax_scale = self .softmax_scale ,
271+ causal = self .causal ,
272+ )
273+ if isinstance (output , tuple ):
274+ output = output [0 ]
275+ return output
0 commit comments