Function test_attention_rewrite_multi_query

tests/python/relax/test_codegen_cutlass.py:2087–2130

()

def test_attention_rewrite_multi_query():
    @I.ir_module(s_tir=True)
    class Module:
        @R.function
        def main(
            q: R.Tensor((4, 16, 32, 16), dtype="float16"),
            k_single: R.Tensor((4, 16, 1, 16), dtype="float16"),
            v_single: R.Tensor((4, 16, 1, 16), dtype="float16"),
        ) -> R.Tensor((4, 16, 32, 8), dtype="float16"):
            with R.dataflow():
                k = R.repeat(k_single, 32, axis=2)
                v = R.repeat(v_single, 32, axis=2)

                lv = R.permute_dims(q, axes=[0, 2, 1, 3])
                lv1 = R.reshape(lv, R.shape([128, 16, 16]))
                lv2 = R.permute_dims(k, axes=[0, 2, 1, 3])
                lv3 = R.reshape(lv2, R.shape([128, 16, 16]))
                lv4 = R.permute_dims(v, axes=[0, 2, 1, 3])
                lv5 = R.reshape(lv4, R.shape([128, 16, 16]))

                lv6 = R.permute_dims(lv3, axes=[0, 2, 1])
                lv7 = R.matmul(lv1, lv6, out_dtype="float16")
                lv3_1 = R.astype(R.const(0.25, "float32"), "float16")
                lv8 = R.multiply(lv7, lv3_1)
                lv11 = R.astype(R.nn.softmax(R.astype(lv8, "float32"), axis=2), "float16")
                lv12 = R.matmul(lv11, lv5, out_dtype="float16")
                lv13 = R.reshape(lv12, R.shape([4, 32, 16, 16]))
                lv6_1 = R.permute_dims(lv13, axes=[0, 2, 1, 3])
                R.output(lv6_1)
            return lv6_1

    q_np = np.random.randn(4, 16, 32, 16).astype("float16")
    k_np = np.random.randn(4, 16, 1, 16).astype("float16")
    v_np = np.random.randn(4, 16, 1, 16).astype("float16")
    args = [q_np, k_np, v_np]
    ref = build_and_run(Module, args, "llvm", legalize=True)

    mod = partition_for_cutlass(Module, use_flash_mqa=True)
    codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80}})
    mod = codegen_pass(mod)

    out = build_and_run(mod, args, "cuda")

    tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)


def _test_batched_var_len_attention(