Accelerating Production LLMs with Combined Token/Embedding Speculators

Davis Wertheimer, Joshua Rosenkranz, Thomas Parnell, Sahil Suneja, Pavithra Ranganathan +2 more
4/29/2024
cs.CL

No official code found. Best implementation: foundation-model-stack/fms-fsdp (286 stars, Apache-2.0, PyTorch, last commit 2025-11-24).

Abstract

This technical report describes the design and training of novel speculative decoding draft models, for accelerating the inference speeds of large language models in a production environment. By conditioning draft predictions on both context vectors and sampled tokens, we can train our speculators to efficiently predict high-quality n-grams, which the base model then accepts or rejects. This allows us to effectively predict multiple tokens per inference forward pass, accelerating wall-clock inference speeds of highly optimized base model implementations by a factor of 2-3x. We explore these initial results and describe next steps for further improvements.

View on arXivView PDF

Code Implementations(1)

PyTorchApache-2.0

πŸš€ Efficiently (pre)training foundation models with native PyTorch features, including FSDP for training and SDPA implementation of Flash attention v2.

28646Shell, PythonFeb 5, 202411 months agoApache-2.0
distributed-trainingllmpytorch

Similar Papers

Cite this paper

@article{wertheimer2024accelerating,
  title  = {Accelerating Production LLMs with Combined Token/Embedding Speculators},
  author = {Davis Wertheimer and Joshua Rosenkranz and Thomas Parnell and Sahil Suneja and Pavithra Ranganathan and Raghu Ganti and Mudhakar Srivatsa},
  year   = {2024},
  eprint = {2404.19124},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2404.19124v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2404.19124.svg)](https://codeofpaper.com/paper/2404.19124)

Discussion