(attention_bias_size)
| 670 | |
| 671 | |
| 672 | def test_attention_bias_offload(attention_bias_size): |
| 673 | b, (s, s_kv), n, (h, h_v), bias_shape = attention_bias_size |
| 674 | concrete_s, concrete_s_kv, concrete_bias_shape = _to_concrete_shape((s, s_kv, bias_shape)) |
| 675 | |
| 676 | q, k, v, bias, ref = get_numpy_attention_ref( |
| 677 | b, concrete_s, concrete_s_kv, n, h, h_v, concrete_bias_shape, "none", "none", "float32" |
| 678 | ) |
| 679 | |
| 680 | q_shape = (b, s, n, h) |
| 681 | k_shape = (b, s_kv, n, h) |
| 682 | v_shape = (b, s_kv, n, h_v) |
| 683 | |
| 684 | mod = get_relax_attention_module( |
| 685 | q_shape, k_shape, v_shape, bias_shape=bias_shape, dtype="float32" |
| 686 | ) |
| 687 | out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias, num_final_bindings=2) |
| 688 | |
| 689 | tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) |
| 690 | |
| 691 | |
| 692 | @pytest.fixture( |
nothing calls this directly
no test coverage detected
searching dependent graphs…