Function kernel_fn

kernels/cuda/softmax.py:245–275

Entry point called by bench.py. Must match reference.softmax_ref signature.

(x: torch.Tensor)

def kernel_fn(x: torch.Tensor) -> torch.Tensor:
    """Entry point called by bench.py. Must match reference.softmax_ref signature."""
    assert x.is_cuda

    # Preserve original shape and dtype
    orig_shape = x.shape
    orig_dtype = x.dtype

    # Ensure contiguous layout
    x = x.contiguous()

    # Flatten to 2-D: (n_rows, n_cols) -- softmax along last dim
    if x.ndim == 1:
        x = x.unsqueeze(0)
    elif x.ndim > 2:
        x = x.view(-1, x.shape[-1])

    # The CUDA launcher handles fp16/bf16/fp32 internally.
    # For unsupported dtypes, cast to fp32 round-trip.
    needs_cast = orig_dtype not in (torch.float16, torch.bfloat16, torch.float32)
    if needs_cast:
        x = x.to(torch.float32)

    mod = _get_module()
    out = mod.softmax_cuda(x)

    # Cast back if we had to promote
    if needs_cast:
        out = out.to(orig_dtype)

    return out.view(orig_shape)