Skip to content

Commit 5b51de5

Browse files
committed
[test]: reconcile attn_qat_infer capability test with #1647 FA4 sm_100/103
#1647 extended ATTN_QAT_INFER to datacenter Blackwell (sm_100/sm_103) via the FP4 FA4 path, so is_attn_qat_infer_available() now resolves True on (10,0)/(10,3) too — the pre-#1647 assertion (10,0)->False is stale. Parametrize over both kernel paths and force both probes "installed" (_get_attn_qat_infer for sm_12x CUTLASS, _fa4_fp4_available for sm_100/103 FA4) so the test isolates the capability gate deterministically regardless of what's importable on the runner.
1 parent 2306230 commit 5b51de5

1 file changed

Lines changed: 18 additions & 1 deletion

File tree

fastvideo/tests/attention/test_attn_qat_infer_capability.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,28 @@
66

77
@pytest.mark.parametrize(
88
"capability,expected",
9-
[((12, 0), True), ((12, 1), True), ((10, 0), False), ((9, 0), False)],
9+
[
10+
((12, 0), True), # sm_120 consumer Blackwell -> fastvideo-kernel CUTLASS (SageAttn3 FP4)
11+
((12, 1), True), # sm_121 (DGX Spark) -> fastvideo-kernel CUTLASS (SageAttn3 FP4)
12+
((10, 0), True), # sm_100 datacenter Blackwell -> FP4 FA4 (flash-attention-fp4), #1647
13+
((10, 3), True), # sm_103 datacenter Blackwell -> FP4 FA4 (flash-attention-fp4), #1647
14+
((9, 0), False), # sm_90 Hopper -> unsupported (falls back to Flash)
15+
((8, 9), False), # sm_89 Ada/L40S -> unsupported (falls back to Flash)
16+
],
1017
)
1118
def test_attn_qat_infer_availability_requires_supported_gpu(monkeypatch, capability, expected):
19+
"""Availability is gated on the active device's compute capability, so a
20+
wheel that carries the kernel on an unsupported host (e.g. import succeeds
21+
on Hopper/Ada) still falls back to Flash instead of crashing at kernel call.
22+
23+
Both kernel probes are forced "installed" so the assertion isolates the
24+
capability gate, independent of what's importable on the CI runner:
25+
* ``_get_attn_qat_infer`` -> the sm_12x CUTLASS extension
26+
* ``_fa4_fp4_available`` -> the sm_100/sm_103 FP4 FA4 path (#1647)
27+
"""
1228
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
1329
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability)
1430
monkeypatch.setattr(attn_qat_infer, "_get_attn_qat_infer", lambda: object())
31+
monkeypatch.setattr(attn_qat_infer, "_fa4_fp4_available", lambda: True)
1532

1633
assert attn_qat_infer.is_attn_qat_infer_available() is expected

0 commit comments

Comments
 (0)