Speculative Decoding in Production: Accelerating LLM Inference via Custom Draft Models — code snippets

snippet-1.py

import torch

def verify_draft_tokens(
    target_probs: torch.Tensor,  # Shape: (K+1, vocab_size)
    draft_probs: torch.Tensor,   # Shape: (K, vocab_size)
    draft_tokens: torch.Tensor,  # Shape: (K,)
) -> tuple[torch.Tensor, int]:
    """
    Executes rejection sampling to verify K draft tokens against target probabilities.
    Returns a tensor of accepted tokens and the index of the first rejected token.
    """
    accepted = []
    K = draft_tokens.shape[0]
    
    for i in range(K):
        token_id = draft_tokens[i].item()
        p = target_probs[i, token_id].item()
        q = draft_probs[i, token_id].item()
        
        # Rejection sampling threshold
        if torch.rand(1).item() < min(1.0, p / (q + 1e-12)):
            accepted.append(token_id)
        else:
            # Token rejected: sample from adjusted distribution and terminate
            diff = torch.clamp(target_probs[i] - draft_probs[i], min=0.0)
            norm = diff.sum()
            adjusted_probs = diff / (norm + 1e-12) if norm > 0 else target_probs[i]
            next_token = torch.multinomial(adjusted_probs, num_samples=1).item()
            accepted.append(next_token)
            return torch.tensor(accepted, dtype=torch.long), i
            
    # If all K tokens are accepted, sample the (K+1)-th token from target_probs[K]
    next_token = torch.multinomial(target_probs[K], num_samples=1).item()
    accepted.append(next_token)
    return torch.tensor(accepted, dtype=torch.long), K

snippet-2.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, temperature: float = 1.5, alpha: float = 0.5):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha  # Weight for KL versus hard cross-entropy
        self.kl_loss = nn.KLDivLoss(reduction="batchmean")
        self.ce_loss = nn.CrossEntropyLoss()

    def forward(
        self, 
        student_logits: torch.Tensor,  # Shape: (batch, seq, vocab)
        teacher_logits: torch.Tensor,  # Shape: (batch, seq, vocab)
        labels: torch.Tensor           # Shape: (batch, seq)
    ) -> torch.Tensor:
        # Scale student and teacher logits by temperature
        soft_student = student_logits / self.temperature
        soft_teacher = teacher_logits / self.temperature
        
        # Calculate KL Divergence using log-probabilities for the student
        kl = self.kl_loss(
            F.log_softmax(soft_student, dim=-1),
            F.softmax(soft_teacher, dim=-1)
        ) * (self.temperature ** 2)
        
        # Hard cross-entropy loss against ground truth labels
        flat_student = student_logits.view(-1, student_logits.size(-1))
        flat_labels = labels.view(-1)
        ce = self.ce_loss(flat_student, flat_labels)
        
        return self.alpha * kl + (1.0 - self.alpha) * ce

snippet-3.py

from vllm import LLM, SamplingParams

def initialize_production_speculative_engine(
    target_model_path: str,
    draft_model_path: str,
    tensor_parallel_size: int = 4
) -> LLM:
    """
    Initializes a speculative vLLM engine with partitioned GPU memory allocation.
    """
    # Note: Target and draft share the GPU. Target is split via Tensor Parallelism.
    engine = LLM(
        model=target_model_path,
        speculative_model=draft_model_path,
        num_speculative_tokens=5,             # Speculative window K=5
        tensor_parallel_size=tensor_parallel_size,
        gpu_memory_utilization=0.85,           # Leave headroom for draft model KV cache
        max_model_len=4096,
        trust_remote_code=True,
        enforce_eager=False,                   # Enables CUDA graphs for fast execution
        speculative_draft_limits={"max_draft_tokens": 5}
    )
    return engine

snippet-4.py

import time
import asyncio
from vllm import LLM, SamplingParams

async def benchmark_inference_session(
    engine: LLM, 
    prompt: str, 
    temperature: float = 0.7
) -> dict:
    """
    Sends an inference query to the engine and metrics the latency characteristics.
    """
    sampling_params = SamplingParams(
        temperature=temperature,
        max_tokens=128,
        use_beam_search=False
    )
    
    start_time = time.perf_counter()
    # Execute generation asynchronously
    results_generator = engine.generate(prompt, sampling_params)
    
    num_tokens = 0
    first_token_time = None
    
    for request_output in results_generator:
        if num_tokens == 0:
            first_token_time = time.perf_counter()
        num_tokens = len(request_output.outputs[0].token_ids)
        
    end_time = time.perf_counter()
    
    total_latency = end_time - start_time
    ttft = first_token_time - start_time if first_token_time else 0.0
    itl = (total_latency - ttft) / (num_tokens - 1) if num_tokens > 1 else 0.0
    
    return {
        "ttft_ms": ttft * 1000,
        "itl_ms": itl * 1000,
        "tokens_generated": num_tokens,
        "throughput_tok_sec": num_tokens / total_latency
    }

snippet-5.py

from prometheus_client import Counter, Histogram

class SpeculativeObservabilityExporter:
    def __init__(self):
        # Track total tokens processed through speculation
        self.accepted_tokens = Counter(
            "speculative_decoding_accepted_tokens_total",
            "Total speculative tokens accepted by target model verification"
        )
        self.rejected_tokens = Counter(
            "speculative_decoding_rejected_tokens_total",
            "Total speculative tokens rejected by target model verification"
        )
        # Track acceptance rate per speculation loop
        self.step_acceptance_histogram = Histogram(
            "speculative_decoding_step_acceptance",
            "Distribution of accepted tokens per validation step",
            buckets=[0, 1, 2, 3, 4, 5]
        )
        # Track processing time for verification step
        self.verification_latency = Histogram(
            "speculative_decoding_verification_latency_seconds",
            "Time spent verifying speculative draft blocks",
            buckets=[0.005, 0.01, 0.015, 0.02, 0.03, 0.05]
        )

    def record_step(self, accepted_count: int, rejected_count: int, duration: float):
        self.accepted_tokens.inc(accepted_count)
        self.rejected_tokens.inc(rejected_count)
        self.step_acceptance_histogram.observe(accepted_count)
        self.verification_latency.observe(duration)