Sequences longer than ~61K tokens (and often less) silently give a wrong answer or crash

Summary

There's a 32-bit integer overflow in the pointer arithmetic of nearly every Triton kernel in the core chunked-scan path (ssd_chunk_state.py, ssd_chunk_scan.py, ssd_combined.py, ssd_bmm.py, plus the same root cause in layer_norm.py/layernorm_gated.py/k_activations.py). Past a threshold that depends on sequence length, batch size, hidden/projection width, and GPU count, kernels compute a pointer offset in 32-bit that should be 64-bit, silently wrap, and either:

• crash with CUDA error: an illegal memory access was encountered, or
• produce finite, plausible-looking, silently wrong output — no crash, no NaN, no warning.

The second failure mode is the one worth worrying about: any long-context training or eval run that didn't crash is not automatically correct.

~61K is the threshold for the narrowest, most conservative case (batch=1, single GPU) on Nemotron 3 Ultra. Any realistic batch size per GPU pushes the real threshold considerably lower — the reported real-world crash below already happened at 40,960.

Concrete repro

Both overflow terms below are demonstrated on Nemotron 3 Ultra (550B-A55B), a real NVIDIA model with a hybrid Mamba2-Attention MoE architecture (public config): hidden_size=8192, expand=2, mamba_num_heads=256, mamba_head_dim=64, n_groups=8, ssm_state_size=128.

Ultra feeds dt/x/B/C/z into the Mamba2 mixer as a non-contiguous view into a much wider fused in_proj tensor (z, x, B, C, dt all concatenated and projected together, then sliced back apart). This is the trigger: an ordinarily-contiguous tensor never gets close to either threshold at any realistic sequence length, which is why this hasn't shown up in typical unit tests.

Fused in_proj width = 2×d_inner + 2×(n_groups×ssm_state_size) + mamba_num_heads, where d_inner = expand × hidden_size = 16,384:

d_in_proj = 16,384 + 16,384 + 1,024 + 1,024 + 256 = 35,072

1. Chunk-axis overflow (present even at batch size 1, single GPU):

(nchunks - 1) * chunk_size * d_in_proj > 2**31 - 1

We confirmed this exactly on GPU, at Ultra's own d_in_proj = 35,072 and chunk_size = 128: seqlen = 61,312 passes, seqlen = 61,313 fails with the same illegal memory access error — a single-token transition, not a fuzzy neighborhood.

General form: seqlen_threshold ≈ (2**31 - 1) / d_in_proj — inversely proportional to hidden/projection width. A wider model breaks at a shorter sequence length; a narrower model can go longer before hitting it.

2. Batch-axis overflow (batch ≥ 2) — reported crash on Ultra, on a single GPU (no data/batch parallelism):

(batch_per_gpu - 1) * seqlen * d_in_proj > 2**31 - 1

Concretely: batch=4, seqlen=40_960, d_in_proj=35_072, all four batch elements on one device → stride_dt_batch = seqlen * d_in_proj = 1_436_549_120, so pid_b=2 alone gives 2_873_098_240 > 2**31 - 1. This crashed inside _chunk_cumsum_fwd_kernel (and its backward twin) during activation-checkpoint recomputation.

General form: seqlen_threshold ≈ (2**31 - 1) × num_gpus / (batch_size × d_in_proj)

This term is:

• inversely proportional to batch size — a larger batch hits it at a shorter sequence length, not a longer one. Doubling batch size roughly halves the sequence length needed to trigger this.
• inversely proportional to hidden/projection width, same as the chunk-axis term.
• proportional to number of GPUs — under ordinary data/batch parallelism, each GPU's kernel launch only ever sees batch_size / num_gpus batch elements locally (pid_b only ranges over what's materialized on that device). So the effective per-GPU batch that actually matters here shrinks as GPU count grows for a fixed global batch size, pushing the sequence-length threshold higher. The reported numbers above are for a single GPU (num_gpus=1) — running the same global batch across more GPUs raises this particular threshold, but does not affect the chunk-axis term at all, which depends only on sequence length and width.

Root cause

Triton's JIT specializes each operand of a pointer-arithmetic expression (pid_c, pid_b, chunk_size, strides, etc.) as a 32-bit int based on that operand's own value at specialization time — not the value their product will reach. Standard integer promotion means the whole chained multiplication stays 32-bit as long as the leading operand is 32-bit, even if a later operand is large. Once the true product exceeds 2**31 - 1, it wraps to a garbage offset.

This is a known, still-open upstream Triton issue: https://github.com/triton-lang/triton/issues/1058

vLLM's own vendored copy of these kernels already casts the leading operand to tl.int64 for exactly this reason (vllm/model_executor/layers/mamba/ops/, present since the file was first added), inherited while porting from this repo — it just was never carried back upstream here.

Fix

Cast the relevant program-id-derived index to tl.int64 immediately after tl.program_id(), so the whole chained multiplication is carried out in 64-bit from the start. Cheap, unconditional, no measurable regression on the common contiguous/short-sequence case (benchmarked).

We've fixed all of this (every affected kernel across ssd_chunk_state.py, ssd_chunk_scan.py, ssd_combined.py, ssd_bmm.py, layer_norm.py, layernorm_gated.py, k_activations.py) with regression tests.

We're splitting the fix into smaller, per-file PRs for easier review, starting with #1038 (ssd_chunk_state.py). The rest is ready but not yet opened as PRs -- it's sitting on branches in our fork (protopia-ai/mamba-ssm) that we'll turn into PRs incrementally:

• fix/int32-overflow-ssd-stacked -- the remaining chunked-scan kernels (ssd_bmm.py, ssd_state_passing.py, ssd_chunk_scan.py/ssd_combined.py, and the causal_conv1d addressing guard)
• fix/int32-overflow-norm-activation -- layer_norm.py/layernorm_gated.py/k_activations.py
• misc/torch-compile-autotune-fix-and-ci-deps -- an unrelated torch.compile/autotune correctness bug found along the way, plus some CI/dependency housekeeping

Affected versions

Present in the current released kernels; not tied to a specific recent regression — this looks like it's been present since these kernels were first written.