FlashMap 공개, LLM 디코딩 최대 6.14배 빨라져

Polymathic AI, 프린스턴대, USC ISI, 뉴욕대, Flatiron Institute 연구진이 추측 디코딩용 드래프터 FlashMap을 공개했다. 토큰 16개 블록을 순전파 한 번에 만들어 내며, 연구진은 Qwen3-8B에서 GSM8K 기준 최대 6.14배 속도 향상을 얻었다고 보고했다. 출력은 대상 모델이 직접 만들 결과와 같다.

Qwen3-8B용과 Qwen3-4B용 드래프터 가중치는 Hugging Face에 올라와 있고, vLLM과 SGLang 플러그인도 제공된다. 가중치 라이선스는 Apache-2.0, 코드는 MIT다. 설치는 uv pip install flashmap으로 하며 torch와 transformers 4.56 이상이 필요하다.

작동 방식

추측 디코딩에서는 작은 드래프터가 토큰 K개를 제안하고, 대상 모델이 한 번의 병렬 계산으로 K개를 검증해 일치하는 가장 긴 앞부분만 채택한다. 속도는 제안한 토큰이 얼마나 채택되는지와 제안 비용이 얼마나 싼지에 달려 있다. EAGLE-3은 자기회귀 방식으로 초안을 만들어 블록이 길수록 비용이 늘고, DFlash는 블록 확산으로 블록 전체를 한 번에 만든다.

DFlash는 블록 안 각 칸을 독립적으로 예측한다. 연구진은 이를 분해 근사라고 부른다. 이어질 말이 "new york"과 "san diego"가 반반이면 독립 예측기는 절반 확률로 "new diego"나 "san york"를 내놓고, 대상 모델은 두 번째 칸에서 이를 거부한다. FlashMap은 Flow Map Language Models를 바탕으로 노이즈 샘플 하나가 블록 전체로 대응되게 해, 노이즈가 모든 칸의 모드를 함께 정한다.

드래프터는 6블록 확산 트랜스포머(DiT)로, 임베딩 차원 2048에 어텐션 헤드 16개다. 대상 모델의 1, 9, 17, 25, 33번째 층 은닉 상태를 읽고, 임베딩 행렬과 LM 헤드는 대상 모델에서 복사해 고정한다. 크기는 Qwen3-8B용 5억, Qwen3-4B용 4억 7,500만 매개변수다.

학습은 두 단계다. 먼저 원핫 토큰 블록 위의 연속 흐름을 교차 엔트로피로 학습시켜 교사 모델을 만들고, 어휘 15만 1,936개에 맞춰 시간 재매개변수화를 적용했다. 다음으로 반군 증류를 통해 이 흐름을 노이즈에서 데이터로 한 번에 건너뛰는 2시점 디노이저, 곧 흐름 지도 학생 모델로 바꿨다. 두 드래프터 모두 Open-PerfectBlend의 샘플 50만 개로 학습했고, 모든 모델과 제거 실험의 연산량은 H200 기준 약 4,600 GPU 시간이었다.

벤치마크 결과

연구진은 EAGLE-3과 DFlash를 같은 데이터와 설정으로 다시 학습해 비교했다. 수학, 코드, 대화 9개 벤치마크를 온도 0과 1에서 측정했고, 논문의 속도는 NVIDIA H200 NVL 1장에서 동시 요청 1건, BF16로 얻었다.

• Qwen3-8B에서는 1단계 FlashMap이 9개 벤치마크 모두, 두 온도 모두에서 가장 빠른 드래프터였다. 다시 학습한 DFlash보다는 평균 약 7% 빨랐다.
• 최고 기록은 온도 0의 GSM8K에서 6.14배(채택 길이 7.69토큰)이며, 최저는 온도 1의 Alpaca에서 2.28배다.
• 채택 길이는 18개 설정 중 15개에서 DFlash를 앞섰고, 매개변수는 DFlash의 절반 수준이다. 온도 0에서 채택 길이가 DFlash보다 늘어난 폭은 수학 5.5%, 대화 11.1%로, 대화에서 가장 컸다.
• 2단계 모드는 18개 설정 모두에서 DFlash보다 채택 길이가 길었고 평균 10% 높았다. 블록 전반에서 채택 길이도 일정하게 유지했다. 4개 노이즈 샘플을 트리로 제안하면 속도가 평균 11% 더 올랐다.
• Qwen3-4B(온도 0)에서는 9개 과제 중 4개에서 DFlash에 뒤졌다.
• 온도 1에서 EAGLE-3은 1단계 FlashMap보다 채택 길이가 9% 길었지만, 실행 속도는 35% 느렸다.

토큰 사이 실제 구조를 쓰는지 확인하는 셔플 실험에서, 2단계 FlashMap은 16단계 교사 모델이 가진 결합 구조 이득의 대부분을 회수했다. 1단계에서는 그 일부만 남았다.

표의 라이선스 칸은 기준이 다르다. FlashMap과 DFlash는 가중치 라이선스이고, EAGLE-3의 Apache-2.0은 코드에 해당하므로 세 항목을 같은 기준으로 비교할 수 없다. 평균 속도는 Marktechpost가 논문 그림 1(c)에서 계산한 값이다. 각 프로젝트가 직접 내세우는 최고 수치는 FlashMap 6.1배, DFlash 6.17배(Qwen3-8B), EAGLE-3 6.5배로 서로 조건이 다르다.

배포 방법

flashmap 패키지의 spec_generate를 호출하면 드래프터와 대상 모델이 GPU 한 장에서 함께 돈다. vLLM 0.31에서는 --speculative-config에 드래프터를 넘기고 num_speculative_tokens를 16으로 둔다. SGLang 0.5.20에서는 --speculative-algorithm에 DSPARK를 쓰고 FlashMap 초안 경로를 붙인다.

num_steps=2로 두면 속도를 일부 내주고 채택률을 높인다. 온도 0을 넘겨 샘플링하려면 vLLM에서 "draft_sample_method": "probabilistic"을 설정해야 한다.