Speculative Decoding and KV-Cache Compression: Accelerating LLM Inference Throughput by 4x

def speculative_decode(target_model, draft_model, ctx, k):
    """
    Perform speculative decoding with bundle size k.
    Returns a list of generated tokens.
    """
    generated = []
    while not stop_condition(generated):
        # 1️⃣ Draft stage – cheap forward pass on draft model
        draft_tokens = draft_model.sample(ctx, max_len=k)

        # 2️⃣ Target verification – single forward pass on target
        #    (includes KV‑cache extension for the whole bundle)
        logits = target_model.forward(ctx + draft_tokens)

        # 3️⃣ Compute acceptance probabilities per token
        accept = True
        for i, tok in enumerate(draft_tokens):
            p = softmax(logits[i])[tok]            # pθ
            q = draft_model.prob(ctx, tok)         # qφ
            ratio = p / q
            if ratio < 1.0:
                accept = False
                break
            ctx.append(tok)        # commit token
            generated.append(tok)

        if not accept:
            # Fallback: generate the next token directly from target
            next_tok = target_model.sample_one(ctx)
            ctx.append(next_tok)
            generated.append(next_tok)

    return generated

+--------------------------------------------------------------+
|                     StreamingLLM Runtime                     |
|   +----------------+   +----------------+   +---------------+ |
|   | Input Buffer   |-->| Tokenizer      |-->| Speculative   | |
|   | (socket/IPC)   |   | (BPE)          |   | Decoder       | |
|   +----------------+   +----------------+   +-------+-------+ |
|                                 |               |           |
|                                 v               v           |
|   +----------------+   +----------------+   +---------------+ |
|   | KV‑Cache Layer |<--| HHE Manager    |<--| Sliding Window| |
|   +----------------+   +----------------+   +---------------+ |
|                                 ^               ^           |
|                                 |               |           |
|   +----------------+   +----------------+   +---------------+ |
|   | Post‑Processor |<--| Output Router  |<--| Beam/Spec Dec | |
|   +----------------+   +----------------+   +---------------+ |
+--------------------------------------------------------------+

// kv_cache_hhe.hpp ---------------------------------------------------------
#pragma once
#include <cuda_runtime.h>
#include <cstdint>
#include <vector>
#include <stdexcept>

/**
 * @brief Count‑Min Sketch for heavy‑hitter detection on the GPU.
 *
 * The sketch is a 2‑D array of width `W` and depth `D`.  Each row uses a
 * distinct 32‑bit hash function (Murmur3 with different seeds).  Counter
 * overflow is handled by saturating addition.
 */
class CountMinSketch {
public:
    // --------------------------------------------------------------------
    // Construction / Destruction
    // --------------------------------------------------------------------
    CountMinSketch(uint32_t width, uint32_t depth, cudaStream_t stream = 0);
    ~CountMinSketch();

    // --------------------------------------------------------------------
    // API
    // --------------------------------------------------------------------
    /**
     * @brief Increment the counters for token `token_id` by `delta`.
     *
     * This kernel runs in a single warp; the caller must ensure that the
     * launch grid contains exactly one block.
     */
    void update(uint32_t token_id, uint32_t delta = 1);

    /**
     * @brief Estimate the frequency of `token_id`.  Returns the minimum
     *        counter across all rows (the CMS guarantee).
     */
    uint32_t estimate(uint32_t token_id) const;

    /**
     * @brief Reset the entire sketch (used when sliding the window).
     */
    void reset();

private:
    // --------------------------------------------------------------------
    // Device data
    // --------------------------------------------------------------------
    uint32_t *d_counters_;   // Device pointer: [depth][width]
    uint32_t  width_;
    uint32_t  depth_;
    cudaStream_t stream_;

    // --------------------------------------------------------------------
    // Helpers – hash functions (Murmur3 variants)
    // --------------------------------------------------------------------
    __device__ __forceinline__ uint32_t hash(uint32_t token_id, uint32_t seed) const;
};

// kv_cache_hhe.cu ---------------------------------------------------------
#include "kv_cache_hhe.hpp"
#include <cuda.h>
#include <cstdio>

__device__ __forceinline__ uint32_t CountMinSketch::hash(uint32_t token_id,
                                                        uint32_t seed) const {
    // Simple 32‑bit mix – not cryptographically secure but fast.
    uint32_t x = token_id ^ seed;
    x ^= x >> 16;
    x *= 0x85ebca6b;
    x ^= x >> 13;
    x *= 0xc2b2ae35;
    x ^= x >> 16;
    return x % width_;
}

// --------------------------------------------------------------------
// Constructor
// --------------------------------------------------------------------
CountMinSketch::CountMinSketch(uint32_t width, uint32_t depth,
                               cudaStream_t stream)
    : width_(width), depth_(depth), stream_(stream) {
    if (width == 0 || depth == 0) {
        throw std::invalid_argument("CMS dimensions must be > 0");
    }
    size_t bytes = static_cast<size_t>(width_) * depth_ * sizeof(uint32_t);
    cudaError_t err = cudaMalloc(&d_counters_, bytes);
    if (err != cudaSuccess) {
        throw std::runtime_error("cudaMalloc failed for CMS counters");
    }
    // Zero‑initialize
    err = cudaMemsetAsync(d_counters_, 0, bytes, stream_);
    if (err != cudaSuccess) {
        cudaFree(d_counters_);
        throw std::runtime_error("cudaMemsetAsync failed for CMS counters");
    }
}

// --------------------------------------------------------------------
// Destructor
// --------------------------------------------------------------------
CountMinSketch::~CountMinSketch() {
    if (d_counters_) cudaFree(d_counters_);
}

// --------------------------------------------------------------------
// Kernel: update single token (called from attention kernel)
// --------------------------------------------------------------------
__global__ static void cms_update_kernel(uint32_t *counters,
                                         uint32_t width,
                                         uint32_t depth,
                                         uint32_t token_id,
                                         uint32_t delta) {
    // One thread per row (depth <= 32 is typical)
    uint32_t row = threadIdx.x;
    if (row >= depth) return;

    // Compute column index via row‑specific hash
    uint32_t col = ((token_id ^ (0x9e3779b9 + row * 0x9e3779b9)) % width);
    uint32_t *cell = &counters[row * width + col];

    // Saturating add (avoid overflow on 32‑bit counters)
    uint32_t old = atomicAdd(cell, delta);
    if (old + delta < old) { // overflowed
        atomicExch(cell, UINT32_MAX);
    }
}

void CountMinSketch::update(uint32_t token_id, uint32_t delta) {
    // Launch a single block with `depth_` threads.
    cms_update_kernel<<<1, depth_, 0, stream_>>>(
        d_counters_, width_, depth_, token_id, delta);
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        throw std::runtime_error("CMS update kernel launch failed");
    }
}

// --------------------------------------------------------------------
// Kernel: estimate frequency (returns min across rows)
// --------------------------------------------------------------------
__global__ static void cms_estimate_kernel(const uint32_t *counters,
                                           uint32_t width,
                                           uint32_t depth,
                                           uint32_t token_id,
                                           uint32_t *out) {
    __shared__ uint32_t s_min;
    if (threadIdx.x == 0) s_min = UINT32_MAX;
    __syncthreads();

    uint32_t row = threadIdx.x;
    if (row < depth) {
        uint32_t col = ((token_id ^ (0x9e3779b9 + row * 0x9e3779b9)) % width);
        uint32_t val = counters[row * width + col];
        atomicMin(&s_min, val);
    }
    __syncthreads();

    if (threadIdx.x == 0) *out = s_min;
}

uint32_t CountMinSketch::estimate(uint32_t token_id) const {
    uint32_t *d_out;
    cudaMalloc(&d_out, sizeof(uint32_t));
    cms_estimate_kernel<<<1, depth_, 0, stream_>>>(
        d_counters_, width_, depth_, token_id, d_out);
    uint32_t host_out;
    cudaMemcpy(&host_out, d_out, sizeof(uint32_t), cudaMemcpyDeviceToHost);
    cudaFree(d_out);
    return host_out;
}

// --------------------------------------------------------------------
// Reset
// --------------------------------------------------------------------
void CountMinSketch::reset() {
    size_t bytes = static_cast<size_t>(width_) * depth_ * sizeof(uint32_t);
    cudaMemsetAsync(d_counters_, 0, bytes, stream_);
}

cms.update(token_id)

nvprof

// attention.cu (excerpt)
template <typename T>
__global__ void attention_kernel(const T *Q, const T *K, const T *V,
                                 T *O,
                                 uint32_t seq_len,
                                 uint32_t head_dim,
                                 CountMinSketch cms) {
    // ... existing block‑wise GEMM code omitted for brevity ...

    // Each thread processes a single query token.
    uint32_t q_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (q_idx >= seq_len) return;

    // Compute attention scores (dot‑product)
    T score = dot_product(&Q[q_idx * head_dim], K, seq_len, head_dim);

    // Record heavy‑hitter statistics
    cms.update(q_idx);

    // ... softmax, weighted sum, write O ...
}

W_h = 8192

W_h

struct KVEntry {
    alignas(64) uint16_t key[HEAD_DIM];   // 2 × HEAD_DIM bytes
    alignas(64) uint16_t value[HEAD_DIM];
    uint32_t token_id;
    uint64_t timestamp;   // monotonically increasing per token
};

W_h

(score, index)

// kv_window_manager.cu ----------------------------------------------------
#include <cuda_runtime.h>
#include <cstdint>
#include <vector>
#include <stdexcept>
#include <thrust/device_vector.h>
#include <thrust/functional.h>

template <typename T>
struct HeapEntry {
    float score;
    uint32_t idx;   // index into KV buffer
};

template <typename T>
class SlidingWindowCache {
public:
    SlidingWindowCache(uint32_t max_window,
                       uint32_t head_dim,
                       cudaStream_t stream = 0);
    ~SlidingWindowCache();

    // Insert a new token (called after each generation step)
    void push(const KVEntry &entry);

    // Evict tokens that exceed the horizon or have low score
    void evict_if_needed(const CountMinSketch &cms,
                         uint64_t cur_timestamp);

    // Accessors for attention kernels
    const KVEntry* device_ptr() const { return d_kv_; }

private:
    uint32_t max_window_;
    uint32_t head_dim_;
    uint32_t size_;                     // current number of entries
    KVEntry *d_kv_;                     // contiguous GPU buffer
    HeapEntry<T> *d_heap_;              // device min‑heap
    cudaStream_t stream_;

    // Helper: recompute score for a given entry
    __device__ float compute_score(const KVEntry &e,
                                   const CountMinSketch &cms,
                                   uint64_t cur_ts) const;
};

template <typename T>
SlidingWindowCache<T>::SlidingWindowCache(uint32_t max_window,
                                          uint32_t head_dim,
                                          cudaStream_t stream)
    : max_window_(max_window), head_dim_(head_dim), size_(0), stream_(stream) {
    size_t kv_bytes = static_cast<size_t>(max_window_) * sizeof(KVEntry);
    if (cudaMalloc(&d_kv_, kv_bytes) != cudaSuccess) {
        throw std::runtime_error("Failed to allocate KV buffer");
    }
    // Heap size = max_window (worst case every entry is in heap)
    size_t heap_bytes = static_cast<size_t>(max_window_) * sizeof(HeapEntry<T>);
    if (cudaMalloc(&d_heap_, heap_bytes) != cudaSuccess) {
        cudaFree(d_kv_);
        throw std::runtime_error("Failed to allocate heap buffer");
    }
}

template <typename T>
SlidingWindowCache<T>::~SlidingWindowCache() {
    cudaFree(d_kv_);
    cudaFree(d_heap_);
}

// --------------------------------------------------------------------
// Device helper for score computation
// --------------------------------------------------------------------
template <typename T>
__device__ float SlidingWindowCache<T>::compute_score(const KVEntry &e,
                                                      const CountMinSketch &cms,
                                                      uint64_t cur_ts) const {
    uint64_t age = cur_ts - e.timestamp;
    uint32_t freq = cms.estimate(e.token_id);
    // Guard against divide‑by‑zero
    freq = max(freq, 1u);
    const float alpha = 0.7f;
    const float beta  = 0.3f;
    return alpha / static_cast<float>(age) + beta / static_cast<float>(freq);
}

// --------------------------------------------------------------------
// push: insert at the tail of the buffer (no heap update yet)
// --------------------------------------------------------------------
template <typename T>
void SlidingWindowCache<T>::push(const KVEntry &entry) {
    if (size_ >= max_window_) {
        throw std::runtime_error("Window overflow – call evict_if_needed first");
    }
    // Copy entry to device (async)
    cudaMemcpyAsync(d_kv_ + size_, &entry, sizeof(KVEntry),
                    cudaMemcpyHostToDevice, stream_);
    ++size_;
}

// --------------------------------------------------------------------
// evict_if_needed: recompute heap, pop low‑score entries
// --------------------------------------------------------------------
template <typename T>
void SlidingWindowCache<T>::evict_if_needed(const CountMinSketch &cms,
                                            uint64_t cur_timestamp) {
    // 1. Build a temporary array of scores on the GPU
    //    (thrust transform + copy to heap)
    thrust::device_vector<float> d_scores(size_);
    thrust::transform(
        thrust::device, d_kv_, d_kv_ + size_,
        d_scores.begin(),
        [cms, cur_timestamp] __device__(const KVEntry &e) {
            uint64_t age = cur_timestamp - e.timestamp;
            uint32_t freq = cms.estimate(e.token_id);
            freq = max(freq, 1u);
            const float alpha = 0.7f, beta = 0.3f;
            return alpha / static

---

## Part 3 – Production‑Grade Serving & Empirical Validation  

*Speculative Decoding and KV‑Cache Compression: Accelerating LLM Inference Throughput by 4×*  

> **Scope** – This final part translates the research‑level constructs from Parts 1‑2 into a battle‑tested production pipeline. We will (1) wire‑up **vLLM** with a **draft Llama‑3.1‑8B** model, (2) expose a repeatable **FSBL (FutureStack Benchmark Lab) benchmark suite**, (3) enumerate operational failure modes with a concrete runbook, and (4) close the loop with a short FAQ and distilled architectural takeaways.  

> **Audience** – Senior staff engineers, platform architects, and performance‑critical MLOps teams who need a “run‑it‑today” reference implementation that can be dropped into an existing Kubernetes‑based inference fleet.  

---  

## 6. Production Serving Setup – Configuring vLLM Speculative Decoding with Llama‑3.1‑8B Draft  

### 6.1 Why vLLM?  

`vLLM` (v0.4.2‑rc1 at time of writing) is the de‑facto reference implementation for **paged KV‑cache**, **tensor‑parallel inference**, and **speculative decoding**. Its design isolates the **draft** and **target** model pipelines into two distinct CUDA streams, enabling **asynchronous token‑level verification** with a single GPU kernel per step.  

Key vLLM features leveraged in our production stack:  

| Feature | Relevance to Speculative Decoding | Implementation Detail |
|---------|-----------------------------------|-----------------------|
| **Paged KV‑cache** | Allows us to compress the cache on‑the‑fly (see Part 2) without sacrificing random‑access latency. | `--enable-prefix-caching` + `--kv-compress-ratio 0.5` |
| **Speculative Decoding API** | Exposes a “draft model” handle that can be swapped at runtime. | `--draft-model <path>` |
| **GPU‑Direct Scheduler** | Binds each request to a dedicated CUDA stream, eliminating host‑side synchronization bottlenecks. | `--scheduler-type=gpu_direct` |
| **Dynamic Batching** | Auto‑tunes batch size per‑GPU based on token‑arrival patterns, crucial for high‑concurrency TTFT guarantees. | `--max-batch-size 256` |

### 6.2 Draft‑Target Model Pairing  

| Role | Model | Parameter Count | Quantization | Checkpoint Format |
|------|-------|----------------|--------------|-------------------|
| **Target** | Llama‑3.1‑70B (FP16) | 70 B | FP16 (no quant) | `safetensors` |
| **Draft** | Llama‑3.1‑8B (Q4_0) | 8 B | 4‑bit `bitsandbytes` | `gguf` (v2) |

> **Rationale** – The 8 B draft is small enough to stay resident in **GPU shared memory** (≈ 2 GB on H100) while still providing a *high‑recall* next‑token distribution. The 70 B target runs in **tensor‑parallel (TP=8)** across the 8 H100 SXM5 GPUs of a single node, each handling a *shard* of the model matrix.  

### 6.3 Docker‑Ready vLLM Service  

Below is a minimal, production‑oriented Dockerfile that builds vLLM with the necessary CUDA extensions and bundles both models.  

```dockerfile
# --------------------------------------------------------------
# Dockerfile – vLLM Speculative Service (FutureStack)
# --------------------------------------------------------------
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04 AS base

# ---- System deps ------------------------------------------------
RUN apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y \
    python3.11 python3-pip git curl libaio1 && \
    rm -rf /var/lib/apt/lists/*

# ---- Python environment -----------------------------------------
ENV PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    VLLM_VERSION=0.4.2rc1

RUN python3 -m pip install --upgrade pip && \
    python3 -m pip install "vllm[speculative]==${VLLM_VERSION}" \
    bitsandbytes==0.43.1 \
    safetensors==0.4.2 \
    transformers==4.41.0 \
    torch==2.4.0+cu124 -f https://download.pytorch.org/whl/cu124/torch_stable.html

# ---- Model assets ------------------------------------------------
# Target model (FP16) – stored on NVMe RAID, mounted at /models/target
# Draft model (Q4_0) – stored on SSD, mounted at /models/draft
# The mount points are injected at runtime via Kubernetes `hostPath` volumes.

# ---- Entrypoint -------------------------------------------------
COPY entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh

ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]

entrypoint.sh

#!/usr/bin/env bash
set -euo pipefail

# Environment variables (override via k8s ConfigMap)
: "${MAX_BATCH_SIZE:=256}"
: "${MAX_NUM_SEQS:=512}"
: "${KV_COMPRESS_RATIO:=0.5}"
: "${DRAFT_MODEL_PATH:=/models/draft/llama3.1-8b-q4_0.gguf}"
: "${TARGET_MODEL_PATH:=/models/target/llama3.1-70b-fp16/}"
: "${PORT:=8000}"

# vLLM launch
exec python -m vllm.entrypoints.openai.api_server \
    --model ${TARGET_MODEL_PATH} \
    --draft-model ${DRAFT_MODEL_PATH} \
    --port ${PORT} \
    --max-batch-size ${MAX_BATCH_SIZE} \
    --max-num-seqs ${MAX_NUM_SEQS} \
    --enable-prefix-caching \
    --kv-compress-ratio ${KV_COMPRESS_RATIO} \
    --speculative-decoding \
    --tensor-parallel-size 8 \
    --gpu-memory-utilization 0.95 \
    --scheduler-type=gpu_direct

charts/vllm-speculative

Deployment

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-speculative
  labels:
    app: vllm
spec:
  replicas: 1                     # Scale horizontally via a Service Mesh (e.g., Istio)
  selector:
    matchLabels:
      app: vllm
  template:
    metadata:
      labels:
        app: vllm
    spec:
      nodeSelector:
        accelerator: h100         # Custom label on GPU‑enabled nodes
      tolerations:
        - key: "nvidia.com/gpu"
          operator: Exists
          effect: NoSchedule
      containers:
        - name: vllm
          image: ghcr.io/futurestack/vllm-speculative:latest
          resources:
            limits:
              nvidia.com/gpu: "8"            # 8× H100 per node
              cpu: "96"
              memory: "512Gi"
          env:
            - name: MAX_BATCH_SIZE
              value: "256"
            - name: KV_COMPRESS_RATIO
              value: "0.5"
          volumeMounts:
            - name: target-model
              mountPath: /models/target
            - name: draft-model
              mountPath: /models/draft
          ports:
            - containerPort: 8000
      volumes:
        - name: target-model
          persistentVolumeClaim:
            claimName: target-model-pvc
        - name: draft-model
          persistentVolumeClaim:
            claimName: draft-model-pvc

--enable-tracing

{
  "datasource": "Prometheus",
  "targets": [
    {
      "expr": "histogram_quantile(0.99, sum(rate(vllm_ttft_seconds_bucket[5m])) by (le, concurrency))",
      "legendFormat": "{{concurrency}}"
    }
  ],
  "type": "timeseries",
  "fieldConfig": {
    "defaults": {
      "unit": "ms",
      "min": 0
    }
  }
}