Kernels:
Trusted publisher
File size: 2,086 Bytes
b0b387a | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # /// script
# requires-python = ">=3.10"
# dependencies = [
# "kernels",
# "numpy",
# "torch",
# ]
# ///
import platform
from pathlib import Path
import kernels
import torch
import torch.nn.functional as F
def main() -> None:
# Load the locally built kernel. The second arg is the *backend* ("cuda"),
# not the kernel name. `result` is the build-output symlink from
# `kernel-builder build`; fall back to `build` if present.
repo = Path(__file__).parent
build_dir = repo / "result" if (repo / "result").exists() else repo / "build"
kernel = kernels.get_local_kernel(build_dir, "cuda")
# Select device
if platform.system() == "Darwin":
device = torch.device("mps")
elif hasattr(torch, "xpu") and torch.xpu.is_available():
device = torch.device("xpu")
elif torch.version.cuda is not None and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
print(f"Using device: {device}")
# Create (B, H, S, D) inputs for scaled dot-product attention
B, H, S, D = 2, 8, 512, 64
q = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
k = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
v = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
# Run kernel (scaled dot-product attention)
result = kernel.attention(q, k, v)
print(f"Output shape: {tuple(result.shape)}")
# Verify result against PyTorch SDPA
expected = F.scaled_dot_product_attention(q, k, v)
assert torch.allclose(result, expected, atol=5e-2, rtol=2e-2), (
"Kernel output doesn't match SDPA!"
)
print("Success!")
# NOTE: the `if __name__ == "__main__"` guard is REQUIRED, not stylistic.
# Helion autotunes in a *spawned* subprocess that re-imports this module; an
# unguarded top-level kernel call would be re-executed on every worker import,
# recursively spawning autotuners until it aborts with NoConfigFound. See
# ISSUES.md ("Autotuner spawn / __main__ guard").
if __name__ == "__main__":
main()
|