File size: 333 Bytes
8e9f35a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | import torch
from kernels import get_kernel
kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)
q = torch.randn(1, 4096, 16, 128, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)
out = kernel.sol_attn(
q,
k,
v,
tau=1.0,
thresh_type="exact",
)
print(out.shape)
|