OmniDraft: Een Cross-vocabulaire, Online Adaptieve Drafter voor On-device Speculatieve Decodering

OmniDraft: A Cross-vocabulary, Online Adaptive Drafter for On-device Speculative Decoding

July 3, 2025

Auteurs: Ramchalam Kinattinkara Ramakrishnan, Zhaocong Yuan, Shaojie Zhuo, Chen Feng, Yicheng Lin, Chenzheng Su, Xiaopeng Zhang

cs.AI

Samenvatting

Speculatief decoderen vereist over het algemeen een klein, efficiënt ontwerp-model dat offline is voorgetraind of gedistilleerd voor een specifieke doelmodelreeks, zoals Llama- of Qwen-modellen. Binnen online implementatieomgevingen zijn er echter twee grote uitdagingen: 1) het gebruik van een doelmodel dat niet compatibel is met het ontwerp-model; 2) de verwachting van latentieverbeteringen tijdens gebruik en over tijd. In dit werk stellen we OmniDraft voor, een uniform raamwerk dat een enkel ontwerp-model in staat stelt te werken met elk doelmodel en zich dynamisch aan te passen aan gebruikersdata. We introduceren een online n-gram cache met hybride distillatie-finetuning om de kruiswoordenschat-mismatch tussen ontwerp- en doelmodellen aan te pakken; en verbeteren verder de decodering snelheid door adaptieve ontwerptechnieken te benutten. OmniDraft is bijzonder geschikt voor on-device LLM-toepassingen waar modelkosten, efficiëntie en gebruikersaanpassing de belangrijkste discussiepunten zijn. Dit benadrukt verder de noodzaak om de bovenstaande uitdagingen aan te pakken en motiveert het paradigma van "één ontwerper voor allen". We tonen de vaardigheid van het OmniDraft-raamwerk aan door online leren uit te voeren op taken voor wiskundig redeneren, coderen en tekstgeneratie. Opmerkelijk is dat OmniDraft een enkel Llama-68M-model in staat stelt te koppelen aan verschillende doelmodellen, waaronder Vicuna-7B, Qwen2-7B en Llama3-8B-modellen voor speculatief decoderen; en biedt daarnaast een snelheidsverbetering van tot 1,5-2x.

English

Speculative decoding generally dictates having a small, efficient draft model that is either pretrained or distilled offline to a particular target model series, for instance, Llama or Qwen models. However, within online deployment settings, there are two major challenges: 1) usage of a target model that is incompatible with the draft model; 2) expectation of latency improvements over usage and time. In this work, we propose OmniDraft, a unified framework that enables a single draft model to operate with any target model and adapt dynamically to user data. We introduce an online n-gram cache with hybrid distillation fine-tuning to address the cross-vocabulary mismatch across draft and target models; and further improve decoding speed by leveraging adaptive drafting techniques. OmniDraft is particularly suitable for on-device LLM applications where model cost, efficiency and user customization are the major points of contention. This further highlights the need to tackle the above challenges and motivates the ``one drafter for all'' paradigm. We showcase the proficiency of the OmniDraft framework by performing online learning on math reasoning, coding and text generation tasks. Notably, OmniDraft enables a single Llama-68M model to pair with various target models including Vicuna-7B, Qwen2-7B and Llama3-8B models for speculative decoding; and additionally provides up to 1.5-2x speedup.