Function test_attn_func

tests/test_interface.py:72–73

(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback)

@pytest.mark.parametrize("is_causal", [False, True])
@pytest.mark.parametrize("force_eager_fallback", [False])
def test_attn_func(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback):
    _test_attn_func(B, H, S_Q, S_KV, D, dtype, device, is_causal, force_eager_fallback)


@pytest.mark.parametrize("B", [1, 2])