Fix int32 overflow in chunk_cumsum kernels' pointer arithmetic

[chipsncheese-ie]

Summary

_chunk_cumsum_fwd_kernel and _chunk_cumsum_bwd_kernel in ssd_chunk_state.py compute pointer offsets as:

dt_ptr += pid_b * stride_dt_batch + pid_c * chunk_size * stride_dt_seqlen

Triton's JIT specializes pid_b, pid_c, chunk_size, and the strides as 32-bit ints based on each operand's own value at specialization time, not the value their product will reach. Both pid_b * stride_dt_batch and pid_c * chunk_size * stride_dt_seqlen are separate multiplications that are each carried out in 32-bit, since their leading operand (pid_b/pid_c) is 32-bit. Once either product exceeds 2**31 - 1, it silently wraps to a garbage value, corrupting the pointer offset added to dt_ptr -- and the kernel then reads/writes through a wildly out-of-bounds address, surfacing as a CUDA "illegal memory access" error that's far removed from the actual root cause.

For an ordinarily contiguous dt (shape (batch, seqlen, nheads)), both stride(0) and stride(1) stay small, so neither product gets close to overflowing at any realistic seqlen/batch. It only becomes a real problem when the caller passes a dt that is a non-contiguous view into a much wider parent tensor -- e.g. transformers' NemotronHMamba2Mixer slices dt/B/C directly out of a much wider in_proj/post-conv1d output. In that case both strides are the parent's width rather than nheads, and either product can overflow int32 well within realistic sequence lengths -- the pid_c term at large seqlen with batch=1 (real width ~18,560 for the model that first surfaced this), and the pid_b term independently at batch >= 2 with much shorter sequences (reported: batch=4, seqlen=40_960, parent_width=35_072, from Nemotron-H Ultra).

We verified this arithmetic against every empirical pass/fail boundary we found, including an exact sign-flip at a single-token transition (N=115,712 passes, N=115,713 fails, for one specific stride_dt_seqlen value) -- this is confirmed mechanism, not a guess. We also verified with a small standalone Triton kernel that casting the leading operand of a chained multiplication to int64 is sufficient to make the whole chain safe (Triton's promotion rules propagate the width forward through subsequent multiplications), while casting only a trailing operand does not help, since the overflow already happens in an earlier sub-product before that operand is multiplied in. This also means pid_b and pid_c each need their own cast, since they lead two separate multiplications rather than one chain -- widening one does not protect the other (this was originally missed; see comment thread on the follow-up audit PR #988 for the report and repro).

Fix

Cast pid_b and pid_c to tl.int64 immediately after their respective tl.program_id() calls in both kernels, so each chained multiplication is carried out entirely in 64-bit from the first multiplication onward. This mirrors vLLM's own vendored copy of this kernel (vllm/model_executor/layers/mamba/ops/), which has had this exact cast (for pid_c) since the file was first added in PR #10909 ("Add Bamba Model", commit aff404571b0d5aba342c46fdf5d7f8a251da9383) -- inherited while porting the kernels from state-spaces/mamba, not added in response to a specific bug report.

See also the related, still-open upstream Triton issue: triton-lang/triton#1058

The casts are unconditional and cheap, so they should not regress the common small-sequence/contiguous-tensor case -- benchmark results confirming this are in a follow-up comment below.

_chunk_cumsum_bwd_kernel has the identical pattern and receives the identical fix here, though it was not the kernel that produced the original crash.

Broader audit: the same class of vulnerability across the rest of mamba_ssm/ops/triton/ (ssd_combined.py, ssd_state_passing.py, ssd_chunk_scan.py, ssd_bmm.py, layer_norm.py, layernorm_gated.py, k_activations.py) is tracked separately in #988.

Test plan

• Added test_chunk_cumsum_fwd_bwd_noncontiguous_wide_view_no_overflow in tests/ops/triton/test_ssd.py, reproducing the exact failing shape (seqlen=115_866, nheads=128, chunk_size=128, parent_width=18_560, batch=1) and exercising both the forward and backward kernel.
• Added test_chunk_cumsum_fwd_bwd_noncontiguous_wide_view_batch_axis_no_overflow covering the separate pid_b overflow (batch=4, seqlen=40_960, parent_width=35_072), comparing output/gradients against the same kernels run on a contiguous copy of identical values (an isfinite-only check can't reliably catch a wrapped-but-still-in-bounds pointer offset, since it can produce finite but silently wrong values rather than a crash or a NaN).
• Verified both tests fail with RuntimeError/AcceleratorError: ... an illegal memory access was encountered on the unpatched kernels, and pass with the fix applied.
• Benchmarked _chunk_cumsum_fwd across a range of typical contiguous shapes before/after the fix -- no regression (results in a follow-up comment).

[chipsncheese-ie]

Regression benchmark

To check the unconditional int64 cast doesn't slow down the common (small-sequence, contiguous-tensor) case, I benchmarked _chunk_cumsum_fwd with torch.utils.benchmark.Timer.blocked_autorange (auto-picks iteration count per shape to fill ~5s), computing mean/median/std and a 95% CI (mean +/- 1.96 * stderr) with numpy.

benchmark_chunk_cumsum_int64_fix.py

#!/usr/bin/env python3
"""Benchmark _chunk_cumsum_fwd/_chunk_cumsum_bwd to check for regressions from
the int64 pid_c cast added in ssd_chunk_state.py (fix for a 32-bit pointer
arithmetic overflow -- see the commit that introduced this file for details).

The cast is unconditional, so the case to check for a regression in is the
*common* one: an ordinarily-contiguous `dt` at realistic seqlens, where the
overflow this fixes could never have occurred in the first place. We time
both the forward cumsum and its autograd backward across a range of
seqlens/nheads/chunk_sizes typical of real usage, plus one non-contiguous
wide-view shape (the shape that originally crashed) to confirm it's now both
correct and not pathologically slow.

Usage:
    CUDA_VISIBLE_DEVICES=<idx> python benchmarks/benchmark_chunk_cumsum_int64_fix.py

To compare against the unpatched kernel, run this script once with the fix
applied and once with `git checkout <pre-fix-sha> -- ssd_chunk_state.py`
(keeping this script), and compare the printed tables. Note the unpatched
kernel will crash outright on the non-contiguous wide-view shape -- that's
the bug, not something this script needs to benchmark around.

Each shape is timed with torch.utils.benchmark.Timer.blocked_autorange,
which auto-picks how many iterations fit in MIN_RUN_TIME and reports
mean/median/std, rather than a single fixed loop count.
"""
import numpy as np
import torch
import torch.utils.benchmark as torch_benchmark

from mamba_ssm.ops.triton.ssd_chunk_state import _chunk_cumsum_fwd

DEVICE = "cuda"
N_WARMUP = 5
MIN_RUN_TIME = 5.0  # seconds, per shape

# (seqlen, nheads, chunk_size, parent_width_or_None)
# parent_width=None means dt is allocated directly (contiguous, stride(1) == nheads).
# parent_width=int means dt is a non-contiguous slice out of a wider parent tensor,
# reproducing the shape that originally overflowed int32.
SHAPES = [
    (2_048, 64, 128, None),
    (8_192, 64, 128, None),
    (8_192, 128, 128, None),
    (32_768, 128, 128, None),
    (115_866, 128, 128, None),        # same seqlen as the crash, but contiguous
    (115_866, 128, 128, 18_560),      # the exact shape that crashed pre-fix
]


def make_dt(seqlen, nheads, parent_width):
    if parent_width is None:
        return torch.randn(1, seqlen, nheads, dtype=torch.bfloat16, device=DEVICE)
    parent = torch.randn(1, seqlen, parent_width, dtype=torch.bfloat16, device=DEVICE)
    return parent[:, :, -nheads:]


def bench_one(seqlen, nheads, chunk_size, parent_width):
    torch.manual_seed(0)
    dt = make_dt(seqlen, nheads, parent_width)
    A = -torch.exp(torch.randn(nheads, dtype=torch.float32, device=DEVICE))
    dt_bias = torch.randn(nheads, dtype=torch.float32, device=DEVICE)

    def run():
        with torch.no_grad():
            return _chunk_cumsum_fwd(dt, A, chunk_size, dt_bias=dt_bias, dt_softplus=True)

    for _ in range(N_WARMUP):
        run()
    torch.cuda.synchronize(DEVICE)

    dA_cumsum, dt_out = run()
    torch.cuda.synchronize(DEVICE)
    assert torch.isfinite(dA_cumsum).all()
    assert torch.isfinite(dt_out).all()

    timer = torch_benchmark.Timer(
        stmt="run()",
        globals={"run": run},
    )
    measurement = timer.blocked_autorange(min_run_time=MIN_RUN_TIME)
    return measurement


def main():
    header = (
        f"{'seqlen':>8} {'nheads':>7} {'chunk':>6} {'parent_width':>13} {'contig':>7} "
        f"{'blocks':>7} {'mean_ms':>10} {'median_ms':>11} {'std_ms':>9} {'95% CI (ms)':>22}"
    )
    print(header)
    for seqlen, nheads, chunk_size, parent_width in SHAPES:
        measurement = bench_one(seqlen, nheads, chunk_size, parent_width)
        # raw_times are per-block totals (each block runs number_per_run
        # iterations); normalize to per-iteration ms before averaging.
        times_ms = np.array(measurement.raw_times) * 1e3 / measurement.number_per_run
        n_blocks = times_ms.size
        mean_ms = float(np.mean(times_ms))
        median_ms = float(np.median(times_ms))
        std_ms = float(np.std(times_ms, ddof=1))
        stderr_ms = std_ms / np.sqrt(n_blocks)
        ci95_halfwidth = 1.96 * stderr_ms
        contig = "yes" if parent_width is None else "no"
        pw = "-" if parent_width is None else parent_width
        ci_str = f"[{mean_ms - ci95_halfwidth:.4f}, {mean_ms + ci95_halfwidth:.4f}]"
        print(
            f"{seqlen:>8} {nheads:>7} {chunk_size:>6} {pw!s:>13} {contig:>7} "
            f"{n_blocks:>7} {mean_ms:>10.4f} {median_ms:>11.4f} {std_ms:>9.4f} {ci_str:>22}"
        )


if __name__ == "__main__":
    main()

Results (patched vs. pre-fix baseline, contiguous shapes; mean ms/iter with 95% CI):

seqlen	nheads	baseline (pre-fix)	patched (this PR)
2,048	64	0.0614 [0.0611, 0.0617]	0.0590 [0.0590, 0.0591]
8,192	64	0.0604 [0.0601, 0.0606]	0.0607 [0.0604, 0.0611]
8,192	128	0.0637 [0.0635, 0.0638]	0.0607 [0.0604, 0.0610]
32,768	128	0.0816 [0.0816, 0.0816]	0.0809 [0.0809, 0.0810]
115,866	128	0.3980 [0.3979, 0.3981]	0.3995 [0.3994, 0.3996]

All differences are within a few microseconds — no meaningful regression from the cast on the common path. For reference, the non-contiguous wide-view shape that crashes outright pre-fix (seqlen=115,866, parent_width=18,560) completes in 0.7724ms [0.7724, 0.7725] with the fix.

(Baseline was measured by checking out the pre-fix blob of ssd_chunk_state.py on top of this branch, keeping this benchmark script, running it against the contiguous shapes only, then restoring the fix — the unpatched kernel crashes outright on the non-contiguous shape so that one isn't in the baseline comparison.)

[KyleMylonakisProtopia]

Any update on landing this?

[Chessing234]

int32 overflow here is nasty. a one-line comment near the cast noting the seq*dim threshold that overflows would help.

[chipsncheese-ie]

Benchmark re-check after adding the pid_b cast

Re-ran the same regression benchmark from the earlier comment on this PR, now with both pid_b and pid_c cast to int64 (the latest commit).

seqlen	nheads	contiguous	mean ms (95% CI)
2,048	64	yes	0.0606 [0.0605, 0.0606]
8,192	64	yes	0.0622 [0.0619, 0.0625]
8,192	128	yes	0.0643 [0.0641, 0.0644]
32,768	128	yes	0.0826 [0.0826, 0.0827]
115,866	128	yes	0.3958 [0.3957, 0.3959]
115,866	128	no (wide view)	0.7769 [0.7769, 0.7769]

These are within noise of the earlier pid_c-only numbers (e.g. 0.0590–0.0614ms and 0.3980–0.3995ms for the same shapes) -- the additional pid_b cast doesn't introduce any measurable regression on top of the original fix.