(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback, is_fp8=False)
| 30 | |
| 31 | @torch.no_grad() |
| 32 | def _test_attn_func(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback, is_fp8=False): |
| 33 | if is_causal and S_Q != S_KV: |
| 34 | pytest.skip("Causal attention is only supported for S_Q == S_KV") |
| 35 | |
| 36 | if is_fp8: |
| 37 | attn_func = fp8_attention |
| 38 | else: |
| 39 | attn_func = vanilla_attention |
| 40 | |
| 41 | torch.manual_seed(0) |
| 42 | query = torch.randn(B, H, S_Q, D, dtype=dtype, device=device) |
| 43 | key = torch.randn(B, H, S_KV, D, dtype=dtype, device=device) |
| 44 | value = torch.randn(B, H, S_KV, D, dtype=dtype, device=device) |
| 45 | |
| 46 | with quantum_attn.config.patch( |
| 47 | { |
| 48 | "attention.force_eager_fallback": force_eager_fallback, |
| 49 | } |
| 50 | ): |
| 51 | try: |
| 52 | attn_out = attn_func(query, key, value, is_causal=is_causal) |
| 53 | except ValueError as e: |
| 54 | pytest.skip(str(e)) |
| 55 | |
| 56 | fa_out = flash_attention(query, key, value, is_causal=is_causal) |
| 57 | |
| 58 | rmse = torch.sqrt(F.mse_loss(attn_out, fa_out)) |
| 59 | print(f"RMSE: {rmse}") |
| 60 | assert rmse < 1e-2, f"RMSE: {rmse}" |
| 61 | |
| 62 | |
| 63 | @pytest.mark.parametrize("B", [1, 2]) |
no test coverage detected