fix(triton): widen softmax/sparsemax row index to int64 (#1497)

## Summary

Widen the softmax and sparsemax program IDs to int64 before row-offset
address arithmetic. With an int32 `tl.program_id(0)`, `row * row_stride`
wraps once `n_rows * n_cols` crosses 2^31, and the kernels hit an
illegal memory access (e.g. `n_cols=32768`, `n_rows > 65536`).

This applies the same safeguard as #251, #804 and #1482 to the four
softmax kernels and the two sparsemax kernels.

Fixes #1496.

## Details

- The softmax multi-block kernels are currently unreachable (see #1252),
but they get the same cast so they stay safe once that path is enabled.
- The sparsemax regression test launches `_sparsemax_forward_kernel` /
`_sparsemax_backward_kernel` directly: at this size the `torch.sort` in
the autograd path would need ~26 GB extra for sorted values and int64
indices.
- The Ascend backends
(`ops/backends/_ascend/ops/{softmax,sparsemax}.py`) are left unchanged
since I have no NPU to test on.

## Testing Done

- Hardware Type: NVIDIA L4 (24 GB)
- [x] `python -m pytest -v test/transformers/test_softmax.py
test/transformers/test_sparsemax.py` (36 passed, including both new
large-tensor tests, which need 17.2 GB free VRAM)
- [x] The two new large-tensor tests fail on `main` with `illegal memory
access was encountered`
- [x] `ruff check` and `ruff format --check` on the changed files
- [ ] `make test` (only the softmax and sparsemax suites were run)
- [ ] `make checkstyle` (ran the equivalent `ruff` commands above on the
changed files)
- [ ] `make test-convergence` (not run; the change only affects address
arithmetic)

Co-authored-by: Vaibhav Jindal <[email protected]>

src/liger_kernel/ops/softmax.py

@@ -17,7 +17,7 @@ def _softmax_single_block_forward_kernel(
     n_cols,
     BLOCK_SIZE: tl.constexpr,
 ):
-    row_id = tl.program_id(0)
+    row_id = tl.program_id(0).to(tl.int64)
     offs = tl.arange(0, BLOCK_SIZE)
     mask = offs < n_cols
 
@@ -38,7 +38,7 @@ def _softmax_multi_block_forward_kernel(
     n_cols,
     BLOCK_SIZE: tl.constexpr,
 ):
-    row_id = tl.program_id(0)
+    row_id = tl.program_id(0).to(tl.int64)
     offs = tl.arange(0, BLOCK_SIZE)
 
     m = tl.float32(-float("inf"))
@@ -71,7 +71,7 @@ def _softmax_single_block_backward_kernel(
     n_cols,
     BLOCK_SIZE: tl.constexpr,
 ):
-    row_id = tl.program_id(0)
+    row_id = tl.program_id(0).to(tl.int64)
     offs = tl.arange(0, BLOCK_SIZE)
     mask = offs < n_cols
 
@@ -93,7 +93,7 @@ def _softmax_multi_block_backward_kernel(
     n_cols,
     BLOCK_SIZE: tl.constexpr,
 ):
-    row_id = tl.program_id(0)
+    row_id = tl.program_id(0).to(tl.int64)
     offs = tl.arange(0, BLOCK_SIZE)
     acc = tl.float32(0.0)
 

src/liger_kernel/ops/sparsemax.py

@@ -20,7 +20,7 @@ def _sparsemax_forward_kernel(
     BLOCK_SIZE: tl.constexpr,
     num_warps: tl.constexpr,
 ):
-    pid_row = tl.program_id(0)
+    pid_row = tl.program_id(0).to(tl.int64)
     ptr_x_data_row = x_ptr + pid_row * x_stride_row
     ptr_sorted_x_data_row = sorted_x_ptr + pid_row * sorted_x_stride_row
     ptr_output_row = o_ptr + pid_row * o_stride_row
@@ -72,7 +72,7 @@ def _sparsemax_forward_kernel(
 def _sparsemax_backward_kernel(
     o_ptr, go_ptr, gi_ptr, stride, n_cols, BLOCK_SIZE: tl.constexpr, num_warps: tl.constexpr
 ):
-    row = tl.program_id(0)
+    row = tl.program_id(0).to(tl.int64)
     o_row = o_ptr + row * stride
     go_row = go_ptr + row * stride
     gi_row = gi_ptr + row * stride

test/transformers/test_softmax.py

@@ -12,6 +12,11 @@ from liger_kernel.utils import infer_device
 device = infer_device()
 set_seed()
 
+_SOFTMAX_OVERFLOW_ROWS = 2**16 + 16
+_SOFTMAX_OVERFLOW_COLS = 2**15
+# x, y, grad_output and x.grad in bf16
+_SOFTMAX_MIN_FREE_BYTES = 4 * _SOFTMAX_OVERFLOW_ROWS * _SOFTMAX_OVERFLOW_COLS * 2
+
 
 @pytest.mark.parametrize(
     "shape",
@@ -101,3 +106,36 @@ def test_liger_softmax_functional(shape, dtype, atol, rtol):
     liger_out.backward(grad_output, retain_graph=True)
 
     assert_verbose_allclose(x1.grad, x2.grad, atol=atol, rtol=rtol)
+
+
+def test_softmax_int32_row_offset_wraps():
+    """The first overflowing row offset wraps in int32."""
+    row_id = torch.tensor(2**16, dtype=torch.int32)
+    row_stride = torch.tensor(_SOFTMAX_OVERFLOW_COLS, dtype=torch.int32)
+
+    assert (row_id * row_stride).item() == -(2**31)
+    assert row_id.to(torch.int64).mul(row_stride).item() == 2**31
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.mem_get_info()[0] < _SOFTMAX_MIN_FREE_BYTES,
+    reason="requires 17.2 GB of free CUDA memory for rows beyond the int32 offset range",
+)
+def test_softmax_large_row_offset():
+    """Forward and backward must address rows whose offset exceeds 2^31 elements."""
+    x = torch.zeros(_SOFTMAX_OVERFLOW_ROWS, _SOFTMAX_OVERFLOW_COLS, dtype=torch.bfloat16, device="cuda")
+    x[-1] = torch.randn(_SOFTMAX_OVERFLOW_COLS, dtype=torch.bfloat16, device="cuda")
+    x.requires_grad_(True)
+
+    y = liger_softmax(x)
+
+    grad_output = torch.zeros_like(y)
+    grad_output[-1] = torch.randn(_SOFTMAX_OVERFLOW_COLS, dtype=torch.bfloat16, device="cuda")
+    y.backward(grad_output)
+
+    x_ref = x[-1].detach().float().requires_grad_(True)
+    y_ref = torch.softmax(x_ref, dim=-1)
+    y_ref.backward(grad_output[-1].float())
+
+    assert_verbose_allclose(y[-1].float(), y_ref, atol=5e-2, rtol=5e-2)
+    assert_verbose_allclose(x.grad[-1].float(), x_ref.grad, atol=5e-2, rtol=5e-2)

test/transformers/test_sparsemax.py

@@ -4,12 +4,20 @@ import torch
 from test.utils import assert_verbose_allclose
 from test.utils import set_seed
 
+from liger_kernel.ops.sparsemax import _sparsemax_backward_kernel
+from liger_kernel.ops.sparsemax import _sparsemax_forward_kernel
+from liger_kernel.ops.utils import calculate_settings
 from liger_kernel.transformers.functional import liger_sparsemax
 from liger_kernel.transformers.sparsemax import LigerSparsemax
 from liger_kernel.utils import infer_device
 
 device = infer_device()
 
+_SPARSEMAX_OVERFLOW_ROWS = 2**16 + 16
+_SPARSEMAX_OVERFLOW_COLS = 2**15
+# forward peak: bf16 input + fp32 sorted input + bf16 output
+_SPARSEMAX_MIN_FREE_BYTES = (2 + 4 + 2) * _SPARSEMAX_OVERFLOW_ROWS * _SPARSEMAX_OVERFLOW_COLS
+
 
 def torch_sparsemax(input_tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
     input_dims = input_tensor.dim()
@@ -109,3 +117,65 @@ def test_liger_sparsemax_functional_correctness(batch_size, seq_len, features, d
     out_l.backward(g)
     out_t.backward(g)
     assert_verbose_allclose(lx.grad, tx.grad, atol=atol, rtol=rtol)
+
+
+def test_sparsemax_int32_row_offset_wraps():
+    """The first overflowing row offset wraps in int32."""
+    row = torch.tensor(2**16, dtype=torch.int32)
+    row_stride = torch.tensor(_SPARSEMAX_OVERFLOW_COLS, dtype=torch.int32)
+
+    assert (row * row_stride).item() == -(2**31)
+    assert row.to(torch.int64).mul(row_stride).item() == 2**31
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.mem_get_info()[0] < _SPARSEMAX_MIN_FREE_BYTES,
+    reason="requires 17.2 GB of free CUDA memory for rows beyond the int32 offset range",
+)
+def test_sparsemax_large_row_offset():
+    """Forward and backward kernels must address rows whose offset exceeds 2^31 elements.
+
+    The kernels are launched directly: at this size the torch.sort in the autograd
+    path would need ~26 GB more memory for the sorted values and int64 indices.
+    """
+    n_rows, n_cols = _SPARSEMAX_OVERFLOW_ROWS, _SPARSEMAX_OVERFLOW_COLS
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    last = torch.randn(n_cols, dtype=torch.bfloat16, device="cuda")
+    x = torch.zeros(n_rows, n_cols, dtype=torch.bfloat16, device="cuda")
+    x[-1] = last
+    x_sorted = torch.zeros(n_rows, n_cols, dtype=torch.float32, device="cuda")
+    x_sorted[-1] = torch.sort(last.float(), descending=True).values
+    out = torch.empty_like(x)
+    _sparsemax_forward_kernel[(n_rows,)](
+        x,
+        x.stride(0),
+        x_sorted,
+        x_sorted.stride(0),
+        out,
+        out.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    del x, x_sorted
+
+    x_ref = last.float().requires_grad_(True)
+    out_ref = torch_sparsemax(x_ref)
+    assert_verbose_allclose(out[-1].float(), out_ref, atol=5e-2, rtol=5e-2)
+
+    grad_out = torch.zeros_like(out)
+    grad_out[-1] = torch.randn(n_cols, dtype=torch.bfloat16, device="cuda")
+    grad_in = torch.empty_like(out)
+    _sparsemax_backward_kernel[(n_rows,)](
+        out,
+        grad_out,
+        grad_in,
+        out.stride(0),
+        n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+
+    out_ref.backward(grad_out[-1].float())
+    assert_verbose_allclose(grad_in[-1].float(), x_ref.grad, atol=5e-2, rtol=5e-2)