P‑EAGLE: 대규모 컨텍스트를 위한 병렬 초안 생성 모델

초록

P‑EAGLE은 기존 EAGLE 초안 모델을 자동회귀에서 병렬 다중 토큰 예측으로 전환하고, 사전 계산된 어텐션 마스크와 시퀀스 분할 기법을 통해 20K 토큰까지의 장기 컨텍스트 학습을 가능하게 하여 GPT‑OSS 120B, 20B 및 Qwen3‑Coder 30B 모델에서 1.10‑1.36배의 추론 속도 향상을 달성한다.

상세 분석

본 논문은 추론 단계에서 토큰당 전체 모델을 한 번씩 통과시켜야 하는 전통적인 자동회귀 디코딩의 비효율성을 해소하고자, 경량 초안 모델이 여러 토큰을 동시에 생성하는 병렬 초안 기법을 제안한다. 기존 병렬 초안 연구(ParallelSpec, P‑ARD)는 긴 시퀀스에서 어텐션 메모리와 마스크 생성 비용이 O((n·K)²) 로 급증해 실용적인 학습이 불가능했다. P‑EAGLE은 두 가지 핵심 혁신을 통해 이 한계를 극복한다. 첫째, “공유 숨은 상태”(shared hidden state) h 를 도입해 MTP(Multi‑Token Prediction) 위치에서 이전 단계의 숨은 벡터가 없을 때 이를 대체한다. 이 설계는 위치별 별도 파라미터를 두지 않아 파라미터 효율성을 높이고, 어텐션 자체가 위치 정보를 충분히 전달한다는 이론적 분석을 통해 정당화된다. 실험적으로도 위치‑특정 설계보다 7‑15% 높은 Acceptance Length 를 기록한다. 둘째, 학습 효율성을 위해 어텐션 마스크를 최대 길이(예: 20K) 기준으로 한 번만 사전 계산하고, 각 배치에서는 단순 슬라이스 연산으로 필요한 부분을 추출한다. 이는 마스크 생성 비용을 상수 시간으로 낮추어 데이터 로딩 오버헤드를 48배 감소시킨다. 또한, COD(Conditional Drop‑token) 샘플링에 의해 깊이별 토큰 수가 감소하는 구조를 유지하면서, 시퀀스를 S개의 세그먼트로 나누어 각 세그먼트별로 순전파‑역전파를 수행하고 그래디언트를 누적한다. 이 “시퀀스 내부 그래디언트 누적” 기법은 깊이‑간 의존성을 보존하도록 세그먼트 할당을 반복적으로 조정하는 알고리즘을 제시한다. 결과적으로 메모리 사용량은 O(L²/S²) 로 감소해 20K 토큰·K=8 설정에서도 GPU 메모리 한계를 넘지 않는다. 아키텍처 측면에서는 LLaMA‑3 스타일의 RoPE를 사용하고, 토큰 임베딩을 고정하지 않고 학습시켜 마스크 토큰 임베딩이 의미 있는 입력으로 작동하도록 한다. 깊이 4의 드래프터가 1‑layer 대비 46% 높은 Acceptance Length 를 달성했으며, 훈련‑추론 깊이 정렬이 성능에 크게 기여함을 확인한다. 최종적으로 vLLM에 통합된 P‑EAGLE은 GPT‑OSS 120B, 20B, Qwen3‑Coder 30B에 대해 기존 자동회귀 EAGLE‑3 대비 1.10×‑1.36×의 추론 속도 향상을 보이며, 장기 추론 시 25% 수준의 수용률 감소 문제를 완화한다.