| import torch | |
| from kernels import get_kernel | |
| mha = get_kernel("flashrt/masked-mha-runtime", version=1) | |
| q = torch.randn((41, 32, 48), device="cuda", dtype=torch.bfloat16) | |
| k = torch.randn_like(q) | |
| v = torch.randn_like(q) | |
| logits = mha.allocate_workspace(q, k) | |
| out = torch.empty_like(q) | |
| mha.forward_static(q, k, v, logits=logits, out=out) | |