Sign In

Batch Speculative Decoding Done Right

Created by
  • Haebom
Category
Empty

저자

Ranran Haoran Zhang, Soumik Dey, Ashirbad Mishra, Hansi Wu, Binbin Li, Rui Zhang

개요

추측 디코딩은 작은 초안 모델을 사용하여 대상 모델이 병렬로 검증하는 여러 토큰을 제안하여 LLM 추론 속도를 높입니다. 이 아이디어를 배치로 확장하는 것은 프로덕션 제공에 필수적이지만, 이는 ragged tensor 문제를 야기합니다. 즉, 동일한 배치 내의 시퀀스가 서로 다른 수의 초안 토큰을 허용하여 오른쪽 정렬을 깨고 위치 ID, 어텐션 마스크 및 KV-캐시 상태를 손상시킵니다. 본 논문에서는 여러 기존 배치 구현이 출력 등가성, 즉 추측 디코딩이 표준 자동 회귀 생성과 동일한 토큰 시퀀스를 생성해야 하는 기본적인 요구 사항을 위반한다는 것을 보여줍니다. 이러한 위반은 정확히 ragged tensor 문제를 제대로 처리하지 못하여 발생합니다. 이에 대응하여 (1) 정확성을 보장하는 동기화 요구 사항을 특징짓고, (2) 재정렬이 오버헤드의 40%를 차지하는 EQSPEC을 제시하고, (3) 시퀀스의 슬라이딩 풀을 유지하고 동일한 길이의 그룹을 동적으로 형성하여 재정렬 오버헤드를 줄이면서 시퀀스별 추측 속도 향상을 유지하는 EXSPEC을 소개합니다. SpecBench 데이터 세트에서 Vicuna-7B/68M, Qwen3-8B/0.6B 및 GLM-4-9B/0.6B target/draft 쌍에 걸쳐 본 연구는 배치 크기 1에 비해 배치 크기 8에서 최대 3배의 처리량 향상을 달성하며, 95%의 출력 등가성을 유지하면서 배치 크기 8까지 효율적으로 확장됩니다.

시사점, 한계점

시사점:
추측 디코딩의 배치 처리를 위한 정확한 동기화 요구 사항을 제시.
EQSPEC과 EXSPEC 두 가지 새로운 배치 추측 디코딩 알고리즘 제안.
다양한 모델 쌍에서 최대 3배의 처리량 향상 달성 (배치 크기 8 기준).
95%의 출력 등가성 유지.
맞춤형 커널 없이 기존 추론 스택과 통합 가능.
한계점:
재정렬 오버헤드 (EQSPEC).
슬라이딩 풀 사용 (EXSPEC).
성능은 모델 및 데이터 세트에 따라 다를 수 있음.
github 코드 링크 제공.
👍