Vllm Debugging Mamba Bug

One Token to Corrupt Them All: A vLLM Debugging Tale

Skip to Main Menu

Skip to Main Content

Skip to Footer

Back to Blog

-->

Back to Blog

TL;DR

While working on a new Jamba model, we noticed it would generate complete gibberish, but just once out of every thousand prompts. Finding and solving the bug sent us deep into the heart of vLLM, eventually touching on how its scheduler interacts with different model architectures. In addition to sharing our fix below, we also share the lessons we learned along the way about debugging a massive codebase like vLLM. We hope this detailed walkthrough makes your next vLLM debugging session a little less daunting.

The problem: gibberish in the haystack

At AI21 Labs, we build and train our own Jamba series of LLMs in-house from scratch. While working on Jamba Reasoning 3B , we noticed something concerning in our reinforcement learning (RL) training pipeline: Our model would occasionally generate complete gibberish. Not even subtle degradation. Just pure nonsense.

We caught it by monitoring max logprobs across generations. Logprobs measure model confidence – values closer to 0 indicate high certainty. Normally, bad generations come with low confidence. But our gibberish had high confidence scores. The model was confidently wrong, signaling something deeply broken rather than just poor output quality.

Each spike represents gibberish generated by the model. The x-axis represents the training step number, and the y-axis represents the absolute logprob difference between the vLLM and transformers outputs.

It was clear the bug wasn’t coming from the model: The same checkpoints worked perfectly with Hugging Face’s transformers, the issue only appeared under specific runtime conditions, and it occurred sporadically after hundreds of requests rather than consistently from the start. This pattern pointed us in the direction of vLLM and, specifically, its request scheduling and cache management.

The bug was serious enough that we couldn’t release the model as long as it persisted. We needed to fix it. The problem was, it would only occur sporadically – think once in a thousand prompts. Finding the needle in this haystack would require patience, methodical debugging, and eventually, instrumenting vLLM itself. You can find the result of all of this work now live in a merged fix to the vLLM project.

If you’re running inference with vLLM, whether for production serving, RL training, or evaluation pipelines, understanding how the scheduler interacts with model architectures like Mamba can save you from silent data corruption. Below, we’re sharing our step-by-step process for debugging vLLM’s massive codebase and the lessons we learned along the way.

The debug script

To systematically detect and measure these failures, we built a comparison script that would become our primary diagnostic tool throughout this investigation, as it allowed us to compare logprobs for the same token IDs; that is, instead of re-generating with transformers, we determine what transformers would have assigned to vLLM’s outputs.

Under normal conditions, vLLM and transformers should produce nearly identical logprobs for the same tokens (minor floating-point differences aside). But when vLLM generates garbage due to state corruption, the logprobs diverge dramatically: vLLM might report high confidence for a nonsense token, while transformers correctly shows it should have been extremely unlikely. This discrepancy reveals that vLLM is “confident” about tokens that make no sense in the true sequence context.

The comparison script works in two phases:

Generate with vLLM: Send batches of prompts through vLLM, collecting the generated tokens and their logprobs.

Verify with transformers: Run the same prompts+their generations through Hugging Face transformers, computing reference logprobs for the exact same token sequences.

Pseudo-code for the debug comparison script

1. Load model

model_name = "ai21labs/AI21-Jamba-Reasoning-3B" vllm_model = load_vllm(model_name) hf_model = load_transformers(model_name) batch_size = 128

2. Generate with vLLM and capture logprobs

prompts = ["Prompt0", "Prompt1", ...] # 1024 prompts all_vllm_outputs = [] for batch_idx in range(0, len(prompts), batch_size): batch = prompts[batch_idx : batch_idx + batch_size] vllm_outputs = vllm_model.generate(batch, return_logprobs=True) all_vllm_outputs.extend(vllm_outputs)

3. For each output, feed prompt + generated tokens to HF and get logits

for prompt_idx, vllm_out in enumerate(all_vllm_outputs): full_sequence = prompt + vllm_out.generated_tokens hf_logits = hf_model.forward(full_sequence) hf_logprobs = compute_logprobs(hf_logits)

4. Compare vLLM's generation logprobs vs HF's logprobs for same sequence

for token_idx, (v_logprob, h_logprob) in enumerate(zip(vllm_out.logprobs, hf_logprobs)): diff = abs(v_logprob - h_logprob) if diff > threshold: print(f"Mismatch at prompt {prompt_idx}, token {token_idx}") print(f" vLLM: {v_logprob}, HF: {h_logprob}, diff: {diff}")

Reproducing the unreproducible

Our first goal sounded simple: Reproduce the issue reliably. The path there proved to be easier said than done, though.

We started with our RL dataset – thousands of prompts of varying lengths – and ran inference using ai21labs/AI21-Jamba-Reasoning-3B . We sent 1024 prompts in batches of 128, mirroring our RL setup.

Nothing. Clean generations across the board.

Implementing memory constraints

We thought about what conditions in RL might differ from our test setup. Then it struck: memory pressure.

During RL training, the GPU memory is heavily utilized, meaning the SSM state cache is nearly full and cache slots get recycled aggressively. Unlike attention models where stale KV cache can’t corrupt new sequences (thanks to sequence length masking), Mamba’s SSM state is a compressed representation of an entire sequence’s history. If a new request reads stale state from a previous sequence, it’s like starting a conversation with someone else’s memories: the state accumulates recursively, so garbage at the start corrupts everything that follows.

The hidden state h h is updated recurrently. If a new request is misclassified as a decode, it loads a stale h h from the previous occupant of that memory slot before it ever performs its first update. Image from Maarten Grootendorst.

vLLM allows control over GPU...