P-EAGLE: vLLM에서 병렬 추측 디코딩

TL;DR

P-EAGLE는 vLLM(v0.16.0부터) 에 통합된 병렬 추측 디코딩 방법으로, 자동회귀 초안 작성의 순차적 병목 현상을 없앱니다. 모든 K개의 초안 토큰을 단일 전방 패스로 생성함으로써, NVIDIA B200 GPU를 사용한 실제 워크로드에서 기존 EAGLE-3 대비 최대 1.69배의 속도 향상을 달성합니다.

자동회귀 초안 작성 병목 해결

EAGLE와 같은 표준 추측 디코딩 방법은 토큰을 자동회귀적으로 초안 작성하는데, K개의 초안 토큰을 생성하려면 초안 모델을 K번 순차적으로 전방 패스해야 합니다. 이는 추측 깊이에 비례해 지연 시간이 선형적으로 증가하게 하여, 초안 작성 오버헤드가 전체 성능 향상을 잠식하지 않도록 시스템이 얼마나 공격적으로 추측할 수 있는지를 제한합니다.

P-EAGLE는 초안 작성 과정을 자동회귀에서 병렬로 전환함으로써 이 한계를 제거합니다. 순차적인 단계 대신, P-EAGLE는 단일 전방 패스에서 전체 초안 토큰 집합을 생성하여 초안 토큰 수와 필요한 전방 패스 수를 분리합니다.

P-EAGLE 아키텍처 및 메커니즘

P-EAGLE는 초안 토큰을 생성하기 위해 두 가지 주요 단계로 동작합니다:

단계 1: 프리필링

대상 모델은 프롬프트를 처리하고 새로운 토큰을 생성합니다. 이 과정에서 P-EAGLE는 내부 은닉 상태를 캡처합니다: 각 프롬프트 위치에 대한 h_prompt와 새로 생성된 토큰에 대한 h_context. 이러한 은닉 상태는 초안 작성자의 예측을 안내합니다.

단계 2: 병렬 초안 작성

초안 작성자는 토큰 임베딩과 은닉 상태의 조합을 사용하여 각 위치에 대한 입력을 병렬로 구성합니다:

  • Prompt Positions: 각 프롬프트 토큰 임베딩 emb(p)은 대상 모델에서 나온 해당 h_prompt와 짝을 이루며, 하나의 위치를 이동시켜 위치 $i$의 토큰을 예측할 수 있게 합니다.
  • Next-Token-Prediction (NTP): 첫 번째 위치는 새로 생성된 토큰 임베딩 emb(new)와 h_context를 짝지습니다.
  • Multi-Token-Prediction (MTP): 위치 2부터 K까지는 아직 토큰 임베딩과 은닉 상태가 존재하지 않으므로, P-EAGLE는 두 개의 학습 가능한 파라미터를 사용합니다: 공유 마스크 토큰 임베딩 emb(mask)와 공유 은닉 상태 h_shared. 이들은 중립적인 자리표시자 역할을 합니다.

그 후 모든 위치는 N개의 트랜스포머 레이어와 언어 모델 헤드를 통과하여 초안 토큰 $t_1$부터 $t_K$까지를 동시에 예측합니다.

긴 시퀀스에 대한 학습

병렬 초안 작성은 학습 중 메모리 요구량을 증가시킵니다. 길이 N인 시퀀스에서 K개의 병렬 그룹을 학습하면 $N \times K$개의 전체 위치가 생성되기 때문입니다. 예를 들어 $N=8,192$이고 $K=8$인 경우, 단일 학습 예시는 65,536개의 위치를 포함하며, 이는 40억 개가 넘는 요소를 가진 거대한 어텐션 행렬을 초래합니다.

이를 해결하기 위해 P-EAGLE는 시퀀스 분할 알고리즘을 도입하여 시퀀스 내부 분할을 수행합니다. 이 알고리즘은 $N \times K$ 위치 시퀀스를 연속적인 청크로 나누면서 경계 간 올바른 어텐션 종속성을 유지하고 동일 시퀀스의 청크들 간에 그래디언트를 누적합니다.

vLLM 구현 세부 사항

vLLM에 병렬 초안 작성을 통합하려면 배치 메타데이터와 메모리 관리와 관련된 여러 기술적 과제를 해결해야 했습니다:

Fused Triton Kernel

병렬 초안 작성은 추가된 MASK 자리표시자 때문에 초안 및 검증 배치 형태 간 일관성을 깨뜨립니다. 여러 GPU 연산을 통해 배치 메타데이터를 재구성하는 오버헤드를 피하기 위해, vLLM은 fused Triton kernel을 구현합니다. 이 커널은 이전 토큰 ID와 위치를 복사하고, 보너스 토큰을 삽입하며, 병렬 초안 슬롯을 MASK 토큰 ID로 채우고, 필요한 메타데이터(거부된 토큰 마스크, 마스크된 토큰 마스크, 은닉 상태 매핑)를 단일 패스로 생성합니다.

은닉 상태 관리

은닉 상태는 토큰 ID보다 훨씬 크기 때문에, vLLM은 전용 복사 커널을 사용하여 학습된 은닉 상태 자리표시자(parallel_drafting_hidden_state_tensor)를 마스크 토큰 슬롯에 브로드캐스트하고, 대상 은닉 상태는 새로운 위치에 매핑합니다.

KV 캐시 및 CUDA 그래프

  • KV Cache: 거부된 토큰은 PADDING_SLOT_ID (-1)로 매핑되어 잘못된 캐시 쓰기를 방지합니다.
  • CUDA Graphs: 캡처 범위가 $K \times \text{max_num_seqs}$만큼 확장되어 더 큰 초안 배치 크기를 수용합니다.

성능 벤치마크

NVIDIA B200 GPU를 사용한 GPT-OSS-20B 평가 결과, P-EAGLE는 MT-Bench, HumanEval, SPEED-Bench 등 여러 벤치마크에서 기존 EAGLE-3보다 크게 우수함을 보여줍니다.

처리량 향상

낮은 동시성(c=1)에서는 P-EAGLE가 55–69% 높은 처리량을 제공하고, 높은 동시성(c=64)에서는 5–25%의 향상이 지속됩니다. 구체적으로, EAGLE-3 대비 속도 향상 비율은 다음과 같습니다:

  • SPEED-Bench: 최대 1.69배 (c=1에서)
  • HumanEval: 최대 1.55배 (c=1에서)
  • MT-Bench: 최대 1.55배 (c=1에서)

수용 길이 (AL)

P-EAGLE는 라운드당 평균 수용 초안 토큰 수(AL)에서 EAGLE-3보다 높습니다. 추측 깊이 $K=7$일 때, HumanEval에서 P-EAGLE의 AL은 30% 더 높으며(3.94 vs 3.03), SPEED-Bench에서는 31% 더 높습니다(3.38 vs 2.59).

추측 깊이 효율성

자동회귀 초안 작성자는 $K=3$에서 처리량이 정점에 도달하지만, P-EAGLE는 모든 동시성 수준에서 $K=7$에서 일관되게 최고 처리량을 달성합니다. 이는 P-EAGLE가 순차 초안 작성과 관련된 선형 지연 페널티 없이 더 깊은 추측으로부터 이점을 얻을 수 있음을 확인합니다.

배포

vLLM에서 병렬 초안 작성을 활성화하려면 SpeculativeConfig 클래스에 "parallel_drafting": true를 설정하면 됩니다. 사전 학습된 P-EAGLE 헤드는 다음 모델에 대해 HuggingFace에서 제공됩니다:

  • GPT-OSS 120B
  • GPT-OSS 20B
  • Qwen3-Coder 30B

Sources