@@ -54,29 +54,117 @@ 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+ # The fork is written against the cutlass-dsl 4.4 API surface; the validated
66+ # install set (GB200-proven) is nvidia-cutlass-dsl==4.4.2 +
67+ # nvidia-cutlass-dsl-libs-base==4.4.2 + quack-kernels==0.4.1 +
68+ # flashinfer-python==0.6.8, with the fork on PYTHONPATH,
69+ # CUTE_DSL_ENABLE_TVM_FFI=1, and FASTVIDEO_FA4=1 (the fork ships no compiled
70+ # FA2, so dense attention paths need the FA4 opt-in). dsl 4.6-era installs
71+ # fail at CuTe JIT trace (cute.make_fragment was removed at module level).
72+ _FA4_INSTALL_HINT = ("install flash-attention-fp4 (branch fp4) from "
73+ "https://github.com/hao-ai-lab/flash-attention-fp4 with "
74+ "nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1, "
75+ "flashinfer-python==0.6.8 and FASTVIDEO_FA4=1; "
76+ "see docs/inference/optimizations.md" )
77+
78+ _fa4_fp4_import_ok : bool | None = None
79+
80+
81+ def _fa4_fp4_available () -> bool :
82+ """flash_attn.cute (FA4) import probe, cached. Reuses #1221's guarded
83+ import chain in fastvideo.attention.utils.flash_attn_cute (which maps
84+ cutlass-dsl version skew to ImportError with a loud warning)."""
85+ global _fa4_fp4_import_ok
86+ if _fa4_fp4_import_ok is None :
87+ try :
88+ from fastvideo .attention .utils .flash_attn_cute import ( # noqa: F401
89+ flash_attn_fp4_func , )
90+ _fa4_fp4_import_ok = True
91+ except ImportError :
92+ _fa4_fp4_import_ok = False
93+ return _fa4_fp4_import_ok
94+
95+
96+ def _active_capability () -> tuple [int , int ] | None :
5997 if not torch .cuda .is_available ():
60- return False
98+ return None
6199 try :
62- return tuple (torch .cuda .get_device_capability ()) in _SUPPORTED_DEVICE_CAPABILITIES
100+ return tuple (torch .cuda .get_device_capability ())
63101 except Exception : # pragma: no cover - defensive: never break backend selection
64- return False
102+ return None
103+
104+
105+ def _resolved_kernel () -> str | None :
106+ """Which ATTN_QAT_INFER kernel serves the active device, or None.
107+
108+ Per-arch resolution (single source of truth -- extend the capability
109+ sets above, do not add equality checks elsewhere):
110+ * sm_12x consumer Blackwell -> fastvideo-kernel CUTLASS extension
111+ (modified SageAttention3 FP4).
112+ * sm_100a/sm_103a datacenter Blackwell -> FP4 FA4 (flash-attention-fp4)
113+ via the merged #1221 plumbing.
114+ """
115+ cap = _active_capability ()
116+ if cap in _SUPPORTED_DEVICE_CAPABILITIES and _get_attn_qat_infer () is not None :
117+ return "cutlass_sm12x"
118+ if cap in _FA4_FP4_CAPABILITIES and _fa4_fp4_available ():
119+ return "fa4_fp4"
120+ return None
121+
122+
123+ def attn_qat_infer_receipt () -> str :
124+ """One-line receipt of the resolution decision (arch + kernel + quant
125+ knobs), for the selection log and for tooling. The FA4 knobs are the
126+ repo's tuned defaults passed through verbatim: qk_mode=nvfp4
127+ (per-16 E4M3 SFs), pv_mode=bf16 -- see flash_attn/cute/README.md in
128+ the kernel repo."""
129+ cap = _active_capability ()
130+ arch = f"sm_{ cap [0 ]} { cap [1 ]} if cap is not None else "no-cuda"
131+ kernel = _resolved_kernel ()
132+ if kernel == "cutlass_sm12x" :
133+ return f"arch={ arch }
134+ if kernel == "fa4_fp4" :
135+ return (f"arch={ arch }
136+ f"pv_mode=bf16 train_sim_mismatch=measured" )
137+ supported = "sm_120a/sm_121a via fastvideo-kernel build.sh; sm_100a/sm_103a via flash-attention-fp4"
138+ if cap is not None and cap in _FA4_FP4_CAPABILITIES :
139+ return f"arch={ arch } { _FA4_INSTALL_HINT }
140+ return f"arch={ arch } { supported }
141+
142+
143+ _receipt_logged = False
144+
145+
146+ def _log_receipt_once () -> None :
147+ # One line per process, not per layer (the validation swap constructs
148+ # one impl per attention layer).
149+ global _receipt_logged
150+ if not _receipt_logged :
151+ _receipt_logged = True
152+ logger .info ("ATTN_QAT_INFER resolved: %s" , attn_qat_infer_receipt ())
65153
66154
67155def 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.
156+ """True only when the active device has a built ATTN_QAT_INFER kernel.
70157
71158 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
159+ carry the sm_12x extension on any host (e.g. H100, GB200), where the
160+ import succeeds, backend selection picks this backend, and the first
161+ kernel call then fails with an unsupported-capability error instead of
162+ ever reaching the documented FlashAttention fallback in
76163 fastvideo.platforms.cuda. Gating on the active device's capability
77- keeps that fallback working on every non-sm_120/121 GPU.
164+ keeps that fallback working on every unsupported GPU, while
165+ sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).
78166 """
79- return _device_capability_supported () and _get_attn_qat_infer () is not None
167+ return _resolved_kernel () is not None
80168
81169
82170class AttnQatInferBackend (AttentionBackend ):
@@ -122,6 +210,12 @@ def __init__(
122210 if dropout_p > 0 :
123211 raise NotImplementedError (f"attn_qat_infer does not support dropout (got dropout_p={ dropout_p }
124212 "The QAT inference kernel applies no stochastic dropout." )
213+ # Kernel resolution is per-forward, not per-construction: callers
214+ # (the validation swap, backend selection) gate on
215+ # is_attn_qat_infer_available() first, and constructing an impl on a
216+ # host without the kernel must stay legal (pre-existing contract the
217+ # validation-swap test pins).
218+ _log_receipt_once ()
125219
126220 def forward (
127221 self ,
@@ -130,10 +224,18 @@ def forward(
130224 value : torch .Tensor ,
131225 attn_metadata : AttentionMetadata ,
132226 ) -> torch .Tensor :
227+ # Dispatch on the single per-arch resolution: importability of the
228+ # bundled sm_12x extension is NOT sufficient (CUDA 13 wheels carry it
229+ # on unsupported hosts, where calling it is the wrong binary).
230+ kernel = _resolved_kernel ()
231+ if kernel == "fa4_fp4" :
232+ return self ._forward_fa4_fp4 (query , key , value )
233+ if kernel is None :
234+ raise ImportError (f"attn_qat_infer is not available ({ attn_qat_infer_receipt ()}
235+ "Please ensure an ATTN_QAT_INFER kernel is installed for this device." )
236+
133237 attn_qat_infer = _get_attn_qat_infer ()
134- if attn_qat_infer is None :
135- raise ImportError ("attn_qat_infer is not available. Please ensure the "
136- "attn_qat_infer kernel package is installed." )
238+ assert attn_qat_infer is not None # kernel == "cutlass_sm12x" implies the import succeeded
137239
138240 query = query .transpose (1 , 2 ).contiguous ()
139241 key = key .transpose (1 , 2 ).contiguous ()
@@ -148,3 +250,42 @@ def forward(
148250 sm_scale = self .softmax_scale ,
149251 )
150252 return output .transpose (1 , 2 ).contiguous ()
253+
254+ def _forward_fa4_fp4 (
255+ self ,
256+ query : torch .Tensor ,
257+ key : torch .Tensor ,
258+ value : torch .Tensor ,
259+ ) -> torch .Tensor :
260+ """sm_100a/sm_103a path: FP4 FA4 with the repo's tuned defaults
261+ (NVFP4 per-16 block-scaled Q/K, BF16 V) -- mirrors
262+ FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
263+ (batch, seqlen, nheads, headdim); no transpose."""
264+ from fastvideo .attention .backends .flash_attn import (
265+ _nvfp4_quantize_for_fa4 , )
266+ from fastvideo .attention .utils .flash_attn_cute import (
267+ flash_attn_fp4_func , )
268+
269+ orig_seqlen_q = query .shape [1 ]
270+ orig_seqlen_k = key .shape [1 ]
271+
272+ q_fp4 , q_sf = _nvfp4_quantize_for_fa4 (query )
273+ k_fp4 , k_sf = _nvfp4_quantize_for_fa4 (key )
274+
275+ # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
276+ # original lengths so padding never biases the softmax.
277+ q_fp4 = q_fp4 [:, :orig_seqlen_q ]
278+ k_fp4 = k_fp4 [:, :orig_seqlen_k ]
279+
280+ output = flash_attn_fp4_func (
281+ q_fp4 ,
282+ k_fp4 ,
283+ value ,
284+ q_sf ,
285+ k_sf ,
286+ softmax_scale = self .softmax_scale ,
287+ causal = self .causal ,
288+ )
289+ if isinstance (output , tuple ):
290+ output = output [0 ]
291+ return output
0 commit comments