| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| def _extract_output(result): |
| if isinstance(result, tuple): |
| return result[0] |
| return result |
|
|
|
|
| def _reference_attention(query, key, value, causal=False): |
| """Reference SDPA on (B, S, H, D) inputs (matches kernels' util).""" |
| query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value)) |
| with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): |
| out = torch.nn.functional.scaled_dot_product_attention( |
| query, key, value, is_causal=causal |
| ) |
| return out.transpose(1, 2).contiguous() |
|
|
|
|
| class AttentionHelionBenchmark(Benchmark): |
| """Mirrors kernels.benchmarks.attention.FlashAttentionBenchmark. |
| |
| Uses the (B, S, H, D) flash-attn layout and calls ``flash_attn_func`` so |
| the same workloads apply to this Helion kernel and to flash-attn3. |
| """ |
|
|
| seed: int = 42 |
|
|
| def setup_small(self): |
| B, S, H, D = 2, 128, 8, 64 |
| self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16) |
|
|
| def benchmark_small(self): |
| self.out = _extract_output( |
| self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False) |
| ) |
|
|
| def verify_small(self) -> torch.Tensor: |
| return _reference_attention(self.q, self.k, self.v, causal=False) |
|
|
| def setup_medium(self): |
| B, S, H, D = 4, 512, 16, 64 |
| self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16) |
|
|
| def benchmark_medium(self): |
| self.out = _extract_output( |
| self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False) |
| ) |
|
|
| def verify_medium(self) -> torch.Tensor: |
| return _reference_attention(self.q, self.k, self.v, causal=False) |
|
|
| def setup_large(self): |
| B, S, H, D = 8, 1024, 32, 128 |
| self.q = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.k = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.v = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16) |
| self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16) |
|
|
| def benchmark_large(self): |
| self.out = _extract_output( |
| self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False) |
| ) |
|
|
| def verify_large(self) -> torch.Tensor: |
| return _reference_attention(self.q, self.k, self.v, causal=False) |
|
|
