vLLM의 AMD GPU에서의 사전 추론: 성능과 방법론

요약

사전 추론은 vLLM에서 빠른 초안 구성 요소가 미래 토큰 여러 개를 제안하고, 목표 모델이 한 번의 전방 계산으로 이를 검증하는 방식으로 작동합니다. AMD Instinct MI300X/MI355X GPU에서는 일부 모델-방법 조합에서 최대 약 2.9배의 처리량 향상을 달성하면서도 원래 모델의 출력 동작을 유지합니다.

vLLM에서의 사전 추론 작동 방식

사전 추론은 원본(목표) LLM 앞에 초안 생성 및 검증 단계를 추가합니다. 초안 구성 요소는 후보 토큰 시퀀스를 생성하고, 목표 모델은 전체 후보 블록을 단일 전방 계산으로 평가합니다. 목표 모델이 수용한 토큰은 확정되며, 처음 거부된 토큰은 목표 모델의 자체 출력으로 대체됩니다. 생성이 끝날 때까지 이 과정이 반복됩니다.

핵심 특성

목표 모델이 진실의 기준이 유지되며, 검증 전까지 어떤 출력 토큰도 생성되지 않습니다.
한 번의 목표 모델 전방 계산으로 여러 개의 초안 토큰을 확정할 수 있어, 비용이 큰 목표 모델 계산 횟수를 줄일 수 있습니다.
수용 여부는 왼쪽에서 오른쪽으로 평가되며, 거부 시 후보 블록의 나머지 부분은 중단됩니다.

평가된 초안 생성 방법

vLLM은 다섯 가지 구체적인 초안 생성 접근 방식을 지원하며, 각각 초안 네트워크가 목표 모델로부터 정보를 어떻게 수신하고, 토큰을 순차적으로 생성할지 또는 병렬로 생성할지에 따라 다릅니다.

방법	초안 구성 요소	목표 모델 정보 사용	토큰 생성 방식
내장 MTP	모델 내부의 보조 예측 경로	목표 모델의 은닉 표현(또는 이전 MTP 단계)	순차적(반복적인 MTP 호출)
Gemma 4 MTP	목표 모델과 별도의 MTP 체크포인트	목표 활성화 + 공유 KV 캐시	순차적
EAGLE‑3	전용 순차적 사전 추론기	초기, 중간, 후기 목표 레이어의 융합 은닉 상태	순차적, 각 초안 토큰이 다음 토큰에 조건을 부여
DFlash	전용 병렬 사전 추론기	목표 은닉 상태를 모든 초안 레이어에서 사용하는 KV 쌍으로 투영	모든 위치를 한 번의 전방 계산에서 동시에 예측
DSpark	DFlash 백본 + 경량 마르코프 헤드	DFlash와 동일한 목표 컨텍스트	병렬 백본 + 순차적 토큰 선택 보정

내장 MTP

목표 모델 내부에 구현되며, 보조 헤드를 사용해 고정된 수의 미래 토큰을 예측합니다. 초안 토큰은 하나씩 차례로 생성되며, 각 단계는 이전 MTP 출력의 은닉 상태를 소비합니다. num_speculative_tokens이 내장 깊이를 초과하면 vLLM은 추가 전방 계산을 통해 MTP 경로를 재사용합니다.

Gemma 4 MTP

별도의 체크포인트(보조 모델)로 구현되며, 목표 모델과 KV 캐시와 활성화를 공유하여 이미 계산된 컨텍스트를 재사용할 수 있습니다. 초안 토큰은 여전히 순차적으로 생성되지만, 초안 구성 요소가 목표 텐서를 재사용하므로 메모리 오버헤드는 적습니다.

EAGLE‑3

EAGLE‑3는 세 개의 목표 레이어 은닉 상태의 융합 표현을 수신하는 사전 추론기를 훈련합니다. 사전 추론기는 순차적으로 작동합니다: 첫 번째 초안 토큰은 융합된 목표 특징과 샘플링된 토큰 임베딩을 결합하고, 이후 토큰은 이전 초안 출력에 조건을 부여합니다. 이로 인해 제안 길이에 비례하는 순차적 초안 작업이 발생합니다.

DFlash

DFlash는 전체 미래 위치 블록을 병렬로 예측합니다. 이미 목표에 의해 검증된 앵커 토큰이 블록을 시작하며, 나머지 마스크된 위치는 동시에 채워집니다. 목표 은닉 상태는 모든 초안 레이어가 접근할 수 있도록 추가 KV 쌍으로 투영되며, 병렬 계산 중에도 목표 컨텍스트를 유지합니다.

DSpark

DSpark는 DFlash의 백본에 경량 마르코프 헤드를 추가하여 병렬 백본 이후 토큰 간 의존성을 도입합니다. 백본은 모든 위치에 대한 기본 로짓을 생성하고, 마르코프 헤드는 이전에 선택된 초안 토큰을 사용하여 각 위치를 조정하여 병렬 속도와 순차적 일관성의 하이브리드를 달성합니다.

vLLM에서 사전 추론 활성화

사전 추론은 --speculative-config 플래그를 통해 구성됩니다. JSON 페이로드는 방법, 선택적 초안 체크포인트, num_speculative_tokens(제안 길이)를 지정합니다.

# 내장 MTP (추가 체크포인트 없음)
vllm serve <target-model> \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 4}'

# Gemma 4 MTP (보조 체크포인트 필요)
vllm serve <target-model> \
  --speculative-config '{"method": "mtp", "model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}'

# EAGLE‑3
vllm serve <target-model> \
  --speculative-config '{"method": "eagle3", "model": "RedHatAI/<target>-speculator.eagle3", "num_speculative_tokens": 3}'

# DFlash (병렬 블록)
vllm serve <target-model> \
  --speculative-config '{"method": "dflash", "model": "z-lab/<target>-DFlash", "num_speculative_tokens": 15}'

# DSpark (병렬 + 마르코프 보정)
vllm serve <target-model> \
  --speculative-config '{"method": "dspark", "model": "RedHatAI/<target>-speculator.dspark", "num_speculative_tokens": 7}'

메모리 고려사항 – 내장 MTP는 목표 모델과 가중치를 공유하므로 GPU 메모리 추가 부담이 거의 없습니다. 다른 모든 방법은 별도의 초안 체크포인트를 로드합니다. 필요 여유 공간은 체크포인트 크기, 정밀도(FP16/INT4), 텐서 병렬도, 런타임 버퍼에 따라 달라집니다.

벤치마크 방법론

저자들은 AMD Instinct MI300X 및 MI355X GPU에서 ROCm™ 스택을 사용하여 출력 토큰 처리량(토큰/초)을 측정했습니다. 벤치마크는 실제 수용 패턴을 반영하기 위해 작업 기반 데이터셋(GSM8K, MATH500, HumanEval, MBPP)을 사용했습니다. 각 목표-방법 조합에 대해 num_speculative_tokens(N)을 스윕하고 다음을 기록했습니다:

처리량 비율(사전 추론 / 베이스라인)
평균 수용 길이(MAL) – 검증 라운드당 평균 수용된 초안 토큰 수
수용률(AR) – 검증을 통과한 제안 토큰의 비율

모든 측정은 동일한 하드웨어 및 소프트웨어 스택에서 실행된 비사전 추론 순차적 베이스라인과의 상대적 비교입니다.

주요 성능 결과

관측된 최고 속도 향상

목표 모델	방법	데이터셋	최적 N	처리량 비율
`gemma‑4‑26B‑A4B‑it`	DFlash	MATH500	7	2.87×
`gemma‑4‑26B‑A4B‑it`	Gemma 4 MTP	MATH500	5	2.74×
`gemma‑4‑26B‑A4B‑it`	EAGLE‑3	GSM8K	5	2.27×
`gemma‑4‑31B‑it`	DFlash	MATH500	7	2.41×
`Qwen3‑8B`	DSpark	GSM8K	7	1.63×
`Qwen3‑6‑35B‑A3B`	DFlash	MATH500	7	2.06×
`Kimi‑K2.5`	DFlash	GSM8K	7	2.37×

방법별 추세

순차적 초안(내장 MTP, Gemma 4 MTP, EAGLE‑3)은 제안 길이가 적당한 수준(N ≈ 3–5)까지 처리량을 향상시킵니다. 그 이상이 되면 추가 순차적 초안 작업이 목표 모델 계산 절약보다 더 커져, 정체 또는 하락을 초래합니다.
병렬 초안(DFlash, DSpark)은 일반적으로 더 큰 N(N ≈ 7–11)에서 최고 비율을 달성합니다. 후속 위치의 수용률은 낮아지지만, 병렬 비용은 낮아서 전체 처리량은 여전히 증가합니다.
DSpark는 경량 마르코프 헤드를 추가하여, 동일한 N에서 순수 DFlash보다 몇 퍼센트 낮은 처리량을 보입니다. 이는 추가 순차적 보정 단계 때문입니다.
모델 계열의 중요성 – Qwen 3.5/3.6 모델에 내장된 MTP는 동일한 모델 크기에서 DFlash보다 성능이 뛰어납니다. 반면 Gemma 4에서는 병렬 방법이 우세합니다.
작업 종속성 – 코드 중심 데이터셋(HumanEval, MBPP)은 후속 토큰이 예측하기 어려워 짧은 제안을 선호하는 경향이 있으며, 수학 중심 데이터셋(GSM8K, MATH500)은 더 긴 제안을 수용할 수 있습니다.

수용 동작

모든 실험에서 첫 번째 초안 토큰은 90% 이상의 수용률을 보입니다. 수용률은 위치에 따라 단조롭게 감소하며, DFlash에서 N = 15일 경우 15번째 토큰의 수용률은 5% 미만일 수 있습니다. 평균 수용 길이(MAL)는 처리량과 강하게 상관되며, 높은 MAL → 적은 목표 모델 계산 → 높은 속도 향상입니다.

실용적인 튜닝 가이드

작게 시작하기 – 어떤 방법이든 num_speculative_tokens = 1로 시작하여 추가 비용 없이 정확성을 검증하세요.
N 스윕하기 – 대표적인 작업 하에서 처리량, MAL, AR을 측정하면서 N을 증가시켜 보세요(e.g., 1–3–5–7–11–15).
위치별 수용률 관찰하기 – 위치 k 이후 수용률이 급격히 떨어진다면, 낭비된 초안 작업을 피하기 위해 N을 k로 줄이세요.
메모리 예산 설정하기 – 초안 체크포인트를 위한 충분한 GPU 메모리 확보; 병렬 방법(DFlash/DSpark)은 일반적으로 가장 많은 메모리를 필요로 합니다.
하드웨어 특화 튜닝 – AMD Instinct GPU는 ROCm 최적화된 어텐션 백엔드(triton_attn)를 통해 최대 효과를 발휘합니다. 병렬 초안은 MI300X/MI355X의 높은 메모리 대역폭을 활용합니다.
배치 크기 및 토큰 제한 – 더 큰 max-num-batched-tokens과 긴 컨텍스트 창은 특히 병렬 방법에서 초안 비용의 분산을 개선합니다.

초안 체크포인트 확보 방법

공급자	지원 방법	예시 체크포인트
Google	Gemma 4 MTP	`google/gemma-4-26B-A4B-it-assistant`, `google/gemma-4-31B-it-assistant`
LightSeek	EAGLE‑3, EAGLE‑3.1	`lightseekorg/kimi-k2.5-eagle3-mla`
Red Hat AI	EAGLE‑3, DFlash, DSpark	`RedHatAI/gemma-4-26B-A4B-it-speculator.eagle3`, `RedHatAI/gemma-4-31B-it-speculator.dspark`
Z-Lab	DFlash	`z-lab/gemma-4-26B-A4B-it-DFlash`, `z-lab/Qwen3.8B-DFlash-b16`
DeepSeek AI	EAGLE‑3, DFlash, DSpark	`deepseek-ai/eagle3_qwen3_8b_ttt7`, `deepseek-ai/dflash_qwen3_8b_block7`
Inferact	EAGLE‑3, DSpark	`Inferact/MiniMax-M3-EAGLE3`, `Inferact/Kimi-K3-DSpark`

새로운 사전 추론기 훈련(고수준 워크플로우)

대상 작업(채팅, 코드, 수학 등)에 적합한 대표적인 프롬프트 수집.
목표 모델 응답 생성 – 추론 시 사용할 정확한 토크나이저, 채팅 템플릿, 샘플링 설정을 사용하여 응답 생성.
은닉 상태 추출 모드 선택 – 온라인(실시간), 오프라인(사전 저장), 하이브리드(캐시 후 첫 번째 에포크).
필요한 목표 레이어 추출 (예: EAGLE‑3의 초기, 중간, 후기 레이어; DFlash/DSpark의 모든 선택된 레이어).
사전 추론기 훈련 – 목표 모델과 동일한 어휘, 은닉 크기, 토큰 임베딩을 사용. 방법별 헤드(병렬 블록, 마르코프 보정 등) 포함.
검증 – 수용률, MAL, 엔드투엔드 처리량 검증; 수용률이 낮으면 프롬프트 혼합 또는 훈련 하이퍼파라미터 조정.
체크포인트 패키징 및 --speculative-config를 사용해 목표 모델과 함께 배포.

미래 방향성

학습되지 않은 사전 추론(예: n-그램 접미사 예측)은 매우 반복적인 코드 편집 작업에 학습된 사전 추론기를 보완할 수 있습니다.
더 넓은 동시성 연구 – 다중 사용자 배치 부하, 다양한 배치 크기, 샘플링 온도 변화 하에서 사전 추론 작동 측정.
사전 추론기 데이터 영향 – 프롬프트 다양성과 도메인 특화 훈련 데이터가 코드, 수학, 채팅, 다국어 작업에서 수용률에 미치는 영향에 대한 체계적 분석.
심층 프로파일링 – AMD GPU에서 초안 생성, 목표 검증, KV 캐시 재사용, 그래프 실행 오버헤드를 분리하여 커널 수준 최적화를 안내.

감사의 말

Hongxia Yang 및 Peng Sun(AMD)과 Pin Siang Tan, Jun Kang Chow, Ye Hur Cheong(Embedded LLM)에게 하드웨어 접근성과 협업에 감사드립니다.

면책 조항

측정은 두 가지 AMD Instinct 구성에서 수행되었습니다:

하드웨어 1: 8× MI300X GPU(gfx942), 듀얼 EPYC 9654 96코어 CPU.
하드웨어 2: 8× MI355X GPU(gfx950), 듀얼 EPYC 9575F 64코어 CPU(MiniMax-M3-MXFP8용). 소프트웨어 스택: Ubuntu 22.04.5 LTS, ROCm 7.2.53211, vLLM 0.23.1rc1.dev1120+g0f0f28b53, PyTorch 2.11.0, Transformers 5.13.1, Python 3.12.13. 결과는 다른 드라이버, 커널 버전, 하드웨어 구성에 따라 달라질 수 있습니다.