@@ -54,29 +54,139 @@ 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+ _FA4_ROUTE_OPS : tuple | None = None
144+
145+
146+ def _import_fa4_route_ops () -> tuple :
147+ """Slow path (own function so tests pin it runs once per process):
148+ resolves the FA4 quantize helper and kernel entry point."""
149+ from fastvideo .attention .backends .flash_attn import (
150+ _nvfp4_quantize_for_fa4 , )
151+ from fastvideo .attention .utils .flash_attn_cute import (
152+ flash_attn_fp4_func , )
153+ return (_nvfp4_quantize_for_fa4 , flash_attn_fp4_func )
154+
155+
156+ def _resolve_fa4_route_ops () -> tuple :
157+ # Lazy but memoized: per-forward resolution graph-breaks dynamo every
158+ # step and blocks fullgraph compilation of the NVFP4 path.
159+ global _FA4_ROUTE_OPS
160+ if _FA4_ROUTE_OPS is None :
161+ _FA4_ROUTE_OPS = _import_fa4_route_ops ()
162+ return _FA4_ROUTE_OPS
163+
164+
165+ _receipt_logged = False
166+
167+
168+ def _log_receipt_once () -> None :
169+ # One line per process, not per layer (the validation swap constructs
170+ # one impl per attention layer).
171+ global _receipt_logged
172+ if not _receipt_logged :
173+ _receipt_logged = True
174+ logger .info ("ATTN_QAT_INFER resolved: %s" , attn_qat_infer_receipt ())
65175
66176
67177def 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.
178+ """True only when the active device has a built ATTN_QAT_INFER kernel.
70179
71180 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
181+ carry the sm_12x extension on any host (e.g. H100, GB200), where the
182+ import succeeds, backend selection picks this backend, and the first
183+ kernel call then fails with an unsupported-capability error instead of
184+ ever reaching the documented FlashAttention fallback in
76185 fastvideo.platforms.cuda. Gating on the active device's capability
77- keeps that fallback working on every non-sm_120/121 GPU.
186+ keeps that fallback working on every unsupported GPU, while
187+ sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).
78188 """
79- return _device_capability_supported () and _get_attn_qat_infer () is not None
189+ return _resolved_kernel () is not None
80190
81191
82192class AttnQatInferBackend (AttentionBackend ):
@@ -122,6 +232,12 @@ def __init__(
122232 if dropout_p > 0 :
123233 raise NotImplementedError (f"attn_qat_infer does not support dropout (got dropout_p={ dropout_p }
124234 "The QAT inference kernel applies no stochastic dropout." )
235+ # Kernel resolution is per-forward, not per-construction: callers
236+ # (the validation swap, backend selection) gate on
237+ # is_attn_qat_infer_available() first, and constructing an impl on a
238+ # host without the kernel must stay legal (pre-existing contract the
239+ # validation-swap test pins).
240+ _log_receipt_once ()
125241
126242 def forward (
127243 self ,
@@ -130,10 +246,18 @@ def forward(
130246 value : torch .Tensor ,
131247 attn_metadata : AttentionMetadata ,
132248 ) -> torch .Tensor :
249+ # Dispatch on the single per-arch resolution: importability of the
250+ # bundled sm_12x extension is NOT sufficient (CUDA 13 wheels carry it
251+ # on unsupported hosts, where calling it is the wrong binary).
252+ kernel = _resolved_kernel ()
253+ if kernel == "fa4_fp4" :
254+ return self ._forward_fa4_fp4 (query , key , value )
255+ if kernel is None :
256+ raise ImportError (f"attn_qat_infer is not available ({ attn_qat_infer_receipt ()}
257+ "Please ensure an ATTN_QAT_INFER kernel is installed for this device." )
258+
133259 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." )
260+ assert attn_qat_infer is not None # kernel == "cutlass_sm12x" implies the import succeeded
137261
138262 query = query .transpose (1 , 2 ).contiguous ()
139263 key = key .transpose (1 , 2 ).contiguous ()
@@ -148,3 +272,39 @@ def forward(
148272 sm_scale = self .softmax_scale ,
149273 )
150274 return output .transpose (1 , 2 ).contiguous ()
275+
276+ def _forward_fa4_fp4 (
277+ self ,
278+ query : torch .Tensor ,
279+ key : torch .Tensor ,
280+ value : torch .Tensor ,
281+ ) -> torch .Tensor :
282+ """sm_100a/sm_103a path: FP4 FA4 with the repo's tuned defaults
283+ (NVFP4 per-16 block-scaled Q/K, BF16 V) -- mirrors
284+ FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
285+ (batch, seqlen, nheads, headdim); no transpose."""
286+ _nvfp4_quantize_for_fa4 , flash_attn_fp4_func = _resolve_fa4_route_ops ()
287+
288+ orig_seqlen_q = query .shape [1 ]
289+ orig_seqlen_k = key .shape [1 ]
290+
291+ q_fp4 , q_sf = _nvfp4_quantize_for_fa4 (query )
292+ k_fp4 , k_sf = _nvfp4_quantize_for_fa4 (key )
293+
294+ # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
295+ # original lengths so padding never biases the softmax.
296+ q_fp4 = q_fp4 [:, :orig_seqlen_q ]
297+ k_fp4 = k_fp4 [:, :orig_seqlen_k ]
298+
299+ output = flash_attn_fp4_func (
300+ q_fp4 ,
301+ k_fp4 ,
302+ value ,
303+ q_sf ,
304+ k_sf ,
305+ softmax_scale = self .softmax_scale ,
306+ causal = self .causal ,
307+ )
308+ if isinstance (output , tuple ):
309+ output = output [0 ]
310+ return output
0 commit comments