JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting

Speculative decoding (SD) accelerates autoregressive Large Language Models (LLMs) by drafting multiple tokens and verifying them in parallel, but it faces a scaling limitation: increasing the draft budget improves speed only when acceptance remains high and drafting overhead stays low. This ceiling has been difficult to break because prior head-based SD methods face a causality-efficiency dilemma. Autoregressive drafters produce path-conditioned candidates that are effective for tree speculative decoding with higher acceptance length, but their drafting cost grows with tree depth. Bidirectional block-diffusion drafters generate all positions in one pass, but their branch-agnostic marginals can form individually plausible yet mutually inconsistent trees, wasting budget and reducing acceptance. We propose JetSpec, a head-based SD framework that combines one-forward drafting efficiency with branch-wise causal conditioning. JetSpec trains a causal parallel draft head over fused hidden states from the frozen target model, producing candidate trees whose scores align with the target model's autoregressive factorization. This enables JetSpec to convert larger draft budgets into longer accepted prefixes and higher end-to-end speedup. Across math, coding, and chat benchmarks on dense and MoE Qwen3 models, JetSpec consistently outperforms bidirectional-head and tree-based SD baselines. On H100 GPUs, JetSpec achieves up to 9.64x speedup on MATH-500 and 4.58x on open-ended conversational workloads, with further latency gains demonstrated through vLLM integration under realistic serving loads. Our code and models are available at https://github.com/hao-ai-lab/JetSpec.

한 줄 요약

JetSpec은 병렬 트리 드래프팅을 통해 투기 디코딩의 확장 한계를 돌파한 헤드 기반 프레임워크이다.

핵심 기여도

JetSpec은 투기 디코딩에서 병렬 드래프팅 효율성과 분기별 인과 조건을 결합한 첫 번째 헤드 기반 프레임워크이다.
Qwen3 모델에서 MATH-500 데이터셋에서 최대 9.64×, 대화형 작업에서 4.58×의 속도 향상을 달성했다.
vLLM 통합을 통해 실제 서빙 환경에서도 지연 시간 감소를 보였다.
투기 디코딩에서 수용 길이가 길어지면서 더 높은 전체 속도 향상을 달성할 수 있음을 입증했다.

핵심 아이디어

JetSpec은 기존 투기 디코딩에서 발생하는 인과성-효율성 딜레마를 해결하기 위해, 병렬 드래프팅의 효율성과 트리 기반 투기 디코딩의 인과 조건을 결합하는 새로운 헤드 구조를 제안한다. 기존의 자동회귀 드래프터는 트리 깊이에 따라 드래프팅 비용이 증가하는 문제를, 양방향 블록 확산 드래프터는 분기 무시한 확률 분포로 인해 불일치 트리를 생성하는 문제를 각각 제기한다. JetSpec은 동결된 타겟 모델의 히든 상태를 기반으로 인과 병렬 드래프트 헤드를 학습하여, 자동회귀 인수분해와 일치하는 후보 트리를 생성한다.

기술적 접근법

JetSpec은 병렬 드래프팅 효율성과 인과 조건을 결합한 헤드 기반 투기 디코딩 프레임워크이다.
투기 디코딩 과정에서 동결된 타겟 모델의 히든 상태를 활용하여 인과 병렬 드래프트 헤드를 학습한다.
후보 트리의 점수는 타겟 모델의 자동회귀 인수분해와 정렬되도록 설계된다.
Qwen3 모델(밀집형 및 MoE)에서 수학, 코딩, 대화 벤치마크를 사용해 평가했다.

주요 결과

MATH-500 데이터셋에서 JetSpec은 H100 GPU에서 최대 9.64×의 속도 향상을 달성했다.
대화형 작업에서는 4.58×의 속도 향상을 보였다.
vLLM 통합을 통해 실제 서빙 환경에서도 지연 시간 감소를 입증했다.
양방향 헤드 및 트리 기반 투기 디코딩 기반 대비 일관된 성능 개선을 보였다.

의의 및 한계

JetSpec은 투기 디코딩에서 확장 한계를 돌파한 첫 번째 헤드 기반 프레임워크로, 대규모 언어 모델의 추론 속도를 현저히 향상시킬 수 있다. 특히, 인과성과 효율성을 동시에 담보하면서 수용 길이를 늘려 전체 속도 향상 폭을 확대한 점이 학술적·실용적 가치를 높인다. 그러나 JetSpec이 모든 모델 아키텍처나 작업 유형에서 동일한 성능을 보장하는지는 추가 실험을 통해 검증이 필요하다.

실용적 활용

JetSpec은 대규모 언어 모델 기반의 실시간 추론이 필요한 산업, 예를 들어 고객 지원 챗봇, 코드 생성 도구, 수학 문제 풀이 시스템 등에 적용 가능하다. 특히, 고성능 GPU 환경에서 서빙 효율성을 극대화할 수 있다.