Entry point called by bench.py. Must match reference.rmsnorm_ref signature.
(x: torch.Tensor, weight: torch.Tensor)
| 169 | |
| 170 | |
| 171 | def kernel_fn(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: |
| 172 | """Entry point called by bench.py. Must match reference.rmsnorm_ref signature.""" |
| 173 | assert x.is_cuda and weight.is_cuda |
| 174 | |
| 175 | # Handle non-fp16 inputs by casting |
| 176 | orig_dtype = x.dtype |
| 177 | if x.dtype != torch.float16: |
| 178 | x = x.to(torch.float16) |
| 179 | if weight.dtype != torch.float16: |
| 180 | weight = weight.to(torch.float16) |
| 181 | |
| 182 | mod = _get_module() |
| 183 | out = mod.rmsnorm_cuda(x, weight) |
| 184 | |
| 185 | # Cast back if needed |
| 186 | if orig_dtype != torch.float16: |
| 187 | out = out.to(orig_dtype) |
| 188 | |
| 189 | return out |
nothing calls this directly
no test coverage detected