AI BriefingDebugging a Mamba Bug in vLLM
Key point
The cause of rare gibberish was traced to vLLM's Mamba scheduler and cache.
Details
AI21 Labs had Jamba Reasoning 3B running in an RL pipeline and discovered that the model would occasionally spit out complete gibberish. Unlike typical low-quality output, the max logprobs were actually high, which raised suspicion of runtime state corruption rather than simple model degradation.
Running the same checkpoint with Hugging Face transformers worked fine, and the issue only appeared intermittently with vLLM. So the authors built a debug script that compared vLLM's output against the transformers baseline for the same token sequence, and tracked the logprob differences for the actually generated tokens.
At first the bug was hard to reproduce, but when they lowered gpu_memory_utilization to 0.2 to simulate the memory pressure of the RL environment, the issue surfaced. After that, the gibberish reproduced on the 854th request, and it repeated identically at temperature=0, confirming the bug was deterministic.
Suspicion turned to the CUDA prefill kernel handling Mamba SSM state. If the state wasn't properly initialized for a new request, or if it incorrectly read a previous request's cache slot, all subsequent tokens could be corrupted. However, bound checks, pointer arithmetic verification, compute-sanitizer checks, and inspection of the SSM state across all 28 Mamba layers all came back clean.
Next, they suspected vLLM's prefill/decode split logic. This was especially relevant for Mamba, since whether a request is assigned to prefill or decode is directly tied to how its state gets initialized. A key clue came from comparing this with the older V0 engine: due to a bug in V0, Mamba1 effectively never went through the decode kernel and only ever used the prefill path.
When they forced all requests to go through prefill only in V1 as well, the gibberish disappeared. This meant the problem wasn't in the Mamba state itself, but in the interaction between the decode path and the scheduler, and subsequent debugging narrowed further onto the decode kernel and the logic surrounding it.
This summary was generated automatically by AI. Check the original for the author's claims and context. Copyright belongs to the original author.
Our guide explains how the AI works. Report summary errors, attribution issues, or removal requests via Contact.