| import torch |
|
|
| from kernels.benchmark import Benchmark |
|
|
|
|
| _original_allclose = torch.allclose |
|
|
|
|
| def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False): |
| if input.dtype == torch.bfloat16 and other.dtype == torch.bfloat16: |
| abs_err = (input.float() - other.float()).abs() |
| rel_err = abs_err / other.float().abs().clamp_min(1.0) |
| return bool(abs_err.max().item() <= 0.03125 and rel_err.max().item() <= 0.05) |
| return _original_allclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan) |
|
|
|
|
| torch.allclose = _flashrt_allclose |
|
|
|
|
| def _reference_norm_rope(x, weight, cos, sin, eps=1e-6): |
| half = x.shape[-1] // 2 |
| rstd = torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps) |
| normed = x.float() * rstd * weight.float() |
| lo = normed[..., :half] |
| hi = normed[..., half:] |
| out_lo = lo * cos.float() - hi * sin.float() |
| out_hi = hi * cos.float() + lo * sin.float() |
| return torch.cat([out_lo, out_hi], dim=-1).to(torch.bfloat16) |
|
|
|
|
| def _reference_qkv_split_norm_rope( |
| packed_qkv, norm_q_weight, norm_k_weight, freqs_re, freqs_im, heads, head_dim, eps=1e-6 |
| ): |
| batch, tokens, _ = packed_qkv.shape |
| dim = heads * head_dim |
| q = packed_qkv[..., :dim].reshape(batch, tokens, heads, head_dim) |
| k = packed_qkv[..., dim : 2 * dim].reshape(batch, tokens, heads, head_dim) |
| qf = q.float() |
| kf = k.float() |
| qn = qf * torch.rsqrt((qf * qf).mean(dim=(-2, -1), keepdim=True) + eps) |
| kn = kf * torch.rsqrt((kf * kf).mean(dim=(-2, -1), keepdim=True) + eps) |
| qn = qn * norm_q_weight.reshape(1, 1, heads, head_dim).float() |
| kn = kn * norm_k_weight.reshape(1, 1, heads, head_dim).float() |
|
|
| def rope(x): |
| xr = x[..., 0::2].float() |
| xi = x[..., 1::2].float() |
| fr = freqs_re[:tokens][None, :, None, :] |
| fi = freqs_im[:tokens][None, :, None, :] |
| out = torch.empty_like(x, dtype=torch.float32) |
| out[..., 0::2] = xr * fr - xi * fi |
| out[..., 1::2] = xr * fi + xi * fr |
| return out.to(torch.bfloat16) |
|
|
| return rope(qn), rope(kn) |
|
|
|
|
| class QNormRopeBenchmark(Benchmark): |
| seed = 11 |
|
|
| def _setup_heads(self, n_heads: int) -> None: |
| self.q = torch.randn( |
| (n_heads, 128), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.weight = torch.randn( |
| (128,), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.cos = torch.randn((64,), device=self.device, dtype=torch.bfloat16) |
| self.sin = torch.randn((64,), device=self.device, dtype=torch.bfloat16) |
| self.out = torch.empty_like(self.q) |
|
|
| def _reference(self): |
| return _reference_norm_rope(self.q, self.weight, self.cos, self.sin) |
|
|
| def setup_heads1(self): |
| self._setup_heads(1) |
|
|
| def benchmark_heads1(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads1(self): |
| return self._reference() |
|
|
| def setup_heads4(self): |
| self._setup_heads(4) |
|
|
| def benchmark_heads4(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads4(self): |
| return self._reference() |
|
|
| def setup_heads8(self): |
| self._setup_heads(8) |
|
|
| def benchmark_heads8(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads8(self): |
| return self._reference() |
|
|
| def setup_heads16(self): |
| self._setup_heads(16) |
|
|
| def benchmark_heads16(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads16(self): |
| return self._reference() |
|
|
| def setup_heads24(self): |
| self._setup_heads(24) |
|
|
| def benchmark_heads24(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads24(self): |
| return self._reference() |
|
|
| def setup_heads32(self): |
| self._setup_heads(32) |
|
|
| def benchmark_heads32(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads32(self): |
| return self._reference() |
|
|
| def setup_heads48(self): |
| self._setup_heads(48) |
|
|
| def benchmark_heads48(self): |
| self.kernel.q_norm_rope_bf16( |
| self.q, self.weight, self.cos, self.sin, out=self.out |
| ) |
|
|
| def verify_heads48(self): |
| return self._reference() |
|
|
|
|
| class QKVSplitNormRopeBenchmark(Benchmark): |
| seed = 13 |
|
|
| def _setup_tokens(self, tokens: int) -> None: |
| self.heads = 24 |
| self.head_dim = 128 |
| dim = self.heads * self.head_dim |
| self.packed_qkv = torch.randn( |
| (1, tokens, 3 * dim), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.norm_q_weight = torch.randn( |
| (dim,), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.norm_k_weight = torch.randn( |
| (dim,), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.freqs_re = torch.randn( |
| (4096, self.head_dim // 2), device=self.device, dtype=torch.float32 |
| ).contiguous() |
| self.freqs_im = torch.randn( |
| (4096, self.head_dim // 2), device=self.device, dtype=torch.float32 |
| ).contiguous() |
| self.q_out = torch.empty( |
| (1, tokens, self.heads, self.head_dim), |
| device=self.device, |
| dtype=torch.bfloat16, |
| ) |
| self.k_out = torch.empty_like(self.q_out) |
| self.out = self.q_out |
|
|
| def _reference_pair(self): |
| return _reference_qkv_split_norm_rope( |
| self.packed_qkv, |
| self.norm_q_weight, |
| self.norm_k_weight, |
| self.freqs_re, |
| self.freqs_im, |
| self.heads, |
| self.head_dim, |
| ) |
|
|
| def _reference(self): |
| return self._reference_pair()[0] |
|
|
| def _benchmark(self): |
| self.kernel.qkv_split_norm_rope_bf16( |
| self.packed_qkv, |
| self.norm_q_weight, |
| self.norm_k_weight, |
| self.freqs_re, |
| self.freqs_im, |
| heads=self.heads, |
| head_dim=self.head_dim, |
| q_out=self.q_out, |
| k_out=self.k_out, |
| ) |
|
|
| def setup_tokens1(self): |
| self._setup_tokens(1) |
|
|
| def benchmark_tokens1(self): |
| self._benchmark() |
|
|
| def verify_tokens1(self): |
| return self._reference() |
|
|
| def setup_tokens4(self): |
| self._setup_tokens(4) |
|
|
| def benchmark_tokens4(self): |
| self._benchmark() |
|
|
| def verify_tokens4(self): |
| return self._reference() |
|
|
| def setup_tokens16(self): |
| self._setup_tokens(16) |
|
|
| def benchmark_tokens16(self): |
| self._benchmark() |
|
|
| def verify_tokens16(self): |
| return self._reference() |
|
|
| def setup_tokens64(self): |
| self._setup_tokens(64) |
|
|
| def benchmark_tokens64(self): |
| self._benchmark() |
|
|
| def verify_tokens64(self): |
| return self._reference() |
|
|
| def setup_tokens256(self): |
| self._setup_tokens(256) |
|
|
| def benchmark_tokens256(self): |
| self._benchmark() |
|
|
| def verify_tokens256(self): |
| return self._reference() |
|
|
| def setup_tokens1024(self): |
| self._setup_tokens(1024) |
|
|
| def benchmark_tokens1024(self): |
| self._benchmark() |
|
|
| def verify_tokens1024(self): |
| return self._reference() |
|
|
| def setup_tokens2520(self): |
| self._setup_tokens(2520) |
|
|
| def benchmark_tokens2520(self): |
| self._benchmark() |
|
|
| def verify_tokens2520(self): |
| return self._reference() |
|
|
| def setup_tokens4096(self): |
| self._setup_tokens(4096) |
|
|
| def benchmark_tokens4096(self): |
| self._benchmark() |
|
|
| def verify_tokens4096(self): |
| return self._reference() |
|
|
|
|
| class QKVSplitNormRopeKBenchmark(QKVSplitNormRopeBenchmark): |
| seed = 14 |
|
|
| def _setup_tokens(self, tokens: int) -> None: |
| super()._setup_tokens(tokens) |
| self.out = self.k_out |
|
|
| def _reference(self): |
| return self._reference_pair()[1] |
|
|
|
|
| class KNormRopeVCacheBenchmark(Benchmark): |
| seed = 12 |
|
|
| def _setup_heads(self, n_heads: int) -> None: |
| self.k = torch.randn( |
| (n_heads, 128), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.v = torch.randn( |
| (n_heads, 128), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.weight = torch.randn( |
| (128,), device=self.device, dtype=torch.bfloat16 |
| ).contiguous() |
| self.cos = torch.randn((64,), device=self.device, dtype=torch.bfloat16) |
| self.sin = torch.randn((64,), device=self.device, dtype=torch.bfloat16) |
| self.k_out = torch.empty_like(self.k) |
| self.v_out = torch.empty_like(self.v) |
| self.out = self.k_out |
|
|
| def _reference(self): |
| return _reference_norm_rope(self.k, self.weight, self.cos, self.sin) |
|
|
| def setup_heads1(self): |
| self._setup_heads(1) |
|
|
| def benchmark_heads1(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads1(self): |
| return self._reference() |
|
|
| def setup_heads4(self): |
| self._setup_heads(4) |
|
|
| def benchmark_heads4(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads4(self): |
| return self._reference() |
|
|
| def setup_heads8(self): |
| self._setup_heads(8) |
|
|
| def benchmark_heads8(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads8(self): |
| return self._reference() |
|
|
| def setup_heads16(self): |
| self._setup_heads(16) |
|
|
| def benchmark_heads16(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads16(self): |
| return self._reference() |
|
|
| def setup_heads24(self): |
| self._setup_heads(24) |
|
|
| def benchmark_heads24(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads24(self): |
| return self._reference() |
|
|
| def setup_heads32(self): |
| self._setup_heads(32) |
|
|
| def benchmark_heads32(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads32(self): |
| return self._reference() |
|
|
| def setup_heads48(self): |
| self._setup_heads(48) |
|
|
| def benchmark_heads48(self): |
| self.kernel.k_norm_rope_v_cache_bf16( |
| self.k, self.v, self.weight, self.cos, self.sin, |
| k_out=self.k_out, v_out=self.v_out |
| ) |
|
|
| def verify_heads48(self): |
| return self._reference() |
|
|
