Function get_numpy_attention_ref

tests/python/relax/test_codegen_cutlass.py:608–635

(
    b,
    s,
    s_kv,
    n,
    h,
    h_v,
    bias_shape,
    qk_scale,
    causal,
    dtype,
    window_size=None,
    num_kv_head=None,
)

def get_numpy_attention_ref(
    b,
    s,
    s_kv,
    n,
    h,
    h_v,
    bias_shape,
    qk_scale,
    causal,
    dtype,
    window_size=None,
    num_kv_head=None,
):
    num_kv_head = num_kv_head or n
    q = np.random.randn(b, s, n, h).astype(dtype)
    k = np.random.randn(b, s_kv, num_kv_head, h).astype(dtype)
    v = np.random.randn(b, s_kv, num_kv_head, h_v).astype(dtype)
    if bias_shape == "none":
        bias = None
    else:
        bias = np.random.randn(*bias_shape).astype(dtype)

    ref = tvm.topi.testing.attention_python(
        q, k, v, bias, qk_scale, causal=causal, window_size=window_size, layout="BSNH"
    )

    return q, k, v, bias, ref


def test_attention_offload(attention_size, attention_dtype):