P-EAGLE: vLLM에서 병렬 투기적 디코딩을 통한 LLM 추론 가속화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

대형 언어 모델의 추론 속도를 높이기 위한 투기적 디코딩 기법 중 하나인 EAGLE은 드래프트 토큰을 순차적으로 생성해야 하는 구조적 한계가 있었다. P-EAGLE(Parallel-EAGLE)은 학습 가능한 마스크 토큰과 공유 은닉 상태를 활용하여 K개의 드래프트 토큰을 단일 포워드 패스에서 동시에 생성함으로써 이 병목을 제거했다. NVIDIA B200 GPU 환경에서 테스트한 결과, 기존 EAGLE-3 대비 최대 1.69배의 처리량 향상을 달성했으며 수락률 또한 개선되었다. 이 기술은 vLLM v0.16.0 이상 버전에 통합되어 간단한 설정 변경만으로 즉시 적용 가능하다.

배경

vLLM v0.16.0 이상 설치, NVIDIA GPU (B200 등 Blackwell 아키텍처 권장), 투기적 디코딩 및 EAGLE 알고리즘에 대한 기본 이해

대상 독자

LLM 서비스의 추론 성능 최적화와 비용 절감을 목표로 하는 MLOps 엔지니어 및 백엔드 개발자

의미 / 영향

P-EAGLE의 등장은 투기적 디코딩의 효율성을 한 단계 높여 고성능 GPU 자원을 더 효율적으로 사용하게 하며, 실시간 응답이 중요한 코드 에이전트나 대화형 AI 서비스의 사용자 경험을 크게 개선할 것이다.

섹션별 상세

기존 EAGLE 방식은 K개의 드래프트 토큰을 생성하기 위해 K번의 순차적인 포워드 패스가 필요하여 드래프트 모델의 오버헤드가 성능 향상을 저해하는 병목 현상이 발생했다.

bash

vllm serve openai/gpt-oss-20b \
  --speculative-config '{"method": "eagle3", "model": "amazon/gpt-oss-20b-p-eagle", "num_speculative_tokens": 5, "parallel_drafting": true}'

vLLM에서 P-EAGLE을 사용하여 병렬 투기적 디코딩을 활성화하는 실행 명령 예시

NVIDIA B200 카드 1개에서 동시성 1일 때 SPEED-BENCH 성능 비교 차트 — ChartP-EAGLE이 858 TPS를 기록하여 Baseline(456 TPS) 및 기존 EAGLE 3(507 TPS) 대비 압도적인 성능 우위를 점하고 있음을 보여준다. 이는 병렬 드래프팅이 단일 요청 처리 속도를 획기적으로 높일 수 있음을 증명한다.

P-EAGLE은 드래프트 모델의 입력을 구성할 때 아직 생성되지 않은 미래 위치에 대해 학습 가능한 파라미터인 공유 마스크 임베딩과 은닉 상태를 배치하여 병렬 생성을 가능하게 했다.

P-EAGLE 아키텍처 개요 다이어그램 — Diagram타겟 모델에서 추출한 은닉 상태와 새로 생성된 토큰 임베딩을 결합하고, 미래 위치를 위한 학습 가능한 마스크 파라미터를 사용하여 단일 패스로 여러 토큰(t1~t4)을 예측하는 과정을 시각화한다. Step 2의 병렬 입력 구성 방식이 핵심 차별점이다.

긴 문맥 학습 시 발생하는 메모리 문제를 해결하기 위해 시퀀스 분할 알고리즘을 도입하여 8K 이상의 긴 시퀀스에서도 효율적으로 병렬 드래프트 모델을 학습시켰다.

vLLM 통합 과정에서 데이터 레이아웃 불일치 문제를 해결하기 위해 GPU 상에서 직접 배치 메타데이터를 재구성하는 융합된 Triton 커널을 구현하여 오버헤드를 최소화했다.

python

# Copy target hidden states to their new positions
self.hidden_states[out_hidden_state_mapping] = target_hidden_states

# Fill masked positions with the learned Parallel Drafting hidden state
mask = self.is_masked_token_mask[:total_num_output_tokens]
torch.where(
    mask.unsqueeze(1),
    self.parallel_drafting_hidden_state_tensor,
    self.hidden_states[:total_num_output_tokens],
    out=self.hidden_states[:total_num_output_tokens],
)

타겟 모델의 은닉 상태를 복사하고 마스크된 위치를 학습된 파라미터로 채우는 내부 로직

벤치마크 결과 P-EAGLE은 MT-Bench, HumanEval 등에서 EAGLE-3보다 높은 토큰 수락률을 보였으며, 특히 낮은 동시성 환경에서 55-69%의 높은 처리량 향상을 기록했다.

동시성 수준에 따른 MT-Bench 처리량 비교 그래프 — Chart동시성이 1에서 64로 증가함에 따라 P-EAGLE이 모든 구간에서 EAGLE-3보다 높은 TPS를 유지함을 보여준다. 특히 낮은 동시성(c=1)에서 1.55배의 가장 큰 가속 효과가 나타난다.

실무 Takeaway

vLLM 환경에서 추론 성능을 극대화하려면 speculative-config에 parallel_drafting: true 옵션을 추가하여 P-EAGLE을 활성화해야 한다.
P-EAGLE은 단일 패스로 다수 토큰을 생성하므로 기존보다 더 깊은 추론 깊이(K=7 이상)를 설정해도 성능 저하 없이 높은 효율을 얻을 수 있다.
코드 생성과 같은 긴 시퀀스 작업에서 P-EAGLE은 기존 방식보다 약 30% 더 높은 토큰 수락률을 제공하여 실질적인 응답 지연 시간을 단축시킨다.

언급된 리소스

논문P-EAGLE ArXiv Paper

DemoHuggingFace P-EAGLE Models

GitHubvLLM Integration PR