[
  {
    "startMs": 17820,
    "text": "So put your model somewhere. Put your agent. You're hoping for the best. You're waiting for something to crash. Everything looks good. Everything's fine, and you see this."
  },
  {
    "startMs": 33340,
    "text": "And that's the problem, right? In these type of bugs, there is no crash. There's no warning, no error, and there's high confidence. That's not a quality issue 'cause, uh, right, this is something that you don't really know what and how and why. Um, welcome to this talk. My name is Yuval. This is Asaf. And together we're gonna take you on a journey of how we ended up"
  },
  {
    "startMs": 63040,
    "text": "fixing those type of bugs. A little bit about us. We work at AI21, which is an AI research lab. We started as a foundation model company, and most famously known for Jamba, which is a hybrid architecture between Transformers and Mamba, which is an SSM, uh, state. And while we were doing those, while we were training those models, while we were shipping those models into production and had users and we had a workload, we"
  },
  {
    "startMs": 93300,
    "text": "got into several interesting bugs, and these are the bugs that I think are the hardest to deal with because this is not a quality problem. It's not something you can take your research team and try to optimize or solve or make the model be better at something. This is an engineering problem. This is an issue where, that there is high confidence, uh, but the output is bad. So"
  },
  {
    "startMs": 123660,
    "text": "let's break... Let's dive deep to the first case, what we call the imposter request, where just to set up the scene, what are we talking about? We are talking about how we, during the training of our Jamba model, more specifically, we did GRPO, which is a type of RL training, and again, this is a hybrid model, layers of Mamba and attention. And just to make sure we're all aligned, what is the type of a request? So life cy- life cy-"
  },
  {
    "startMs": 153180,
    "text": "life of a request, so we start with a prompt, tokeni- tokenization, and then in the forward pass, we're doing both prefill and then decode. After that, we finish the forward pass, uh, detokenization to go back to text. And the thing about the crime here is that it's bad on so many levels, but mainly on these three. This is what we call the one in thousand gibberish. It's not something that will"
  },
  {
    "startMs": 182720,
    "text": "happen in the first five hundred or nine hundred requests, but it will happen in the one thousand, which is rare enough to duplicate it easily, but not, uh, right, too common to ship it. Sorry. Uh, also, it only happened in vLLM, not in other, uh, not in any other inference framework, and it's something which is late onset. It's not something that will happen if you have only few requests. You"
  },
  {
    "startMs": 212720,
    "text": "need some sort of workload. So it's rare, it's late onset, and it's very engine specific. It's a very, very hard task, and we had to bring one of our best detectives to handle that. So I'll give it to Asaf to explain how."
  },
  {
    "startMs": 227700,
    "text": "All right. Hey, guys. Thank you, Yuval. Uh, thanks, Yuval. Uh, so we're gonna start with, um, trying to reproduce something with-- that was very difficult to reproduce. Um, and basically, um, vLLM has got a lot of, um, a lot of flags, a lot of the CLI flags, and a lot of knobs you guys can alter and tweak. Um, and one of the things that helped us understand how to even reproduce it, because when we tried to reproduce it on the first time, just sending prompts here, prompts there, a few batches, um, it didn't really help us manage to get gibberish back from our model. The model"
  },
  {
    "startMs": 257260,
    "text": "responded back just fine. So what we did was, is we tried to make it happen in a very, very short, um, amount of time, so we'll get a, a quick feedback loop when we try to debu- uh, debug it. So what we did was, is we took one of the, one of the, um, most default and most common, um, g- um, flag that v- the vLLM allows you to play with, which is GPU memory utilization, which basically allows you to, uh, out, to choose how much memory, how much GPU memory you want to allocate for your, for your weight, for your activations, and for your, um,"
  },
  {
    "startMs": 287340,
    "text": "KV cache and so on, and reduced it from, uh, ninety percent to twenty percent. And once we did that, and then we started, uh, running a lot of, uh, requests, uh, simultaneously, um, uh, all of a sudden, request number, let's say eight hun- eight hundred and fifty-four suddenly returned, uh, gibberish. And when we did that, we, uh, we sampled all of the batches with temperature zero, so we'll be able to deterministically and constantly get the same request to g- to return, to return and respond with gibberish."
  },
  {
    "startMs": 318080,
    "text": "Um, so, um, like Yuval said, it happened only in vLLM, and, um, we used, um, in order to, uh... A, another way to reproduce it and to understand where the issue really came from, we used, um, uh, Hugging Face's Transformers as a baseline since, uh, Transformers is a very, had a very, um, vanilla and, uh, plain implementation of our Mamba kernels, as opposed to vLLM, which, uh, all the kernels and all the engine have gone through a lot of changes and modifications to support a lot of, um, cool features that vLLM"
  },
  {
    "startMs": 347760,
    "text": "supports. So we used Transformers as our baseline to understand whether or not there is an issue with our, um, with our inference or not, with the model or not. So what we did was, we took vLLM and we sent, uh, all of our prompts through vLLM, and we generated a response, uh, all the responses. We got it... And now in our hands we have the response along with the logprobs, 'cause v- 'cause in vLLM, you're, you, you're able to get your logprobs out and inspect them. Then what we did was we took the full sequence, the prompt and the generation, and we, uh, uh, passed it over to,"
  },
  {
    "startMs": 377988,
    "text": "to, uh, to Hugging Face's forward pass. But what all we did was run just the prefill, um, and then we, uh, uh, sample... We took this, the logits out of the prefill, um, um, response. We ran it through softmax, and then we were able to, um, to compare the divergence, uh, in the distributions of our tokens. Um, that's a short, uh, pseudocode of how that look like. You can see here that we, uh, take up the prompt. Uh, we, we run it through gen- uh, vLLM's generate. We get the, uh, the"
  },
  {
    "startMs": 407728,
    "text": "response back along with the, uh, with the log probs. We pass it over to, to Hugging Face's, uh, forward pass. We only run it with prefill. Um, we've created some function called com- compute log, log probs, which, um, runs the softmax. Um, then you calculate the difference between them, and then you'll be able to, uh, tell the divergence between every one of the tokens' log probs. All right. So now that we have the tools in our hand to understand how... where the issue could maybe come from, um, we started to look at different"
  },
  {
    "startMs": 437388,
    "text": "suspects in vLLM's engine. So the first thing we looked at was, um, the CUDA prefill kernel of Mamba. Um, we looked at it. We, uh, inspected all of the, all of the math that's being done here, th- that's being done there. Um, and we looked at the tensor in, the tensors out before the, be- before we call the prefill and after. Everything looks just fine. Second thing we did was running NVIDIA's Compute Sanitizer tool to really see if we have any, uh, out-of-bound memory, um, any other, uh, memory, uh,"
  },
  {
    "startMs": 467368,
    "text": "uh, bugs or issues. Looked okay to me. Um, then what we did was we tried to isolate between the decode kernels and the prefill kernels. Now, we saw that the prefill kernels were working just fine, so we tried to not call the decode kernels 'cause in Mamba you're able to do that. Um, so what we did was, um, we moved all of our calls and all of our computations to go through the p- through the prefill kernel, and there you have it. The gibberish all of a sudden kind of vanished. So we were like, \"Okay, it's gotta be the"
  },
  {
    "startMs": 497388,
    "text": "decode kernels.\" But you know how it is in software. You get excited too quickly, and then you figure out it's not what happened. So what we did was, uh, we tried to start playing with, uh, with vLLM's engine, and we kind of needed to go and, you know, lift the hood up and see what we can do to maybe get a bit better understanding and, uh, maybe, you know, get our hands dirty because vLLM didn't really give us more, um, tools to really, um, debug our kernel and our, and our forward pass. So once, so once,"
  },
  {
    "startMs": 527308,
    "text": "uh, once a tensor, once the request gets all the way to your forward pass and before it goes into your prefill and decode kernels, they don't really have any identity. You can't really tell what prompt, uh, is currently being processed. It, uh, it's all just tensors and, and numbers and matrices. So what we did was is we added to, uh, the request ID to some class called Forward Context that, uh, that we propagated all the way down to, uh, to Mamba's, uh, forward pass just before the, the prefill and the decode kernels were called."
  },
  {
    "startMs": 556988,
    "text": "And there we just managed to, uh, you know, have a simple c- uh, if condition with a request ID, the one that gave us gibberish, and put a breakpoint there. And then we were able to, to, to infer and to really, um, inspect all the metadata that comes along with it. And the second we did that, we saw that the request was, um, for the first time when it went through the, uh, through the forward pass, it's actually doing decode before prefill. The scheduler, um, decided that it's, um, that the, that"
  },
  {
    "startMs": 586868,
    "text": "this request should be doing, uh, decode before prefill. And as Yuval said earlier, in our, in the life cycle of a prompt, a prompt should first be, uh, going through prefill and then decode. And what happens was is that when, when in Mamba, um, you run a request fir- uh, uh, with a, uh, with decode first, after a long... a lot of other requests were already computed, the state was already kind of, um, overused, and we were"
  },
  {
    "startMs": 616728,
    "text": "using, uh, the data and the computations of stale requests, requests that came before it. So now we were actually running decode on, on previous requests, and that kind of generated gibberish for us. So the kernels weren't doing the wrong thing. They were called at the wrong time for the wrong requests. And now... And why did it matter only for Mamba? The reason is, was is that in, uh, in, in attention, uh, when you write the tokens KV, you write it, you, you write the tokens KVs before you actually, you"
  },
  {
    "startMs": 646648,
    "text": "read it. So even if you have stale data, it's being overwritten. But for Mamba, as, uh, as I said, it, when, uh, when, uh, when you first go through the decode kernels, you first read the state, and then you compute over it. So what happens was is you just use over, uh, you use stale data, um, um, when you do the, wh-when you do a decode. And the fix was relatively simple. We just needed to make sure that what we do is that when a request first... wh-when the req- when the scheduler first classifies a request, it's gotta, it's gotta make"
  },
  {
    "startMs": 676628,
    "text": "sure that, um, that, um, that when, that it sets, that, that if, that if, that if you see the request that's, whose tokens were never been computed, um, we, and, and they're, and it, and they're zero, to mark them as, to mark them as, uh, as prefill, um, as, as prefill so the, so when they get to the forward pass, they'll actually just be, um, uh, used for prefill and not decode and not, um, and not chunked. You can see that it was merged after some time. Um, and that really leads us to... And then we thought everything was fixed, right? We thought everything was fixed,"
  },
  {
    "startMs": 706608,
    "text": "and there you have it. No more issues. But that, but that was almost the case because after a little bit of time, it gets us to, uh, case number two, which, um, surfaced another issue that we've faced in our, in our, uh, RL, in our inference. Um, so we ran RL, and, and when, while we were running, um, our, our trainings, our post-trainings, um, and we looked at our evaluations and all of our benchmarks, we thought that we had some logprob spikes, uh, between the rollout and the FSDP step. Uh, so b- then... And that"
  },
  {
    "startMs": 736608,
    "text": "was before any weight up- update. So same weights, uh, same inputs, and the two logprob should be, uh, identical. Now they weren't Um, we saw that every 12 step cons- uh, constantly, um, there was a logprob spike, and that was kind of weird. Now, what would you guys do, right? What can-- what, what's the ne- what's the first thing to do here? So we wanted to find some lever that changes how things fail and not just how much they fail. We wanna see how much, um..."
  },
  {
    "startMs": 766648,
    "text": "We wanna tweak some knobs that, that don't just tell us, \"Hey, this error, um, this error is very, very ba- uh, uh, this error happens, uh, this many times or, or, or... and, and so, and so on.\" We wanted to tweak some knobs that kind of tell us that once we tweak that knob, um, we understand how, how it's wired to, uh, anything in the, in, in vLLM's engine, and so we'll be able to specifically go and, uh, and debug that specific part. Um, that's some, uh, some, some cool meme that"
  },
  {
    "startMs": 796488,
    "text": "Yuval wanted to put in."
  },
  {
    "startMs": 799708,
    "text": "Um, uh, so what we did was is we decided to, um, to increase the, the, the amount of rollouts per prompt. Uh, since we saw, uh, in our, in our default RL engine we have, um, eight rollouts per prompt, and we saw that it happened consi- uh, deterministically every 12 steps, we decided, \"Okay, let's try to tweak it up a bit and, uh, and increase the amount of rollouts per prompt.\" So we started doubling it from s- eight to 16 to 64, 32, and 128. And you can see here that it's almost, um, almost,"
  },
  {
    "startMs": 829508,
    "text": "um... it's-- there's, there's a pattern here. That, that the m- the more we increased it, the closer, um, it happened. Because what we wanted to achieve here, we wanted to try to reproduce the issue as fast as possible so we'd have a faster debug, uh, uh, debug loop, uh, feedback loop. So when we, when we re- when we ran it on 128 rollouts, uh, per prompt, it happened immediately on step one, and we didn't have to wait for step 12 and step 24 and so on. Now, you might think, \"Okay, so you guys played with the, uh, with the GPU memory utilization before. You, you tweaked it. You"
  },
  {
    "startMs": 859728,
    "text": "decreased it. It looks like, you know, when you test on pressure, um, it really surfaced things up.\" So, so we, we, we thought that as well. And when we reduced the GPU memory from zero, uh, from 0.9, uh, to 0.2, it actually caused the issue to go away. So we actually pulled the wrong lever here. Um, and the reason is, is because we noticed that Mamba kernels, um, used, um, int32, unsigned int, uh, 32 index pattern, uh, pointer. So once the offset went past"
  },
  {
    "startMs": 889808,
    "text": "some, you know, 4 billion, um, uh, in, in numbers, it wrapped around instead of throwing an error. Um, so, uh, when we shrank the GPU memory, vLLM allocated a small state buffer and, um, and the cache index never got large to hit that slot. So we were, so we were just not reaching, um, far enough for the buffer to trigger an overflow. So again, the fix was rather simple. All we needed to do was just change one word, one v- one data type variable, um, from"
  },
  {
    "startMs": 919668,
    "text": "uint32 to size_t, which basically means for most modern ar- architectures, hard- uh, hardware architectures, size_t would mean to, uh, it would be now changed to unsigned 64, uh, bit, and that's a very large number. We didn't-- We never reached that number, and that overflow now never happened. So  what we can see here is that we had, um, two scenes and, uh, one criminal. Um, both kind of, you know, they had similar symptoms. Both had silent"
  },
  {
    "startMs": 949608,
    "text": "gibberish and, and silent logprob spikes, which also kind of pro- sometimes generated gibberish. They were both around the Mamba state cache. Um, they were both surfaced by memory pressure, whether it was for worse or for the best, and, uh, both found via logprob's, um, forensics. Um, stateful inference, inference systems don't fail loudly. They lie to you confidently. I mean, obviously, sometimes you get crashed, you get out of bounds errors, you get other, you know, exceptions and so on,"
  },
  {
    "startMs": 980008,
    "text": "but sometimes there are some errors that don't surface up, and you don't get a trace log. You don't get anything. You have to go and dig and understand why things happen. Um, so if there are some takeaways to take from this, um, presentation is build a logprob comparison script. If you need to compare your quality, you need to compare it to understand whether you modelize the issues or not, um, logprob comparison script with a baseline of some other inference framework that you have or built is always great. Um, reproducing under"
  },
  {
    "startMs": 1010068,
    "text": "pressure and constrained memory, uh, crank the scale up, play with other knobs that the inference framework gives you, and really try to understand where the issue comes from. Uh, look for what moves, um, the failure shape, the timing, the space, and the location. And when things don't really have identity, uh, thread identity through. And what also I want you to take from this, and don't be afraid to even, you know, for complex, uh, systems like vLLM or any other, um, complex framework, don't be afraid"
  },
  {
    "startMs": 1039968,
    "text": "to go dig in the code, get your hands dirty. Um, sometimes, you know, model languages, um, LLMs are, um, they might tell you how things work, but, you know, without you seeing it in your own eyes and getting your hands dirty, you won't get full understanding of what's going on. Um, thank you. You guys can add us on LinkedIn. Scan the, uh, QR code to read the actual blog that we've, uh, published with this finding. Um, yeah, that's it."
  }
]