Adaptive Speculative Decoding for Efficient Large Language Model Inference

Problem Statement

Standard autoregressive decoding in Large Language Models requires a costly sequential forward pass for every single generated token. As models scale, memory bandwidth bottlenecks generation speed, causing high latency.

This project implements Speculative Decoding to parallelize token verification, dramatically increasing Tokens-Per-Second (TPS) without compromising output perplexity or quality.

By applying an adaptive loop with a GPT-2 target model and a DistilGPT-2 draft model, this implementation achieves near 2x latency improvements while strictly retaining 100% output equivalence.

Key Features

• Confidence-based Early Rejection: Halts the drafting process immediately if the draft model's confidence logic falls below a threshold (e.g., 60%), preventing the target model from wasting FLOPS on verifying likely-incorrect tokens.
• Dynamic K-Scaling: Automatically adjusts the draft length (k) based on the acceptance rate. Aggressive speculation scales actively with the contextual alignment between the draft and target models.
• Memory-Isolated Wrappers: Clean abstraction layers for Hugging Face Transformers.
• Interactive UI: A real-time Streamlit dashboard that visualizes accepted vs. rejected tokens.
• Asynchronous API: FastAPI-powered inference endpoints for production-ready integration.

Architecture & Core Algorithm

graph TD;
    subgraph Speculative Decoding Loop
        A[Start Generation] --> B[Draft Model<br/>Generates k tokens];
        B --> C[Target Model<br/>Verifies prefix + draft tokens];
        C --> D{Tokens Match?};
        D -- Yes --> E[Accept Tokens];
        E --> F{100% Accepted?};
        F -- Yes --> G[Increase k];
        F -- No --> H[Adjust K based on rate];
        G --> I[Next Iteration];
        H --> I;
        D -- No --> J[Reject at First Mismatch];
        J --> K[Discard subsequent drafts];
        K --> H;
    end

Loading

The system relies on exact-match verification:

1. The tiny, fast Draft model generates k sequential tokens.
2. The large Target model calculates the logits for the sequence prefix + draft_tokens in a single parallel forward pass.
3. If target_token == draft_token, it is accepted.
4. Upon the first rejection, all subsequent draft tokens are discarded, and the target model's generated correction is appended "for free."

Repository Structure

.
├── api/             # FastAPI layer for asynchronous inference
├── decoding/        # Adaptive speculative loop handling early rejection & dynamic k-scaling
├── evaluation/      # Baselines and metric tracking (TPS, Latency, Perplexity)
├── models/          # Memory-isolated wrappers for HuggingFace transformers
├── ui/              # Streamlit dashboard for interactive token analysis
└── README.md

Setup Instructions

Prerequisites

• Python 3.10+
• pip package manager

1. Install Dependencies

pip install torch transformers fastapi uvicorn pydantic streamlit pandas

2. Run the FastAPI Service

Start the backend inference API:

uvicorn api.inference_api:app --host 0.0.0.0 --port 8000

The API will be available at http://localhost:8000. You can view the Swagger UI at http://localhost:8000/docs.

3. Run the Streamlit Dashboard

In a new terminal window, start the interactive UI:

streamlit run ui/app.py

The dashboard will be available at http://localhost:8501.

Results

Using our adaptive k-scaling and confidence-based early rejection, the speculative decoding implementation successfully decouples the generation bottleneck from the large target model.

• The Draft Model handles syntactic structure with high accuracy.
• The Target Model intervenes selectively on complex semantic nodes.
• Latency is halved compared to the naive autoregressive baseline without any degradation in response quality.