File size: 1,724 Bytes
aa3ac98 81fff32 aa3ac98 81fff32 aa3ac98 |
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 64 |
# /// script
# requires-python = ">=3.10"
# dependencies = [
# "numpy",
# "torch==2.8.0",
# "kernels-benchmark-tools",
# "kernels",
# ]
#
# [tool.uv.sources]
# kernels-benchmark-tools = { path = "../../../../../tools", editable = true }
# ///
import torch
import sys
import kernels_benchmark_tools as kbt
from kernels import get_kernel
# Load the activation kernel
activation = get_kernel("kernels-community/activation")
def hf_kernels_swiglu(input_tensor):
"""HuggingFace Kernels SwiGLU implementation"""
hidden_dim = input_tensor.shape[-1] // 2
out_shape = input_tensor.shape[:-1] + (hidden_dim,)
out = torch.empty(out_shape, dtype=input_tensor.dtype, device=input_tensor.device)
return activation.silu_and_mul(out, input_tensor)
# Register the implementation
kbt.add(
"hf_kernels_swiglu",
hf_kernels_swiglu,
tags={"family": "hf-kernels", "backend": "triton", "compile": "none"},
)
if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
print("HF Kernels SwiGLU requires CUDA - skipping benchmark")
sys.exit(0)
dtype = "bfloat16"
# Generate workloads - using a subset for faster testing
wl = list(kbt.activation.llama_workloads(dtype=dtype))[:3] # First 3 workloads
print(f"Running SwiGLU benchmarks on {device} with {dtype}")
print(f"Testing {len(wl)} workloads")
# Run benchmark
kbt.run(
wl,
jsonl="activation.jsonl",
reps=5,
warmup=2,
gen=kbt.activation.gen_inputs,
ref=kbt.activation.ref_swiglu,
cmp=kbt.activation.cmp_allclose,
profile_trace=True
)
kbt.summarize(["activation.jsonl"]) |