Entry point called by bench.py. Must match reference.layernorm_ref signature.
(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor)
| 318 | |
| 319 | |
| 320 | def kernel_fn(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: |
| 321 | """Entry point called by bench.py. Must match reference.layernorm_ref signature.""" |
| 322 | assert x.is_cuda |
| 323 | |
| 324 | # Flatten to 2D for row-parallel processing |
| 325 | orig_shape = x.shape |
| 326 | if x.ndim == 1: |
| 327 | x = x.unsqueeze(0) |
| 328 | elif x.ndim > 2: |
| 329 | x = x.view(-1, x.shape[-1]) |
| 330 | |
| 331 | n_rows, n_cols = x.shape |
| 332 | assert weight.shape[0] == n_cols |
| 333 | assert bias.shape[0] == n_cols |
| 334 | |
| 335 | # Handle non-fp16 inputs by casting to fp16 for the CUDA kernel |
| 336 | orig_dtype = x.dtype |
| 337 | if x.dtype != torch.float16: |
| 338 | x = x.to(torch.float16) |
| 339 | if weight.dtype != torch.float16: |
| 340 | weight = weight.to(torch.float16) |
| 341 | if bias.dtype != torch.float16: |
| 342 | bias = bias.to(torch.float16) |
| 343 | |
| 344 | mod = _get_module() |
| 345 | y = mod.layernorm_cuda(x, weight, bias) |
| 346 | |
| 347 | # Cast back to original dtype if needed |
| 348 | if orig_dtype != torch.float16: |
| 349 | y = y.to(orig_dtype) |
| 350 | |
| 351 | return y.view(orig_shape) |
nothing calls this directly
no test coverage detected