(attention_scale_size, attention_scale)
| 706 | |
| 707 | |
| 708 | def test_attention_scale_offload(attention_scale_size, attention_scale): |
| 709 | b, (s, s_kv), n, (h, h_v), bias_shape = attention_scale_size |
| 710 | q, k, v, bias, ref = get_numpy_attention_ref( |
| 711 | b, s, s_kv, n, h, h_v, bias_shape, attention_scale, "none", "float32" |
| 712 | ) |
| 713 | |
| 714 | q_shape = (b, s, n, h) |
| 715 | k_shape = (b, s_kv, n, h) |
| 716 | v_shape = (b, s_kv, n, h_v) |
| 717 | |
| 718 | mod = get_relax_attention_module( |
| 719 | q_shape, k_shape, v_shape, dtype="float32", bias_shape=bias_shape, qk_scale=attention_scale |
| 720 | ) |
| 721 | if bias is None: |
| 722 | out = get_result_with_relax_cutlass_offload(mod, q, k, v, num_final_bindings=2) |
| 723 | else: |
| 724 | out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias, num_final_bindings=2) |
| 725 | tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2) |
| 726 | |
| 727 | |
| 728 | @pytest.fixture( |
nothing calls this directly
no test coverage detected
searching dependent graphs…