(attention_size, attention_dtype)
| 636 | |
| 637 | |
| 638 | def test_attention_offload(attention_size, attention_dtype): |
| 639 | b, (s, s_kv), n, (h, h_v) = attention_size |
| 640 | concrete_s, concrete_s_kv = _to_concrete_shape((s, s_kv)) |
| 641 | q, k, v, _, ref = get_numpy_attention_ref( |
| 642 | b, concrete_s, concrete_s_kv, n, h, h_v, "none", "none", "none", attention_dtype |
| 643 | ) |
| 644 | |
| 645 | q_shape = (b, s, n, h) |
| 646 | k_shape = (b, s_kv, n, h) |
| 647 | v_shape = (b, s_kv, n, h_v) |
| 648 | |
| 649 | mod = get_relax_attention_module(q_shape, k_shape, v_shape, dtype=attention_dtype) |
| 650 | out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2) |
| 651 | |
| 652 | tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) |
| 653 | |
| 654 | |
| 655 | @pytest.fixture( |
nothing calls this directly
no test coverage detected
searching dependent graphs…