Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PRISM: Parametrically Refactoring Inference for Speculative Sampling Draft Models

Xuliang Wang, Yuetao Chen, Maochan Zhen, Fang Liu, Xinzhou Zheng, Xingwu Liu, Hong Xu, Ming Li

TL;DR

PRISM addresses the trade-off between draft quality and computational overhead in speculative decoding by distributing draft-step computation across specialized parameter sets, decoupling capacity from per-step cost. The approach enlarges draft capacity without increasing per-step compute, using a step-wise surjection to assign draft steps to processing modules and training with context-aligned losses, implemented in an optimized engine. The paper provides extensive system- and data-scale experiments showing 2.4x+ end-to-end speedups and better data efficiency, along with scaling-law validation. This work offers a practical, architecture-level principle for scalable speculative decoding and system-level deployment.

Abstract

Large Language Models (LLMs), constrained by their auto-regressive nature, suffer from slow decoding. Speculative decoding methods have emerged as a promising solution to accelerate LLM decoding, attracting attention from both systems and AI research communities. Recently, the pursuit of better draft quality has driven a trend toward parametrically larger draft models, which inevitably introduces substantial computational overhead. While existing work attempts to balance the trade-off between prediction accuracy and compute latency, we address this fundamental dilemma through architectural innovation. We propose PRISM, which disaggregates the computation of each predictive step across different parameter sets, refactoring the computational pathways of draft models to successfully decouple model capacity from inference cost. Through extensive experiments, we demonstrate that PRISM outperforms all existing draft architectures, achieving exceptional acceptance lengths while maintaining minimal draft latency for superior end-to-end speedup. We also re-examine scaling laws with PRISM, revealing that PRISM scales more effectively with expanding data volumes than other draft architectures. Through rigorous and fair comparison, we show that PRISM boosts the decoding throughput of an already highly optimized inference engine by more than 2.6x.

PRISM: Parametrically Refactoring Inference for Speculative Sampling Draft Models

TL;DR

PRISM addresses the trade-off between draft quality and computational overhead in speculative decoding by distributing draft-step computation across specialized parameter sets, decoupling capacity from per-step cost. The approach enlarges draft capacity without increasing per-step compute, using a step-wise surjection to assign draft steps to processing modules and training with context-aligned losses, implemented in an optimized engine. The paper provides extensive system- and data-scale experiments showing 2.4x+ end-to-end speedups and better data efficiency, along with scaling-law validation. This work offers a practical, architecture-level principle for scalable speculative decoding and system-level deployment.

Abstract

Large Language Models (LLMs), constrained by their auto-regressive nature, suffer from slow decoding. Speculative decoding methods have emerged as a promising solution to accelerate LLM decoding, attracting attention from both systems and AI research communities. Recently, the pursuit of better draft quality has driven a trend toward parametrically larger draft models, which inevitably introduces substantial computational overhead. While existing work attempts to balance the trade-off between prediction accuracy and compute latency, we address this fundamental dilemma through architectural innovation. We propose PRISM, which disaggregates the computation of each predictive step across different parameter sets, refactoring the computational pathways of draft models to successfully decouple model capacity from inference cost. Through extensive experiments, we demonstrate that PRISM outperforms all existing draft architectures, achieving exceptional acceptance lengths while maintaining minimal draft latency for superior end-to-end speedup. We also re-examine scaling laws with PRISM, revealing that PRISM scales more effectively with expanding data volumes than other draft architectures. Through rigorous and fair comparison, we show that PRISM boosts the decoding throughput of an already highly optimized inference engine by more than 2.6x.
Paper Structure (12 sections, 11 equations, 8 figures, 5 tables)

This paper contains 12 sections, 11 equations, 8 figures, 5 tables.

Table of Contents

  1. Introduction
  2. background
  3. methodology
  4. Motivation
  5. Prism Architecture
  6. Integrating with SGLang
  7. Experiments
  8. Effectiveness
  9. Re-examine Scaling Laws of Prism
  10. Hyper-Parameters
  11. Related Works
  12. conclusion

Figures (8)

  • Figure 1: Emperical results showing the average acceptance rate step-wise on LLaMA-3-8B.
  • Figure 2: An illustration of computation and gradient flow of HASS, EAGLE-3 and Prism.
  • Figure 3: An illustration of the Prism architecture and its computing paths. (a) The target model prefills or verifies multiple tokens and the output hidden states of its last-layer transformer is used as an input for the draft model, fused with the input token embedding. The fused features is fed to the first transformer of the Prism drafter for prefill. (b) After fusing the embedding of the previous generated token and the previous generated hidden states, apply the fused feature to the second transformer layer for the first decoding step. Note that KV caches from transfomer 1 is transfered to transformer 2.
  • Figure 4: Scaling curves of Prism and comparative drafters over 6 benchmarks and greedy and non-greedy sampling.
  • Figure 5: Prism vs. EAGLE-3
  • ...and 3 more figures