vLLM의 Mamba 버그 디버깅

AI21 Labs는 Jamba Reasoning 3B를 RL 파이프라인에 올려두고, 가끔 모델이 완전한 gibberish를 내뱉는 문제를 발견했다. 평소의 저품질 출력과 달리 max logprobs는 오히려 높게 나와, 단순한 모델 성능 저하가 아니라 런타임 상태 오염을 의심하게 만들었다.

같은 체크포인트를 Hugging Face transformers로 돌리면 정상이었고, 문제는 vLLM에서만 간헐적으로 나타났다. 그래서 저자들은 vLLM 출력과 transformers 기준값을 같은 토큰 시퀀스에 대해 비교하는 디버그 스크립트를 만들고, 실제 생성 토큰의 logprob 차이를 추적했다.

처음에는 재현이 잘 되지 않았지만, RL 환경의 메모리 압박을 재현하려고 gpu_memory_utilization을 0.2로 낮추자 문제가 드러났다. 이후에는 854번째 요청에서 gibberish가 재현됐고, temperature=0에서도 동일하게 반복돼 버그가 결정적임이 확인됐다.

의심은 Mamba SSM state를 다루는 CUDA prefill kernel로 향했다. 새 요청이 들어올 때 상태를 제대로 초기화하지 못했거나, 이전 요청의 cache slot을 잘못 읽는 상황이면 이후 토큰 전체가 망가질 수 있었기 때문이다. 하지만 바운드 체크, 포인터 계산 검증, compute-sanitizer 검사, 28개 Mamba layer의 SSM 상태 검사까지 모두 이상이 없었다.

다음으로는 vLLM의 prefill/decode split logic를 의심했다. 특히 Mamba에서는 request가 prefill과 decode 중 어디로 배정되느냐가 상태 초기화 방식과 직결되기 때문이다. 그 과정에서 오래된 V0 engine과의 차이가 중요한 단서가 됐는데, V0에서는 버그 때문에 Mamba1이 사실상 decode kernel을 타지 않고 prefill 경로만 사용했다.

V1에서도 모든 요청을 강제로 prefill로만 보내자 gibberish가 사라졌다. 즉, 문제는 Mamba의 상태 자체가 아니라 decode 경로와 scheduler의 상호작용에 있었다는 뜻이며, 이후의 디버깅은 decode kernel과 그 주변 로직으로 더 좁혀졌다.