(attention_causal_size, attention_causal)
| 743 | |
| 744 | |
| 745 | def test_attention_causal_offload(attention_causal_size, attention_causal): |
| 746 | b, (s, s_kv), n, (h, h_v), bias_shape = attention_causal_size |
| 747 | q, k, v, bias, ref = get_numpy_attention_ref( |
| 748 | b, s, s_kv, n, h, h_v, bias_shape, "none", attention_causal, "float16" |
| 749 | ) |
| 750 | |
| 751 | q_shape = (b, s, n, h) |
| 752 | k_shape = (b, s_kv, n, h) |
| 753 | v_shape = (b, s_kv, n, h_v) |
| 754 | |
| 755 | mod = get_relax_attention_module( |
| 756 | q_shape, |
| 757 | k_shape, |
| 758 | v_shape, |
| 759 | dtype="float16", |
| 760 | bias_shape=bias_shape, |
| 761 | causal_mask=attention_causal, |
| 762 | ) |
| 763 | |
| 764 | if bias is None: |
| 765 | out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2) |
| 766 | else: |
| 767 | out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias, num_final_bindings=2) |
| 768 | tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) |
| 769 | |
| 770 | |
| 771 | @memoize("topi.tests.test_codegen_cutlass.test_stacked_attention_offload") |
nothing calls this directly
no test coverage detected
searching dependent graphs…