vLLM CUDA Integer Overflow

While training Jamba 3B with GRPO, a problem arose where the logprob computed by rollout didn't match the value recomputed by FSDP on the same sequence. At first it was unclear whether the cause was in training, inference, or synchronization, but the actual cause was a silent integer overflow in an internal vLLM CUDA kernel.

The problem surfaced through a sanity check commonly used in RL systems: comparing log π_old, saved by the rollout engine, against log π_train, recomputed by the training model on the same input. These two values should be nearly identical given the same weights and the same input, but at certain steps the discrepancy periodically spiked, revealing the anomaly.

The key turning point in debugging was rollouts_num. With the default setting, spikes appeared only around step 12, but when the number of rollouts was increased to 8, 16, 32, 64, 128, the spike period shifted to match that value exactly. Notably, at 128, the discrepancy appeared from the very first rollout, showing that the problem lay not in training dynamics but in the rollout path.

From this point, the issue was narrowed down from an RL training bug to an inference-engine bug. After a standalone vLLM reproduction and further ablations, it was confirmed that lowering GPU memory utilization to 50% made the problem disappear, while it reproduced above that threshold. Subsequently, the problem did not occur with attention-only models, so the cause converged on the Mamba-1 selective scan CUDA kernel rather than the KV-cache.

The actual defect was in the kernel's pointer arithmetic. index_t in SSMParamsBase was defined as uint32_t, and the product of cache_index and ssm_states_batch_stride was computed in 32-bit, silently wrapping around the moment it exceeded UINT32_MAX. When the batch stride was 89,600, the overflow occurred once cache_index exceeded roughly 47,935, and about 31% of the total 69,776 cache slots ended up being written to the wrong memory location, leaving the intended slot at zero.

The fix was just two characters.

• using index_t = uint32_t;
• using index_t = size_t;

This change promoted the stride multiplication to 64-bit arithmetic, eliminating the overflow. The conclusion is clear: in distributed RL, a logprob mismatch that appears on the surface to be training instability can actually be caused by a low-level CUDA bug underneath, and the best way to solve such a problem is to find the variable that changes the structure of the failure and narrow it down to a step-zero reproduction.