Draft Less, Retrieve More: Speculative Decoding을 위한 하이브리드 트리 구성

동적 깊이 prune은 드래프트 비용을 낮추지만 MAT를 떨어뜨려 Pareto frontier를 굳건히 만든다. Graft는 pruning으로 해방된 예산을 retrieval으로 보충해 동일 예산에서 후보 범위를 확장하고, lossless한 검증 경로를 유지하며 짧은 컨텍스트와 긴 컨텍스트 양쪽에서 실용적 속도 향상을 달성한다. 대형 모델에서의 평균 속도 up to 5.41×를 기록하고, Qwen3-235B에서 EAGLE-3 대비 평균 속도 향상을 최대 21.8% 달성했다. 또한 LLaMA3.1-8B의 장-context에서 3.22×의 평균 속도를 달성하고 Qwen3-14B에서 EAGLE3-64K 대비 16.6%를 넘었다. DFlash와의 확장 가능성에 대한 초기 탐색도 제시한다.

왜 중요한가

핵심 기여

Budget-compensation 관점의 동적 트리 prune

동적 깊이 pruning이 도입하는 MAT 손실을 보완하기 위해 pruning이 해방한 슬롯을 retrieval으로 보충하는 관점으로, 고정 예산 내에서 후보 범위를 확장하고 end-to-end latency를 줄이는 원리를 제시한다.

Graft: GPU-친화적 prune-then-graft 구조

root-centered retrieval를 병렬로 실행해 pruning으로 해방된 슬롯에 retrieval을 삽입하고, merging 후 lossless하게 검증하는 하이브리드 트리 구성이다. GPU-resident adjacency matrix를 사용하고 online 업데이트를 통해 컨텍스트에 맞춘 후보를 재정렬한다.

실무 배포 설정에의 확장성

short-context, long-context, 고사양 대규모 모델에서 Pareto frontier를 새로 정의한다. Qwen3-235B에서 최대 21.8%의 평균 속도 향상을 기록하고, LLaMA3.1-8B에서 3.22×의 long-context 평균 속도향상을 달성했다.

핵심 아이디어 이해하기

단계별 요지 1) 시작점은 SD에서의 단일 혹은 다수의 드래프트 토큰 체인인 트리 구조이다. Dense 트리에 비해 Dynamic-depth pruning은 계산을 줄이지만, 구조적으로 Dense Tree의 후보 경로 전체를 활용하지 못하게 되고 MAT가 상한에 묶인다. 2) 해결 원리는 prune으로 남겨진 슬롯을 retrieval으로 채워 budget를 재분배하는 것이다. GPU 상의 adjacency matrix를 통해 root 토큰에서 retrieval이 병렬로 시작되며, pruning checkpoint에서 released budget에 맞춰 stage-specific retrieval 템플릿이 호출된다. 3) 결과적으로 pruning으로 얻은 예산은 retrieval에 의해 보완되고, 최종 검증 경로의 예산은 변경 없이 유지된다. 이로써 단일 fixed-budget에서 더 넓은 후보 공간과 높은 MAT를 모두 달성할 수 있다.

방법론

첫째, Dynamic-depth pruning으로 budget를 해방한다. 깊이 d에서 c_d를 계산해 prune 여부를 결정하고, 각 pruning 스테이지 s마다 예산 분할을 결정한다(Kmax = K_draft_s + K_ret_s). 둘째, GPU-resident adjacency matrix M를 이용해 stage-specific retrieval 템플릿 G_ret_s를 root에서부터 매칭된 prefix로 채운다. 이때 retrieval은 루트 주변의 후보를 선점해 전체 예산을 확장하며, g_s는 prune 여부를 gating한다. 셋째, retained draft nodes와 retrieval 노드를 하나의 hybrid 트리로 합친 뒤, 표준 target 모델의 verification 경로에서 검증한다. 네째, 온라인 업데이트를 통해 verified tree의 logits를 이용해 M를 지속적으로 갱신한다. 다섯째, long-context에서 프롬프트 evidence를 확장하고 YaRN-adapted EAGLE-64K를 도입해 retrieval의 초기 신뢰도를 높인다.

주요 결과

주요 벤치마크에서 Graft는 Short-context에서 1.83×~5.41×의 속도향상을 기록했다. Vicuna-13B에서 평균 속도 향상은 5.41×에 달했고, large-scale Qwen3-235B에서 EAGLE-3 대비 평균 향상은 최대 21.8%였다. LLaMA3-1B-8B의 long-context에서 평균 속도향상은 3.22×이며, Qwen3-14B에서 EAGLE3-64K 대비 16.6%를 넘었다. Long-context에서의 평균 속도향상은 3.22×로, Qwen3-235B에서도 2.09×의 평균 속도향상을 달성했다. Ablation 연구에서 w/o retrieval은 LLaMA3-8B에서 평균 3.74×, Qwen3-8B에서 2.20×로 하락했고, w/o prune 시에는 3.62×, 1.95×로 추가 하락했다. Warm-up five 라운드 도입 시 HumanEval/GSM8K/MT-Bench에서 각각 2.46×/2.86×/2.46×의 속도향상을 보였다. DFLASH(16) 대비 Graft-DFLASH은 Qwen3-8B에서 평균 3.71×로 상승했고, DFLASH 대비 9.1%의 상대 향상을 보였다. 고병렬(batch) 서빙에서도 Graft는 EAGLE3 대비 throughput 및 MAT를 개선했다(예: HumanEval에서 319.54 tokens/s vs 303.97 tokens/s; MAT 3.22 vs 3.10).

기술 상세

Graft의 아키텍처는 (i) budget-releasing dynamic-depth pruning, (ii) GPU-resident retrieval adjacency matrix, (iii) stage-adaptive retrieval templates, (iv) hybrid verification 및 온라인 업데이트의 네 축으로 구성된다. Pruning은 depth별 confidence c_d를 기반으로 𝜏_d 임계치를 사용해 d0, d1, d5에서 후보를 제거한다. 예산 분해는 Kmax = K_draft_s + K_ret_s로 표기되며, s가 작아질수록 더 많은 budget이 retrieval에 할당된다. Retrieval 템플릿은 base-tree Envelope와 호환되도록 stage별 prefix를 미리 준비하고, pruning stage에 맞춘 prefix를 M에서 행렬 조회로 채운다. Retreival과 retained draft를 합쳐 T_s = T_draft_s ∪ G_ret_s로 구성하고, |T_s| = Kmax를 유지한다. 검증은 기존의 tree-attention 경로를 따라 하나의 target forward pass에서 이루어지며, retrieval 노드 역시 proposal로 취급되며 최종 분포는 lossless하게 보장된다. online 업데이트는 verified tree에서 얻은 logits를 이용해 M를 지속적으로 업데이트하며, long-context에서는 YaRN-adapted EAGLE-64K를 도입해 retrieval의 초기 신뢰성을 높인다.

실무 활용

Graft는 prune으로 비용을 절감한 후 retrieval로 채워 고정-budget에서 더 넓은 후보를 확보하고 검증 경로를 유지하는 방식으로, 실무에서도 대형 모델의 실시간 서비스에서 속도와 품질의 균형을 개선한다.

Long-context 코드 자동완성 및 대규모 문서 요약에서 속도와 품질의 균형 개선
대규모 다중모달/멀티턴 대화 시스템에서 speculative decoding의 throughput 향상
실시간 자동응답 서비스에서 fixed-budget 하에 latency를 낮추고 MAT를 높이는 고려사항의 예시

코드 공개 여부: 미확인

키워드

speculative decodingdraft-then-verifyGraftpruningretrievalGPU-resident adjacency matrixlong-contextshort-context

용어 해설

예산 해제(Budget Release): — 프래닝에서 제거된 후보 슬롯을 재활용하는 아이디어로, pruning으로 줄인 후보 영역을 retrieval으로 채워 주어 fixed-budget에서 속도와 MAT 사이의 트레이드를 완화시키는 원리이다.
루트 중심 검색(Root-Centered Retrieval): — 루트를 기준으로 후보를 우선 탐색하는 retrieval 방식으로, pruning으로 남긴 슬롯의 근접 위치에서 고품질 후보를 채택하도록 설계된 검색 전략이다.
GPU-내 인접 행렬(GPU-Resident Adjacency Matrix): — 토큰 간 전이 후보를 top-k 형태로 GPU에 상주시키는 데이터 구조로, root-centered retrieval를 빠르게 수행하고 CPU-최근접 조회를 제거한다.
손실 없는 추측적 디코딩(Lossless Speculative Decoding): — 추측 제안이 최종 검증에서 손실 없이 모델 출력 분포를 재현하도록 보장하는 원리로, 제안 토큰의 채택 여부가 타깃 모델의 검증 규칙에 의해 결정된다.

Draft Less, Retrieve More: Speculative Decoding을 위한 하이브리드 트리 구성

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

RTX 3090에서 Qwen3.8-27B 추론 가속

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

RTX 3090에서 Qwen3.8-27B 추론 가속