Kernels:

HelionDSL
/

attention

Trusted publisher

attention / example.py

Upload Helion attention kernel (noarch triton, pre-tuned sm100+sm90)

b0b387a verified 7 days ago

2.09 kB

	# /// 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()