File size: 1,293 Bytes
8e9f35a | 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 | import importlib
import torch
from kernels.benchmark import Benchmark
class SolAttnBenchmark(Benchmark):
seed: int = 42
def _setup(self, tokens, heads):
self.q = torch.randn(
1,
tokens,
heads,
128,
device=self.device,
dtype=torch.bfloat16,
)
self.k = torch.randn_like(self.q)
self.v = torch.randn_like(self.q)
module = importlib.import_module(f"{self.kernel.__name__}.triton_ref")
self.triton_sol_attn = module.sol_attn
def _run(self):
self.out = self.kernel.sol_attn(
self.q,
self.k,
self.v,
tau=1.0,
thresh_type="exact",
)
def _reference(self):
return self.triton_sol_attn(
self.q,
self.k,
self.v,
tau=1.0,
thresh_type="exact",
)
def setup_base(self):
self._setup(tokens=4096, heads=8)
def benchmark_base(self):
self._run()
def verify_base(self):
return self._reference()
def setup_video(self):
self._setup(tokens=16384, heads=16)
def benchmark_video(self):
self._run()
def verify_video(self):
return self._reference()
|