vLLM 버그로 인한 '알 수 없는' 출력 발생

이러한 유형의 조용한 실패는 근본 원인을 식별하는 것을 매우 어렵게 만듭니다. 즉각적인 시스템 경고가 없으면 잘못된 결과가 장기간 감지되지 않아 다운스트림 애플리케이션과 사용자 신뢰에 영향을 미칠 수 있습니다.

예를 들어, Jamba는 트랜스포머와 Mamba(SSM) 아키텍처의 요소를 결합하여 복잡한 상호 작용을 생성합니다. 팀은 이러한 복잡한 문제를 해결하기 위해 순수한 모델 최적화에서 심층적인 엔지니어링 수준 디버깅으로 초점을 전환했습니다.

'임포스터 요청' 버그는 약 10,000건의 요청 중 한 번꼴로 발생하는 간헐적인 '알 수 없는' 출력으로 나타났습니다. 이러한 희귀성 때문에 테스트 환경에서 일관되게 재현하기 어려웠습니다. 이 문제는 프레임워크에 특화되어 vLLM에서만 나타났으며, 워크로드 규모에 매우 민감하여 트리거하려면 상당한 계산 스트레스가 필요했습니다. 이는 기본 모델 자체는 건전했기 때문에 전통적인 모델 품질 문제는 아니었습니다.

GPU 메모리 사용량을 90%에서 20%로 줄여 시스템이 처리하는 요청 밀도를 극적으로 높였습니다. 이 방법은 더 빠른 피드백 루프를 생성하여 잠재적인 수정 사항에 대한 더 빠른 반복을 가능하게 했습니다.

vLLM의 고도로 수정된 사용자 지정 커널과 달리 Transformers는 Mamba 커널의 바닐라 구현을 제공했습니다. 이 비교 분석은 예상 모델 추론 동작과 vLLM 엔진 내 실제 실행 간의 불일치를 식별하는 데 도움이 되었습니다.

vLLM에서 생성된 로짓과 Hugging Face의 순방향 통과(forward pass)를 비교하여 로짓 발산을 추적하는 사용자 지정 진단 도구가 개발되었습니다. 이 도구는 발산의 기준선을 설정하기 위해 특히 프리필(prefill) 단계에 중점을 두었습니다. 'compute_log_props' 함수는 소프트맥스 차이를 정확하게 계산하기 위해 생성되어 엔지니어가 두 구현 간에 토큰 분포가 어디에서 벗어났는지 정확히 찾아낼 수 있도록 했습니다.

엔지니어들은 프리필(prefill) 호출 전후의 텐서 입력 및 출력을 검사했으며, 모든 것이 정상으로 보였습니다. Nvidia compute sanitizer도 커널 내의 메모리 경계 초과 문제를 확인하는 데 사용되었지만, 즉각적인 문제는 감지되지 않았습니다.

모든 계산이 프리필(prefill) 커널을 통해서만 라우팅되었을 때, 무의미한 출력은 사라졌습니다. 이 발견은 디코드(decode) 커널이 '임포스터 요청' 버그의 주요 원인임을 강력히 시사하며, 조사의 초점을 전환시켰습니다.

기존 vLLM 구조는 요청이 순방향 통과에 도달하면 텐서 뒤에 원시 프롬프트 데이터를 숨겨 그 식별을 추적하기 어렵게 만들었습니다. 이를 해결하기 위해 엔지니어들은 요청 ID 메타데이터를 새로운 '전방 컨텍스트' 클래스에 주입했습니다. 이 컨텍스트는 순방향 통과로 전파되어 중단점(breakpoint) 사용 및 프리필(prefill) 및 디코드(decode) 커널을 통한 개별 요청의 보다 정밀한 추적을 가능하게 했습니다.

어텐션 메커니즘과 달리 Mamba 커널은 계산 전에 상태를 읽는데, 이는 디코드가 먼저 실행되면 이전 계산의 오래된 데이터를 사용하게 된다는 의미입니다. 이 부적절한 시퀀스는 새로운 토큰이 잘못되고 오래된 컨텍스트 정보로 처리되었기 때문에 '알 수 없는' 출력으로 이어졌습니다.

엔지니어들은 스케줄러가 요청을 처음 처리할 때, 아직 계산되지 않은 토큰을 정확하게 식별하도록 했습니다. 이러한 미계산 토큰은 프리필(prefill) 작업이 필요하다고 명시적으로 표시되어야 합니다. 이 조정은 스케줄러가 새로운 요청을 디코드(decode) 커널로 잘못 라우팅하는 것을 방지하여 모든 토큰이 적절한 초기 처리를 받도록 보장했습니다. 이 논리적 간극 수정의 성공적인 구현은 '임포스터 요청' 버그를 해결하는 통합된 수정으로 이어졌습니다.