Speculative Decoding & Medusa Architecture: Multi-Token Parallel LLM Acceleration

1. The Memory Bandwidth Wall in Autoregressive LLM Inference

In standard autoregressive Large Language Model (LLM) generation, generating each new token requires loading the entire model weights (e.g. 140GB for a 70B FP16 model) from High Bandwidth Memory (HBM) into GPU SRAM. Because arithmetic intensity during single-batch decoding is extremely low ($pprox 1 ext{ FLOP/byte}$), GPU Tensor Cores sit idle > 90% of the time waiting for memory transfers.

Speculative Decoding breaks this bottleneck by shifting the generation paradigm from sequential token-by-token execution to parallel batch verification. A small, fast draft model generates $ candidate tokens rapidly, which the large target model verifies simultaneously in a single compute-bound forward pass.

// Mathematical Speedup Formulation in Speculative Decoding
// Let alpha = Acceptance rate per draft token (0.7 - 0.85)
// Let K = Draft lookahead length (typically 4 - 6 tokens)
// Expected accepted tokens per step E = (1 - alpha^(K+1)) / (1 - alpha)
// Wall-clock speedup factor S = E / (1 + (T_draft * K / T_target))

2. Speculative Sampling Mathematics & Distribution Preservation

To ensure that speculative decoding outputs are mathematically identical to the target model native probability distribution (x)$, Leviathan et al. introduced Speculative Sampling.

For each candidate token $ generated by draft distribution (x_i)$, the target model accepts $ with probability $\min\left(1, rac{p(x_i)}{q(x_i)} ight)$. If a token is rejected at index $, the target model resamples the replacement token from the normalized positive residual distribution $\max(0, p(x) - q(x))$, guaranteeing zero loss in output quality or perplexity.

// Python: Speculative Sampling Rejection and Resampling Algorithm
import numpy as np

def speculative_sample_step(target_logits, draft_logits, draft_token_id):
    p = np.exp(target_logits) / np.sum(np.exp(target_logits))
    q = np.exp(draft_logits) / np.sum(np.exp(draft_logits))
    
    # Accept / Reject Probability
    p_val = p[draft_token_id]
    q_val = q[draft_token_id]
    accept_prob = min(1.0, p_val / q_val)
    
    if np.random.rand() < accept_prob:
        return True, draft_token_id
    else:
        # Resample from residual distribution (p - q)+
        residual = np.maximum(0.0, p - q)
        residual_prob = residual / np.sum(residual)
        resampled_id = np.random.choice(len(residual_prob), p=residual_prob)
        return False, resampled_id

3. Medusa Architecture: Multi-Head Self-Speculation

Deploying two separate models (draft + target) introduces operational complexity, VRAM fragmentation, and vocabulary synchronization overhead. Medusa solves this by augmenting the original model with multiple single-layer feed-forward heads (Medusa Heads) trained on top of the last hidden states.

Head 1 predicts token +1$, Head 2 predicts +2$, and Head 3 predicts +3$ concurrently from the same feature representation, generating a multi-token speculative candidate tree with zero additional base model invocations.

// PyTorch: Medusa Multi-Head Extension Architecture
import torch
import torch.nn as nn

class MedusaHead(nn.Module):
    def __init__(self, hidden_size, vocab_size, num_heads=4):
        super().__init__()
        self.heads = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_size, hidden_size),
                nn.SiLU(),
                nn.Linear(hidden_size, vocab_size, bias=False)
            )
            for _ in range(num_heads)
        ])
        
    def forward(self, last_hidden_states):
        # Returns list of logits for predictions at t+1, t+2, t+3, t+4
        return [head(last_hidden_states) for head in self.heads]

4. Tree-Attention Kernels & High-Throughput Serving

Instead of validating a single linear token candidate chain, Medusa constructs a candidate Tree representing top-$ paths across all speculative heads. To process all tree nodes inside a single target attention forward pass without cross-branch contamination, the engine applies Tree Attention with a custom causal 2D Boolean mask.

Integrated into high-throughput inference engines like vLLM and TensorRT-LLM, Tree-Attention speculative decoding achieves 2.2x - 2.8x speedup in real-world chat and coding benchmarks.

# Launching vLLM with Speculative Decoding Acceleration
vllm serve meta-llama/Llama-3.1-70B-Instruct \n  --speculative-model meta-llama/Llama-3.1-8B-Instruct \n  --num-speculative-tokens 5 \n  --gpu-memory-utilization 0.92 \n  --max-model-len 8192

Written by Zyekh Abdul Qadir Jailani

Digital Forensics & Incident Response (DFIR) Specialist & Security Researcher specializing in Linux kernel hardening, threat hunting, and system security research.

LinkedIn • GitHub • Discord • Security Policy • PGP Key