"""Test NVFP4 FA4 (FP4 quantized QK flash attention) vs BF16 FA4. Requires: Blackwell GPU (sm100a/sm103a), flash-attention-fp4, flashinfer, cutlass-dsl. Run: CUTE_DSL_ENABLE_TVM_FFI=1 pytest tests/local_tests/test_nvfp4_fa4.py -v -s """ import os import torch import pytest os.environ["CUTE_DSL_ENABLE_TVM_FFI"] = "1" def _make_impl(nheads, headdim, nvfp4=False): from fastvideo.attention.backends.flash_attn import FlashAttentionImpl return FlashAttentionImpl( num_heads=nheads, head_size=headdim, causal=False, softmax_scale=headdim**-0.5, nvfp4_fa4=nvfp4, ) def _cuda_timer(fn, warmup=5, iters=20): """Time a function using CUDA events.""" for _ in range(warmup): fn() torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(iters): fn() end.record() torch.cuda.synchronize() return start.elapsed_time(end) / iters # ms @pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)], reason="Requires Blackwell GPU (sm100a or sm103a)") class TestNVFP4FA4: # Wan2.1-T2V-1.3B: num_attention_heads=12, attention_head_dim=128 MODEL_NHEADS = 12 MODEL_HEADDIM = 128 MODEL_SEQLEN = 32760 # 480x832 video, 81 frames @pytest.fixture(scope="class", autouse=True) def _flashinfer_warmup(self): # Warm up flashinfer JIT before any FA4 imports from flashinfer.quantization import SfLayout, nvfp4_quantize _warmup = nvfp4_quantize( torch.randn(4, 128, device="cuda", dtype=torch.bfloat16), torch.ones(1, device="cuda", dtype=torch.float32), sfLayout=SfLayout.layout_128x4, do_shuffle=False, ) del _warmup def test_quantize_fn_shapes(self): """_nvfp4_quantize_for_fa4 produces correct shapes and strides.""" from fastvideo.attention.backends.flash_attn import _nvfp4_quantize_for_fa4 batch, seqlen, nheads, headdim = 1, 256, self.MODEL_NHEADS, self.MODEL_HEADDIM t = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) fp4, sf = _nvfp4_quantize_for_fa4(t) assert fp4.dtype == torch.float4_e2m1fn_x2 assert fp4.shape == (batch, seqlen, nheads, headdim // 2) rest_m = seqlen // 128 rest_k = (headdim // 16) // 4 assert sf.shape == (32, 4, rest_m, 4, rest_k, nheads, batch) assert sf.dtype == torch.uint8 assert sf.stride()[3] == 1, f"SF stride[3] must be 1, got {sf.stride()[3]}" def test_output_shape_and_dtype(self): """NVFP4 FA4 output has correct shape and dtype.""" batch, nheads, headdim = 1, self.MODEL_NHEADS, self.MODEL_HEADDIM seqlen = 256 q = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) k = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) v = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) impl = _make_impl(nheads, headdim, nvfp4=True) out = impl.forward(q, k, v, attn_metadata=None) assert out.shape == (batch, seqlen, nheads, headdim) assert out.dtype == torch.bfloat16 assert not torch.isnan(out).any(), "Output contains NaN" def test_cross_attention(self): """NVFP4 FA4 handles cross-attention (seqlen_q != seqlen_k).""" nheads, headdim = self.MODEL_NHEADS, self.MODEL_HEADDIM q = torch.randn(1, self.MODEL_SEQLEN, nheads, headdim, device="cuda", dtype=torch.bfloat16) k = torch.randn(1, 512, nheads, headdim, device="cuda", dtype=torch.bfloat16) v = torch.randn(1, 512, nheads, headdim, device="cuda", dtype=torch.bfloat16) impl = _make_impl(nheads, headdim, nvfp4=True) out = impl.forward(q, k, v, attn_metadata=None) assert out.shape == (1, self.MODEL_SEQLEN, nheads, headdim) assert not torch.isnan(out).any() def test_vs_bf16_accuracy(self): """NVFP4 FA4 output is close to BF16 FA4.""" nheads, headdim = self.MODEL_NHEADS, self.MODEL_HEADDIM from flash_attn.cute.interface import flash_attn_func as fa4_func torch.manual_seed(42) q = torch.randn(1, self.MODEL_SEQLEN, nheads, headdim, device="cuda", dtype=torch.bfloat16) k = torch.randn(1, self.MODEL_SEQLEN, nheads, headdim, device="cuda", dtype=torch.bfloat16) v = torch.randn(1, self.MODEL_SEQLEN, nheads, headdim, device="cuda", dtype=torch.bfloat16) fp4_out = _make_impl(nheads, headdim, nvfp4=True).forward(q, k, v, attn_metadata=None) bf16_out = fa4_func(q, k, v, softmax_scale=headdim**-0.5, causal=False) if isinstance(bf16_out, tuple): bf16_out = bf16_out[0] cos = torch.nn.functional.cosine_similarity(fp4_out.flatten().unsqueeze(0).float(), bf16_out.flatten().unsqueeze(0).float()).item() print(f"cos={cos:.4f}") assert cos > 0.95, f"Cosine similarity too low: {cos:.4f}" def test_attention_speedup(self): """NVFP4 FA4 kernel is faster than BF16 FA4 (excluding quantization overhead).""" from flash_attn.cute.interface import flash_attn_func as fa4_func from fastvideo.attention.backends.flash_attn import _nvfp4_quantize_for_fa4 batch, seqlen, nheads, headdim = 1, self.MODEL_SEQLEN, self.MODEL_NHEADS, self.MODEL_HEADDIM q = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) k = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) v = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.bfloat16) # BF16 kernel timing bf16_ms = _cuda_timer(lambda: fa4_func(q, k, v, softmax_scale=headdim**-0.5, causal=False)) # FP4 kernel timing (pre-quantized, pad V to match K) q_fp4, q_sf = _nvfp4_quantize_for_fa4(q) k_fp4, k_sf = _nvfp4_quantize_for_fa4(k) seqlen_padded = ((seqlen + 127) // 128) * 128 v_padded = torch.nn.functional.pad(v, (0, 0, 0, 0, 0, seqlen_padded - seqlen)) if seqlen_padded != seqlen else v fp4_ms = _cuda_timer(lambda: fa4_func( q_fp4, k_fp4, v_padded, softmax_scale=headdim**-0.5, causal=False, mSFQ=q_sf, mSFK=k_sf, )) speedup = bf16_ms / fp4_ms print(f"BF16={bf16_ms:.2f}ms FP4={fp4_ms:.2f}ms speedup={speedup:.2f}x") assert speedup > 1.0, f"Expected speedup > 1.0, got {speedup:.2f}x" if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) @pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)], reason="Requires Blackwell GPU (sm100a or sm103a)") class TestAttnQatInferFA4Route: """ATTN_QAT_INFER resolves to the FP4 FA4 kernel on sm_100/sm_103. Hardware-validated on both GB200 (sm_100a) and GB300 (sm_103a): route resolution, SDPA parity, and cross-attention lengths pass on real silicon for each capability in the FA4-FP4 set. """ def test_resolves_to_fa4(self): from fastvideo.attention.backends.attn_qat_infer import (_resolved_kernel, attn_qat_infer_receipt, is_attn_qat_infer_available) assert is_attn_qat_infer_available() assert _resolved_kernel() == "fa4_fp4" receipt = attn_qat_infer_receipt() assert "qk_mode=nvfp4(per-16-e4m3-sf)" in receipt assert "pv_mode=bf16" in receipt def test_forward_parity_vs_sdpa(self): from fastvideo.attention.backends.attn_qat_infer import AttnQatInferImpl torch.manual_seed(0) b, s, h, d = 1, 4096, 12, 128 q = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) k = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) v = torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) impl = AttnQatInferImpl(num_heads=h, head_size=d, causal=False, softmax_scale=d**-0.5) out = impl.forward(q, k, v, attn_metadata=None) ref = torch.nn.functional.scaled_dot_product_attention( q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float()).transpose(1, 2) cos = torch.nn.functional.cosine_similarity(out.float().flatten(), ref.flatten(), dim=0).item() # Same bound the sm_120 CUTLASS kernel test uses; FA4 NVFP4 QK # measures cos_sim ~0.99 vs BF16 (kernel repo README precision table). assert cos >= 0.97, f"cos_sim={cos:.4f} < 0.97" def test_cross_attention_lengths(self): from fastvideo.attention.backends.attn_qat_infer import AttnQatInferImpl torch.manual_seed(0) b, h, d = 1, 12, 128 q = torch.randn(b, 384, h, d, device="cuda", dtype=torch.bfloat16) k = torch.randn(b, 512, h, d, device="cuda", dtype=torch.bfloat16) v = torch.randn(b, 512, h, d, device="cuda", dtype=torch.bfloat16) impl = AttnQatInferImpl(num_heads=h, head_size=d, causal=False, softmax_scale=d**-0.5) out = impl.forward(q, k, v, attn_metadata=None) assert out.shape == (b, 384, h, d)