EAGLE-3 추측 디코딩 on AMD Instinct GPUs
EAGLE-3 추측 디코딩 on AMD Instinct GPUs
vLLM과 AMD Quark 팀은 AMD Instinct GPU에서 EAGLE-3 추측 디코딩을 위한 엔드투엔드 파이프라인을 도입하여, 대형 혼합 전문가(MoE) 및 주의 집중 모델에 대한 무손실 추론 가속을 가능하게 합니다. 이 통합은 AMD Instinct MI355X GPU에서 Kimi-K2.5에 대해 1.69x에서 2.00x 범위의 처리량 속도 향상을 달성하고, MiniMax-M2.5에 대해 1.38x에서 1.79x의 향상을 달성합니다.
추측 디코딩과 EAGLE-3 접근법
추측 디코딩은 비용이 많이 드는 대상 모델 디코드 반복 횟수를 줄이는 무손실 가속 기술입니다. 토큰을 하나씩 생성하는 대신, 가벼운 드래프트 모델이 여러 미래 토큰을 제안하고, 대상 모델이 이를 한 번의 순방향 패스로 검증합니다. 대상 모델의 분포와 일치하는 토큰은 수락되며, 생성은 첫 번째 거부된 토큰에서 계속됩니다.
EAGLE-3은 대상 모델의 다층 특징을 활용하여 이전 버전을 개선합니다. 관련 없는 작은 언어 모델을 사용하는 대신, EAGLE-3은 대상 모델과 긴밀히 연결된 드래프트 모듈을 훈련시켜 저-, 중-, 고-level 의미적 특징을 활용하여 드래프트의 수락률과 전체 추론 속도를 높입니다.
AMD Quark MXFP4 양자화
대형 MoE 모델의 메모리 대역폭 및 용량 제약을 해결하기 위해, AMD Quark 팀은 MXFP4(OCP 마이크로스케일링 4비트 부동소수점) 형식을 사용하여 주류 LLM에 대한 'day-0' 양자화를 제공합니다.
MXFP4 구현의 주요 기술 세부 사항은 다음과 같습니다:
- 하드웨어 가속: AMD Instinct MI350-series GPU(MI350X/MI355X)는 네이티브 FP4 행렬 가속을 제공합니다.
- 메모리 효율성: MXFP4는 공유 스케일 팩터를 가진 작은 블록에 4비트 요소를 그룹화하여 메모리 footprint를 줄이면서 INT4보다 우수한 수치적 동작을 유지합니다.
- vLLM 통합: 이러한 체크포인트는 FP4 ASM GEMM 경로(
VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1) 및 AITER MoE 커널을 통해 ROCm에서 vLLM에 의해 소비됩니다.
vLLM 중심 EAGLE-3 훈련 파이프라인
높은 수락률을 가진 드래프트 모델 훈련은 시스템 문제로 간주되며, vLLM이 5단계에 걸친 훈련 루프의 핵심 역할을 합니다.
- 온-폴리시 데이터 합성: vLLM은 정확한 서빙 채팅 템플릿을 사용하여 훈련과 배포 간의 일관성을 보장하기 위해 AMD Quark MXFP4 대상 모델을 통해 응답을 생성합니다.
- 숨은 상태 추출: vLLM의 숨은 상태 추출 훅은 실행 중인 대상 엔진에서 보조 레이어(저-, 중-, 고-level 숨은 상태 및
fc_norm)을 직접 노출합니다. 이는 온라인, 오프라인, 스트리밍 모드를 지원하며, 후자는 단일 노드에서 420B MXFP4 MoE 대상 모델의 훈련을 가능하게 합니다. - 콜드스타트 FSDP2 훈련: 단일 레이어 EAGLE-3 드래프트 헤드는 FSDP2 하에서 훈련-시간-테스트(TTT) 손실과 위치-감쇠 가중치를 사용하여 처음부터 훈련됩니다.
- 루프 내 서브-평가: 현재 체크포인트는 주기적으로 내보내지고 vLLM 추측 디코딩 하에서 제공되어 실제 수용 길이를 측정하며, 선택된 체크포인트가 생산 엔진에 최적화되도록 보장합니다.
- 내보내기 및 배포: 최종 드래프트는 Hugging Face 형식으로 내보내지고 vLLM-ROCm을 통해 배포됩니다.
드래프트 품질 및 컨텍스트 스케일링
SPEED-Bench 평가에서 MiniMax-M3 EAGLE-3 드래프트는 11개 도메인에서 평균 수용 길이(AL)가 2.77입니다. 가장 높은 수용률은 구조화된 기술 콘텐츠에서 발견됩니다:
- 코딩: 3.16 AL
- 수학: 3.12 AL
- RAG: 3.11 AL
- 다국어: 3.07 AL
주목할 점은 프롬프트 길이가 1K에서 32K 토큰으로 증가함에 따라(2.64에서 2.63) 수용 길이가 안정적으로 유지된다는 점이며, 이는 장시간 컨텍스트에서 속도 향상이 감소하지 않음을 나타냅니다.
AMD Instinct MI355X에서의 성능 벤치마크
AMD Instinct MI355X (TP=4)에서 1K 입력 및 1K 출력 토큰(1K/1K)으로 수행된 벤치마크는 비추측 기준선에 비해 상당한 처리량 향상을 보여줍니다:
Kimi-K2.5 결과
- BF16 드래프트 경로: 처리량 증가 1.69x에서 1.90x
- AMD Quark FP8 드래프트 경로: 처리량 증가 1.76x에서 2.00x
MiniMax-M2.5 결과
- BF16 드래프트 경로: 처리량 증가 1.38x에서 1.79x
두 모델 모두에서 성능 향상은 낮은 동시성 수준에서 가장 두드러집니다.