Vllm Cuda Integer Overflow

How a 32-bit Overflow Corrupted a vLLM CUDA Kernel

Skip to Main Menu

Skip to Main Content

Skip to Footer

Back to Blog

-->

Back to Blog

TL;DR

While training our Jamba 3B model with GRPO, we hit a mysterious logprob mismatch between rollout and training. The root cause turned out to be a silent integer overflow deep inside a vLLM CUDA kernel – one that only triggered when the number of cache slots exceeded ~47,935. The fix was two characters away. The journey to find it was not.

One of the trickiest parts of debugging reinforcement learning systems is that the bug almost never introduces itself politely.

In reinforcement learning from human feedback, correctness is load-bearing in ways that aren’t always visible. A common sanity check is to compare the log probabilities which the rollout engine assigned to its own completions, against what the freshly-loaded training model computes on the same sequences – before any weight update has happened. In a healthy system, these should be nearly identical: same weights, same inputs, same numbers. When they diverge, it’s a canary signal that something upstream has already gone wrong.

Recently, while working on our GRPO training setup for our Jamba 3B model, we ran into exactly this kind of issue. Jamba is AI21’s hybrid attention-Mamba model – a 3B parameter architecture that combines standard attention with Mamba state-space layers. That hybrid nature, as we’ll see, is exactly what made this bug so hard to find. The training stack included Ray for orchestration, an LM training path with FSDP for weight updates, and a rollout path used for generating completions via vLLM and scoring them using our internal evaluation service. What we saw was a mismatch in log probabilities between rollout and the FSDP side. At first glance, it was not obvious whether the problem came from training, inference, synchronization, or something more subtle in the interaction between them.

Spoiler alert: It was a typical unsigned integer overflow of an index deep inside a CUDA kernel. The interesting part was the journey , not the destination, so let’s dig in.

The first question: where does the bug actually live?

In a system like this, GRPO is not really one thing, it is not your average model.fit(). It is a pipeline made of distinct stages, each with its own assumptions:

rollout generation

FSDP-based training

distributed coordination through Ray

handoff between inference-time and training-time components

In a typical GRPO training flow the process looks as below:

Figure 1: A typical GRPO training loop using vLLM and FSDP. vLLM handles rollout generation and computes log π_old; the FSDP model recomputes log π_train on the same completions before any weight update. The mean absolute difference between the two (sanity check, right) should be ≈0 at the start of each iteration.

When logprob spikes between the vLLM engine and those of the just-loaded model in FSDP started appearing – before any weight updates – the natural temptation was to inspect everything at once. But that usually leads nowhere. In distributed RL systems, if you do not reduce the search space aggressively, you end up debugging the entire stack instead of debugging the bug.

So the real task became: can we determine whether this issue belongs to the rollout side, the FSDP side, or the weight sync in-between?

Conventional debugging wisdom says: find a single hyperparameter that lets you disturb the system in a controlled way, and observe whether the symptom moves with it.

Finding the lever

Watching how the logprobs difference behaved, a key observation was that the spikes were not random. They appeared periodically – recurring damage and recovery at a fixed interval. With our default experiment configs, the spike didn’t occur until roughly step 12, which added more questions than it answered.

Figure 2: Mean absolute logprob difference between vLLM rollout and FSDP recompute (rollouts_num=8, default config). Spikes appear periodically – roughly every 12 steps – rather than randomly, suggesting a structured, recurring source of error rather than training noise.

From looking at the model rollouts around those steps, there were a number of gibberish completions, but not all. We took the weights from such steps and tried reproducing the issue in a standalone vLLM instance. The generation looked fine. Dead end.

We then started a methodical search for a lever. The eureka moment came when we ran experiments with a doubled number of rollouts per prompt – 8, 16, 32, up to 128 – and found that the periodicity of the spikes tracked exactly with that number. That changed the debugging story completely.

If the logprob diff spikes move in a pattern tied to rollout count, then the bug is very unlikely to be some generic instability in FSDP training. The symptom is being structured by rollout behavior. That doesn’t fully prove where the bug lives, but it gives you a much sharper hypothesis: the issue is probably introduced before the training path ever consumes the data.

That was the turning point.

Figure 3: Spike periodicity as a function of rollouts_num (8, 16, 32, 64, 128). As the number of rollouts per prompt increases, spikes shift earlier in training. With rollouts_num=128, the spike triggers on the very first step – a key diagnostic signal pointing away from training dynamics and toward the rollout path itself.

From “training instability” to “reproducible rollout issue”

Before that observation, the issue looked like an RL-training bug. After it, it looked more like a rollout-path bug that happened to surface during RL training.

That distinction matters a lot. A full RL training run is a terrible place to debug low-level correctness issues. There are too many layers involved, too much asynchronous machinery, too many places where the original signal gets blurred. The ideal debugging path is always to peel away layers until the issue survives in the smallest possible environment.

So that became the goal: reduce the problem from “something is wrong during GRPO training” to “we can reproduce this on demand, on the earliest possible rollout.”

The next milestone was getting the issue to reproduce at step zero – and looking closely at the chart, the 128 rollout-per-prompt run spiked on the very first rollout.

Figure 4: Zoomed view of step-one...