vLLM and Jamba: debugging silent inference bugs

Asaf Gardin and Yuval Belfer describe two inference bugs they encountered at AI21 while working with Jamba, a hybrid of attention and Mamba layers. The talk covers debugging methods relevant to self-hosted model serving, where incorrect output can arrive without a crash or warning.

In the first case, they report rare gibberish under load in vLLM. Reducing GPU memory utilization from 90% to 20% and sampling at temperature zero made the failure repeatable. They compared token log probabilities against Hugging Face Transformers, inspected the Mamba prefill kernel and ran NVIDIA Compute Sanitizer. Passing request IDs into the forward context finally exposed a scheduling error: a fresh request entered decode before prefill. Mamba then read stale state from earlier requests. Their fix classified requests with no computed tokens as prefill.

The second case involved logprob spikes every twelve training steps, before any weight update. Increasing rollouts per prompt from eight to 128 brought the failure to step one. Reducing memory hid this bug because the smaller state buffer avoided an overflowing unsigned 32-bit index. The speakers say replacing that index type with size_t resolved it. Their practical advice is to compare against a reference implementation and vary settings that change when or where a failure occurs.