@@ -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,13 @@ def forward(
130224 value : torch .Tensor ,
131225 attn_metadata : AttentionMetadata ,
132226 ) -> torch .Tensor :
227+ if _resolved_kernel () == "fa4_fp4" :
228+ return self ._forward_fa4_fp4 (query , key , value )
229+
133230 attn_qat_infer = _get_attn_qat_infer ()
134231 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." )
232+ raise ImportError (f "attn_qat_infer is not available ( { attn_qat_infer_receipt () } ). "
233+ "Please ensure an ATTN_QAT_INFER kernel is installed for this device ." )
137234
138235 query = query .transpose (1 , 2 ).contiguous ()
139236 key = key .transpose (1 , 2 ).contiguous ()
@@ -148,3 +245,42 @@ def forward(
148245 sm_scale = self .softmax_scale ,
149246 )
150247 return output .transpose (1 , 2 ).contiguous ()
248+
249+ def _forward_fa4_fp4 (
250+ self ,
251+ query : torch .Tensor ,
252+ key : torch .Tensor ,
253+ value : torch .Tensor ,
254+ ) -> torch .Tensor :
255+ """sm_100a/sm_103a path: FP4 FA4 with the repo's tuned defaults
256+ (NVFP4 per-16 block-scaled Q/K, BF16 V) -- mirrors
257+ FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
258+ (batch, seqlen, nheads, headdim); no transpose."""
259+ from fastvideo .attention .backends .flash_attn import (
260+ _nvfp4_quantize_for_fa4 , )
261+ from fastvideo .attention .utils .flash_attn_cute import (
262+ flash_attn_fp4_func , )
263+
264+ orig_seqlen_q = query .shape [1 ]
265+ orig_seqlen_k = key .shape [1 ]
266+
267+ q_fp4 , q_sf = _nvfp4_quantize_for_fa4 (query )
268+ k_fp4 , k_sf = _nvfp4_quantize_for_fa4 (key )
269+
270+ # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
271+ # original lengths so padding never biases the softmax.
272+ q_fp4 = q_fp4 [:, :orig_seqlen_q ]
273+ k_fp4 = k_fp4 [:, :orig_seqlen_k ]
274+
275+ output = flash_attn_fp4_func (
276+ q_fp4 ,
277+ k_fp4 ,
278+ value ,
279+ q_sf ,
280+ k_sf ,
281+ softmax_scale = self .softmax_scale ,
282+ causal = self .causal ,
283+ )
284+ if isinstance (output , tuple ):
285+ output = output [0 ]
286+ return output
0 commit comments