Kernels:
Trusted publisher
File size: 1,116 Bytes
ffa8327 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | # /// script
# requires-python = ">=3.10"
# dependencies = [
# "helion",
# "kernels",
# "torch",
# ]
# ///
from __future__ import annotations
from pathlib import Path
import kernels
import torch
import torch.nn.functional as F
def main() -> None:
project = Path(__file__).parent
build = project / "result" if (project / "result").exists() else project / "build"
linear_attention = kernels.get_local_kernel(build, "cuda")
b, t, h, d = 2, 512, 8, 64
q = torch.randn(b, t, h, d, device="cuda", dtype=torch.bfloat16)
k = F.normalize(
torch.randn(b, t, h, d, device="cuda", dtype=torch.float32), dim=-1
).to(q)
v = torch.randn_like(q)
g = -torch.rand(b, t, h, d, device="cuda", dtype=torch.float32) * 0.1
beta = torch.rand(b, t, h, device="cuda", dtype=torch.bfloat16)
output, final_state = linear_attention.chunk_kda(
q,
k,
v,
g,
beta,
output_final_state=True,
)
print("output:", tuple(output.shape))
print("final state:", tuple(final_state.shape))
if __name__ == "__main__":
main()
|