Entry point called by bench.py. Must match reference.softmax_ref signature.
(x: torch.Tensor)
| 55 | |
| 56 | |
| 57 | def kernel_fn(x: torch.Tensor) -> torch.Tensor: |
| 58 | """Entry point called by bench.py. Must match reference.softmax_ref signature.""" |
| 59 | assert x.is_cuda |
| 60 | |
| 61 | # Flatten to 2D for row-parallel processing |
| 62 | orig_shape = x.shape |
| 63 | if x.ndim == 1: |
| 64 | x = x.unsqueeze(0) |
| 65 | elif x.ndim > 2: |
| 66 | x = x.view(-1, x.shape[-1]) |
| 67 | |
| 68 | n_rows, n_cols = x.shape |
| 69 | output = torch.empty_like(x) |
| 70 | |
| 71 | # Block size must be a power of 2 >= n_cols |
| 72 | BLOCK_SIZE = triton.next_power_of_2(n_cols) |
| 73 | |
| 74 | grid = (n_rows,) |
| 75 | softmax_kernel[grid]( |
| 76 | x, output, |
| 77 | n_cols, |
| 78 | x.stride(0), |
| 79 | output.stride(0), |
| 80 | BLOCK_SIZE=BLOCK_SIZE, |
| 81 | ) |
| 82 | |
| 83 | return output.view(orig_shape) |
nothing calls this directly
no outgoing calls
no test coverage detected