Entry point called by bench.py. Must match reference.rmsnorm_ref signature.
(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6)
| 43 | |
| 44 | |
| 45 | def kernel_fn(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: |
| 46 | """Entry point called by bench.py. Must match reference.rmsnorm_ref signature.""" |
| 47 | assert x.is_cuda |
| 48 | M, N = x.shape |
| 49 | out = torch.empty_like(x) |
| 50 | |
| 51 | BLOCK_SIZE = triton.next_power_of_2(N) |
| 52 | |
| 53 | rmsnorm_kernel[(M,)]( |
| 54 | x, weight, out, |
| 55 | M, N, |
| 56 | x.stride(0), x.stride(1), |
| 57 | out.stride(0), out.stride(1), |
| 58 | eps, |
| 59 | BLOCK_SIZE=BLOCK_SIZE, |
| 60 | ) |
| 61 | return out |
nothing calls this directly
no outgoing calls
no test coverage detected