DEV Community: Yuvraj singh Bhadoria

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Yuvraj singh Bhadoria — Sun, 30 Aug 2026 08:48:51 +0000

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Speculative decoding is one of the most widely discussed inference optimizations in recent LLM engineering, and frequently one of the most misunderstood. The core proposition sounds ideal: achieving a 2–3× boost in decoding throughput with mathematically identical output distributions—yielding performance gains via a lightweight secondary model.

In practice, speculative decoding functions as a trade-off rather than a guaranteed acceleration: you pay the computational overhead of running a smaller draft model with the expectation that its outputs align sufficiently with the target model to yield a net speedup.

This post details the complete system stack—covering core transformer architecture, memory bandwidth constraints, the draft-then-verify loop, state-of-the-art methodology taxonomies, and empirical benchmark evaluations on GPT-2 weights—to highlight where performance gains originate and where they risk regressing.

Executive Summary

Autoregressive Decoding is Serial & Memory-Bound: Single-token autoregressive generation is limited by memory bandwidth during low-batch inference. Speculative decoding directly targets this single-stream latency bottleneck.
The Core Mechanism is Draft-Then-Verify: A fast draft model proposes $k$ candidate tokens sequentially. The target model then verifies all $k+1$ positions in a single parallel forward pass. Rejection sampling preserves the exact target output distribution.
The Speedup Has a Hard Inequality Constraint ($\alpha > c$): Net performance gains require the draft model's token acceptance rate ($\alpha$) to strictly exceed its relative computational cost fraction ($c$). Failing this condition increases overall latency.
Architectural Taxonomy: Current approaches (Speculative Sampling, Medusa, EAGLE, Self-Speculative, Lookahead) vary primarily in their draft generation mechanism. End-to-end performance depends directly on draft–target distribution alignment rather than the verification loop itself.
Empirical Findings: Evaluations on un-aligned weights yielded a speedup of 0.2–0.9× (a net performance penalty). This negative result highlights the operational necessity of validating draft alignment prior to deployment.

Visual Overview

Figure 0: High-level visual summary showing serial memory constraints, drafting execution, taxonomy, framework integration, and performance benchmarking.

1. Model Setup & Architecture

To evaluate speculative decoding, we must analyze the hardware execution costs of a single forward pass. A causal language model generates next-token probability distributions by executing tensor operations across its transformer stack. This pipeline is bounded by two distinct factors:

Matrix multiplication operations (GEMM FLOPS across transformer layers).
Memory bus transfers (fetching model weights and KV-cache states from DRAM to SRAM/registers).

For small batch sizes ($B=1$) on single-stream inference, memory transfers dominate total step latency: compute units execute quickly and subsequently stall while waiting for weight loading. Speculative decoding specifically targets this architectural bottleneck.

[Target Model (Ground Truth): GPT-2 124M]
d_model=768, depth=12, heads=12, Vocab=50k
                       ▲
                       │ (pt)
             ┌───────────────────┐
             │ Target Final Head │
             └─────────▲─────────┘
                       │
             ┌───────────────────┐
             │  GPT-2 (12 Layers)│
             └─────────▲─────────┘
                       │
             ┌───────────────────┐
             │  Input Embeddings │
             └───────────────────┘

Key Takeaway: A draft mechanism provides net throughput benefits only when it is computationally inexpensive (low layer count/parameter footprint) and statistically aligned with the target distribution.

Figure 1: Structural comparison between the primary target model (GPT-2 124M), an independent draft model (tiny-GPT2 10M), and self-speculative early-exit configurations.

2. The Serial Memory Bottleneck

Transformer architectures process input sequences in parallel during context encoding (prefill phase), but execute sequentially across step iterations during auto-regressive decoding (decoding phase). Token $t_i$ cannot be evaluated until token $t_{i-1}$ is generated, due to causal self-attention dependencies across historic Key-Value (KV) states.

$$\text{Latency}{\text{autoregressive}} = N{\text{tokens}} \times t_{\text{per-token}}$$

Panel A: Autoregressive Bottleneck (Serial)
[Step t]  ──(Fetch Weights & KV)──► [Token 1] ──┐
[Step t+1] ──(Fetch Weights & KV)──► [Token 2] ──┼─► High Memory Stall / Idle Compute
[Step t+2] ──(Fetch Weights & KV)──► [Token 3] ──┘

Panel B: Speculative Verification (Parallel)
[Single Forward Pass] ──(Fetch Weights Once)──► [Verify Tokens 1, 2, 3, 4 Simultaneously]

At low batch sizes, single-token generation iterations fail to fully saturate GPU compute pipelines.

Figure 2: Comparison of memory-bandwidth-bound serial autoregressive generation (Panel A) against batch-parallel target verification across k+1 token positions (Panel B).

3. The Speculative Sampling Mechanism

The speculative sampling execution pipeline follows a three-step cycle:

Draft Phase: The lightweight draft model generates $k$ candidate tokens sequentially: $$\hat{x}{1}, \hat{x}{2}, \dots, \hat{x}{k} \sim p{d}(x \mid \text{context})$$
Verification Phase: The target model processes the concatenated sequence $\text{context} \cup {\hat{x}{1} \dots \hat{x}{k}}$ in one forward pass, computing target logits for all $k+1$ token positions in parallel.
Correction Phase: Rejection sampling is applied sequentially across candidate tokens. The first rejected token $\hat{x}i$ is resampled from the corrected residual distribution: $$p{\text{adjusted}}(x) = \text{relu}\left(p_{t}(x) - p_{d}(x)\right)$$

This mathematical formulation guarantees that the final output distribution remains provably identical to sampling directly from the target model.

Figure 3: Detailed control-flow loop showing sequential draft generation, parallel target model scoring, and distribution-preserving rejection sampling.

Sampling & Rejection Implementation

import torch
import torch.nn.functional as F

def spec_sample(seq, k, draft_fn, target_fn, temp=0.8):
    """Executes speculative decoding with exact target distribution preservation."""
    draft_tokens, draft_logits = [], []
    current_seq = list(seq)

    # 1. Draft Step: Generate k candidates sequentially
    for _ in range(k):
        logits = draft_fn(current_seq)
        next_logit = logits[-1]
        draft_logits.append(next_logit)
        token = torch.multinomial(F.softmax(next_logit / temp, dim=-1), 1).item()
        draft_tokens.append(token)
        current_seq.append(token)

    # 2. Verify Step: Parallel validation over k+1 positions
    target_logits = target_fn(seq + draft_tokens)
    prefix_offset = len(seq)
    accepted_count = 0

    # 3. Correction Step: Rejection sampling loop
    for i in range(k):
        pt = F.softmax(target_logits[prefix_offset - 1 + i] / temp, dim=-1)
        pd = F.softmax(draft_logits[i] / temp, dim=-1)
        candidate = draft_tokens[i]

        # Accept condition
        if torch.rand(1).item() < min(1.0, (pt[candidate] / pd[candidate]).item()):
            accepted_count += 1
        else:
            # Reject: Resample candidate from adjusted distribution (pt - pd)+
            residual = torch.clamp(pt - pd, min=0.0)
            residual /= residual.sum()
            resampled_token = torch.multinomial(residual, 1).item()
            return seq + draft_tokens[:accepted_count] + [resampled_token], accepted_count

    # Bonus token sampling if all k drafted tokens are accepted
    bonus_token = torch.multinomial(F.softmax(target_logits[prefix_offset + k - 1] / temp, dim=-1), 1).item()
    return seq + draft_tokens + [bonus_token], k

4. The Speculative-Decoding Family

While all speculative decoding variants rely on the same fundamental parallel verification framework, they differ in their structural draft generation mechanisms:

Approach	Draft Generation Architecture	Target Acceleration	Primary Operational Trade-off
Speculative Sampling	Independent auxiliary Small LM	2.0–3.0×	Requires hosting separate draft model & tokenizer alignment
Medusa	Multi-head prediction heads on target	2.3–3.0×	Requires parameter fine-tuning of prediction heads
EAGLE	Feature-level drafting with tree attention	2.0–3.0×	High draft acceptance rate; higher system complexity
Self-Speculative	Target model internal early-exiting	1.3–1.8×	Zero additional weight hosting; limited by layer alignment
Lookahead	Jacobi iteration / N-gram retrieval	1.8–2.3×	Parameter-free; highly sequence/task dependent

Figure 4: Classification taxonomy of speculative decoding variants structured by draft generation method.

Production Framework Integrations

Hugging Face Transformers (`assisted_generation`)

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
target_model = AutoModelForCausalLM.from_pretrained("gpt2")
assistant_model = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")

inputs = tokenizer("The future of artificial intelligence is", return_tensors="pt")
outputs = target_model.generate(
    inputs.input_ids,
    assistant_model=assistant_model,
    do_sample=True,
    temperature=0.8,
    max_new_tokens=50
)

vLLM Native Scheduler Integration

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-8B-Instruct \
  --speculative-model meta-llama/Llama-3-1B-Instruct \
  --num_speculative_tokens 5 \
  --max-model-len 4096

5. Benchmarking & Empirical Performance

Mathematical Speedup Condition ($\alpha > c$)

Let $c$ represent the ratio of draft model execution cost relative to target model execution cost per token:

$$c = \frac{\text{Cost}{\text{draft}}}{\text{Cost}{\text{target}}}$$

Let $\alpha$ represent the average token acceptance rate across speculative steps. The theoretical speedup factor $S$ relative to standard autoregressive execution is modeled as:

$$S \approx \frac{1 + k \cdot \alpha}{1 + k \cdot c}$$

To achieve a net speedup ($S > 1$), the pipeline must satisfy the inequality:

$$\alpha > c$$

If candidate acceptance falls below the cost threshold ($\alpha < c$), the computational overhead of draft generation and verification outpaces the benefits of sequence amortization, leading to increased latency.

Figure 5: Measured wall-clock performance curve showing speedup vs. acceptance rate ($\alpha$) relative to relative cost fraction ($c$). Misaligned drafts fail to clear the baseline threshold.

Measured Experimental Results

Evaluating speculative decoding across unaligned draft configurations demonstrates the real-world operational impact of the $\alpha > c$ constraint:

Draft Configuration	Relative Cost Fraction ($c$)	Acceptance Rate ($\alpha$)	Speedup ($k=1$)	Speedup ($k=3$)	Net Performance Impact
`tiny-gpt2` (10M vs 124M)	~0.08	4.0%	0.68×	0.44×	Performance Penalty
`gpt2` Early-Exit (Layer 3/12)	~0.25	12.0%	0.68×	0.38×	Performance Penalty
`gpt2` Early-Exit (Layer 10/12)	~0.83	29.0%	0.38×	0.35×	Performance Penalty

    Speedup (x Baseline)
    1.2x ┼─────────────────────────────────────────────────── (Break-even: 1.0x)
    1.0x ┼───────────────────────────────────────────────────
    0.8x ┼───── Top Performance (k=1, alpha=4%..12%): ~0.68x
    0.6x ┼───────────────────────────────────────────────────
    0.4x ┼───────────────── Top Performance (k=3): ~0.35x..0.44x
    0.2x ┼───────────────────────────────────────────────────
         └───────┬───────────────┬───────────────┬───────────
                k=1             k=2             k=3

System Failure Modes

Distributional Mismatch ($\alpha < c$): Using an unaligned draft model causes frequent rejection steps, incurring severe draft-overhead penalties.
High Batch Aggregation ($B \gg 1$): At high query volumes, hardware compute pipelines shift from memory-bound to compute-bound states. Under these conditions, speculative verification offers diminishing latency returns while increasing total FLOP utilization.
High-Entropy Generation Tasks: Tasks with high logical complexity (e.g., code generation or complex mathematical reasoning) exhibit lower acceptance rates ($\alpha$), reducing maximum attainable sequence extensions per step.

Conclusion

Speculative decoding is a powerful technique for accelerating single-stream LLM inference, but its benefits are fundamentally conditional. Achieving real-world speedups requires strict optimization of draft alignment ($\alpha$) relative to system cost ($c$). When deploying speculative pipelines in production systems, profile target-draft acceptance rates under actual workload distributions before enabling parallel verification logic.

References

Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. International Conference on Machine Learning (ICML). arXiv:2211.17192
Chen, C., Borgeaud, S., Zhou, G., Steiner, D., & Chen, Z. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv preprint. arXiv:2302.01318
Cai, T., Li, Y., Geng, Z., Peng, L., Li, F., & Xiao, W. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. International Conference on Machine Learning (ICML). arXiv:2401.10774
Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. International Conference on Machine Learning (ICML). arXiv:2401.15077
Zhang, J., Wang, J., Huang, H., Chen, Y., & Zhou, W. (2023). Self-Speculative Decoding with Self-Draft and Self-Verification. Empirical Methods in Natural Language Processing (EMNLP). arXiv:2311.08466
Fu, Y., Bailis, P., Stoica, I., & Zhang, H. (2024). Break the Sequential Dependency of LLM Inference Using Lookahead Decoding. arXiv preprint. arXiv:2402.02057

Original blog: https://YuvrajSinghBhadoria2.github.io/spec-decoding-blog/

Tags: #llm #machinelearning #performance #inference

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Yuvraj singh Bhadoria — Sun, 30 Aug 2026 08:41:43 +0000

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Executive Summary

Autoregressive Decoding is Serial & Memory-Bound: Single-token autoregressive generation is limited by memory bandwidth during low-batch inference. Speculative decoding directly targets this single-stream latency bottleneck.
The Core Mechanism is Draft-Then-Verify: A fast draft model proposes $k$ candidate tokens sequentially. The target model then verifies all $k+1$ positions in a single parallel forward pass. Rejection sampling preserves the exact target output distribution.
The Speedup Has a Hard Inequality Constraint ($\alpha > c$): Net performance gains require the draft model's token acceptance rate ($\alpha$) to strictly exceed its relative computational cost fraction ($c$). Failing this condition increases overall latency.
Architectural Taxonomy: Current approaches (Speculative Sampling, Medusa, EAGLE, Self-Speculative, Lookahead) vary primarily in their draft generation mechanism. End-to-end performance depends directly on draft–target distribution alignment rather than the verification loop itself.
Empirical Findings: Evaluations on un-aligned weights yielded a speedup of 0.2–0.9× (a net performance penalty). This negative result highlights the operational necessity of validating draft alignment prior to deployment.

Visual Overview

Figure 0: High-level visual summary showing serial memory constraints, drafting execution, taxonomy, framework integration, and performance benchmarking.

1. Model Setup & Architecture

Matrix multiplication operations (GEMM FLOPS across transformer layers).
Memory bus transfers (fetching model weights and KV-cache states from DRAM to SRAM/registers).

[Target Model (Ground Truth): GPT-2 124M]
d_model=768, depth=12, heads=12, Vocab=50k
                       ▲
                       │ (pt)
             ┌───────────────────┐
             │ Target Final Head │
             └─────────▲─────────┘
                       │
             ┌───────────────────┐
             │  GPT-2 (12 Layers)│
             └─────────▲─────────┘
                       │
             ┌───────────────────┐
             │  Input Embeddings │
             └───────────────────┘

2. The Serial Memory Bottleneck

$$\text{Latency}{\text{autoregressive}} = N{\text{tokens}} \times t_{\text{per-token}}$$

Panel A: Autoregressive Bottleneck (Serial)
[Step t]  ──(Fetch Weights & KV)──► [Token 1] ──┐
[Step t+1] ──(Fetch Weights & KV)──► [Token 2] ──┼─► High Memory Stall / Idle Compute
[Step t+2] ──(Fetch Weights & KV)──► [Token 3] ──┘

Panel B: Speculative Verification (Parallel)
[Single Forward Pass] ──(Fetch Weights Once)──► [Verify Tokens 1, 2, 3, 4 Simultaneously]

At low batch sizes, single-token generation iterations fail to fully saturate GPU compute pipelines.

Figure 2: Comparison of memory-bandwidth-bound serial autoregressive generation (Panel A) against batch-parallel target verification across k+1 token positions (Panel B).

3. The Speculative Sampling Mechanism

The speculative sampling execution pipeline follows a three-step cycle:

Draft Phase: The lightweight draft model generates $k$ candidate tokens sequentially: $$\hat{x}{1}, \hat{x}{2}, \dots, \hat{x}{k} \sim p{d}(x \mid \text{context})$$
Verification Phase: The target model processes the concatenated sequence $\text{context} \cup {\hat{x}{1} \dots \hat{x}{k}}$ in one forward pass, computing target logits for all $k+1$ token positions in parallel.
Correction Phase: Rejection sampling is applied sequentially across candidate tokens. The first rejected token $\hat{x}i$ is resampled from the corrected residual distribution: $$p{\text{adjusted}}(x) = \text{relu}\left(p_{t}(x) - p_{d}(x)\right)$$

This mathematical formulation guarantees that the final output distribution remains provably identical to sampling directly from the target model.

Figure 3: Detailed control-flow loop showing sequential draft generation, parallel target model scoring, and distribution-preserving rejection sampling.

Sampling & Rejection Implementation

import torch
import torch.nn.functional as F

def spec_sample(seq, k, draft_fn, target_fn, temp=0.8):
    """Executes speculative decoding with exact target distribution preservation."""
    draft_tokens, draft_logits = [], []
    current_seq = list(seq)

    # 1. Draft Step: Generate k candidates sequentially
    for _ in range(k):
        logits = draft_fn(current_seq)
        next_logit = logits[-1]
        draft_logits.append(next_logit)
        token = torch.multinomial(F.softmax(next_logit / temp, dim=-1), 1).item()
        draft_tokens.append(token)
        current_seq.append(token)

    # 2. Verify Step: Parallel validation over k+1 positions
    target_logits = target_fn(seq + draft_tokens)
    prefix_offset = len(seq)
    accepted_count = 0

    # 3. Correction Step: Rejection sampling loop
    for i in range(k):
        pt = F.softmax(target_logits[prefix_offset - 1 + i] / temp, dim=-1)
        pd = F.softmax(draft_logits[i] / temp, dim=-1)
        candidate = draft_tokens[i]

        # Accept condition
        if torch.rand(1).item() < min(1.0, (pt[candidate] / pd[candidate]).item()):
            accepted_count += 1
        else:
            # Reject: Resample candidate from adjusted distribution (pt - pd)+
            residual = torch.clamp(pt - pd, min=0.0)
            residual /= residual.sum()
            resampled_token = torch.multinomial(residual, 1).item()
            return seq + draft_tokens[:accepted_count] + [resampled_token], accepted_count

    # Bonus token sampling if all k drafted tokens are accepted
    bonus_token = torch.multinomial(F.softmax(target_logits[prefix_offset + k - 1] / temp, dim=-1), 1).item()
    return seq + draft_tokens + [bonus_token], k

4. The Speculative-Decoding Family

While all speculative decoding variants rely on the same fundamental parallel verification framework, they differ in their structural draft generation mechanisms:

Approach	Draft Generation Architecture	Target Acceleration	Primary Operational Trade-off
Speculative Sampling	Independent auxiliary Small LM	2.0–3.0×	Requires hosting separate draft model & tokenizer alignment
Medusa	Multi-head prediction heads on target	2.3–3.0×	Requires parameter fine-tuning of prediction heads
EAGLE	Feature-level drafting with tree attention	2.0–3.0×	High draft acceptance rate; higher system complexity
Self-Speculative	Target model internal early-exiting	1.3–1.8×	Zero additional weight hosting; limited by layer alignment
Lookahead	Jacobi iteration / N-gram retrieval	1.8–2.3×	Parameter-free; highly sequence/task dependent

Production Framework Integrations

Hugging Face Transformers (`assisted_generation`)

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
target_model = AutoModelForCausalLM.from_pretrained("gpt2")
assistant_model = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")

inputs = tokenizer("The future of artificial intelligence is", return_tensors="pt")
outputs = target_model.generate(
    inputs.input_ids,
    assistant_model=assistant_model,
    do_sample=True,
    temperature=0.8,
    max_new_tokens=50
)

vLLM Native Scheduler Integration

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-8B-Instruct \
  --speculative-model meta-llama/Llama-3-1B-Instruct \
  --num_speculative_tokens 5 \
  --max-model-len 4096

5. Benchmarking & Empirical Performance

Mathematical Speedup Condition ($\alpha > c$)

Let $c$ represent the ratio of draft model execution cost relative to target model execution cost per token:

$$c = \frac{\text{Cost}{\text{draft}}}{\text{Cost}{\text{target}}}$$

Let $\alpha$ represent the average token acceptance rate across speculative steps. The theoretical speedup factor $S$ relative to standard autoregressive execution is modeled as:

$$S \approx \frac{1 + k \cdot \alpha}{1 + k \cdot c}$$

To achieve a net speedup ($S > 1$), the pipeline must satisfy the inequality:

$$\alpha > c$$

Figure 5: Measured wall-clock performance curve showing speedup vs. acceptance rate ($\alpha$) relative to relative cost fraction ($c$). Misaligned drafts fail to clear the baseline threshold.

Measured Experimental Results

Evaluating speculative decoding across unaligned draft configurations demonstrates the real-world operational impact of the $\alpha > c$ constraint:

Draft Configuration	Relative Cost Fraction ($c$)	Acceptance Rate ($\alpha$)	Speedup ($k=1$)	Speedup ($k=3$)	Net Performance Impact
`tiny-gpt2` (10M vs 124M)	~0.08	4.0%	0.68×	0.44×	Performance Penalty
`gpt2` Early-Exit (Layer 3/12)	~0.25	12.0%	0.68×	0.38×	Performance Penalty
`gpt2` Early-Exit (Layer 10/12)	~0.83	29.0%	0.38×	0.35×	Performance Penalty

    Speedup (x Baseline)
    1.2x ┼─────────────────────────────────────────────────── (Break-even: 1.0x)
    1.0x ┼───────────────────────────────────────────────────
    0.8x ┼───── Top Performance (k=1, alpha=4%..12%): ~0.68x
    0.6x ┼───────────────────────────────────────────────────
    0.4x ┼───────────────── Top Performance (k=3): ~0.35x..0.44x
    0.2x ┼───────────────────────────────────────────────────
         └───────┬───────────────┬───────────────┬───────────
                k=1             k=2             k=3

System Failure Modes

Distributional Mismatch ($\alpha < c$): Using an unaligned draft model causes frequent rejection steps, incurring severe draft-overhead penalties.
High Batch Aggregation ($B \gg 1$): At high query volumes, hardware compute pipelines shift from memory-bound to compute-bound states. Under these conditions, speculative verification offers diminishing latency returns while increasing total FLOP utilization.
High-Entropy Generation Tasks: Tasks with high logical complexity (e.g., code generation or complex mathematical reasoning) exhibit lower acceptance rates ($\alpha$), reducing maximum attainable sequence extensions per step.

Conclusion

References

Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. International Conference on Machine Learning (ICML). arXiv:2211.17192
Chen, C., Borgeaud, S., Zhou, G., Steiner, D., & Chen, Z. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv preprint. arXiv:2302.01318
Cai, T., Li, Y., Geng, Z., Peng, L., Li, F., & Xiao, W. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. International Conference on Machine Learning (ICML). arXiv:2401.10774
Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. International Conference on Machine Learning (ICML). arXiv:2401.15077
Zhang, J., Wang, J., Huang, H., Chen, Y., & Zhou, W. (2023). Self-Speculative Decoding with Self-Draft and Self-Verification. Empirical Methods in Natural Language Processing (EMNLP). arXiv:2311.08466
Fu, Y., Bailis, P., Stoica, I., & Zhang, H. (2024). Break the Sequential Dependency of LLM Inference Using Lookahead Decoding. arXiv preprint. arXiv:2402.02057

Original blog: https://YuvrajSinghBhadoria2.github.io/spec-decoding-blog/

Tags: #llm #machinelearning #performance #inference

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Yuvraj singh Bhadoria — Sun, 30 Aug 2026 03:35:19 +0000

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Executive Summary

Autoregressive Decoding is Serial & Memory-Bound: Single-token autoregressive generation is limited by memory bandwidth during low-batch inference. Speculative decoding directly targets this single-stream latency bottleneck.
The Core Mechanism is Draft-Then-Verify: A fast draft model proposes $k$ candidate tokens sequentially. The target model then verifies all $k+1$ positions in a single parallel forward pass. Rejection sampling preserves the exact target output distribution.
The Speedup Has a Hard Inequality Constraint ($\alpha > c$): Net performance gains require the draft model's token acceptance rate ($\alpha$) to strictly exceed its relative computational cost fraction ($c$). Failing this condition increases overall latency.
Architectural Taxonomy: Current approaches (Speculative Sampling, Medusa, EAGLE, Self-Speculative, Lookahead) vary primarily in their draft generation mechanism. End-to-end performance depends directly on draft–target distribution alignment rather than the verification loop itself.
Empirical Findings: Evaluations on un-aligned weights yielded a speedup of 0.2–0.9× (a net performance penalty). This negative result highlights the operational necessity of validating draft alignment prior to deployment.

Visual Overview

1. Model Setup & Architecture

Matrix multiplication operations (GEMM FLOPS across transformer layers).
Memory bus transfers (fetching model weights and KV-cache states from DRAM to SRAM/registers).

[Target Model (Ground Truth): GPT-2 124M]
d_model=768, depth=12, heads=12, Vocab=50k
                       ▲
                       │ (pt)
             ┌───────────────────┐
             │ Target Final Head │
             └─────────▲─────────┘
                       │
             ┌───────────────────┐
             │  GPT-2 (12 Layers)│
             └─────────▲─────────┘
                       │
             ┌───────────────────┐
             │  Input Embeddings │
             └───────────────────┘

Experimental Setup Parameters

Target Model: gpt2 (124M parameters, 12 layers).
Draft Models:
1. sshleifer/tiny-gpt2 (10M parameters, 2 layers) serving as an independent small LM draft.
2. gpt2 early-exit heads at intermediate layers $m \in {3, 10}$ serving as a self-speculative draft.
Decoding Strategy: Rejection sampling (Chen et al., 2023) at temperature $T = 0.8$.
Evaluation Metric: Wall-clock output generation rate (tokens/sec) evaluated across benchmark prompts.
Reproducibility: Benchmarks are programmatically generated via spec_experiment.py.

Key Takeaway: A draft mechanism provides net throughput benefits only when it is computationally inexpensive (low layer count/parameter footprint) and statistically aligned with the target distribution.

2. The Serial Memory Bottleneck

$$\text{Latency}{\text{autoregressive}} = N{\text{tokens}} \times t_{\text{per-token}}$$

Panel A: Autoregressive Bottleneck (Serial)
[Step t]  ──(Fetch Weights & KV)──> [Token 1] ──┐
[Step t+1] ──(Fetch Weights & KV)──> [Token 2] ──┼─> High Memory Stall / Idle Compute
[Step t+2] ──(Fetch Weights & KV)──> [Token 3] ──┘

Panel B: Speculative Verification (Parallel)
[Single Forward Pass] ──(Fetch Weights Once)──> [Verify Tokens 1, 2, 3, 4 Simultaneously]

At low batch sizes, single-token generation iterations fail to fully saturate GPU compute pipelines.

Executing a single target forward pass over a sequence of length $L+k$ requires nearly the same memory bandwidth overhead as executing a forward pass over sequence length $L$. Speculative decoding leverages this invariant by amortizing weight transfer costs across $k+1$ candidate tokens simultaneously.

3. The Speculative Sampling Mechanism

The speculative sampling execution pipeline follows a three-step cycle:

Draft Phase: The lightweight draft model generates $k$ candidate tokens sequentially: $$\hat{x}{1}, \hat{x}{2}, \dots, \hat{x}{k} \sim p{d}(x \mid \text{context})$$
Verification Phase: The target model processes the concatenated sequence $\text{context} \cup {\hat{x}{1} \dots \hat{x}{k}}$ in one forward pass, computing target logits for all $k+1$ token positions in parallel.
Correction Phase: Rejection sampling is applied sequentially across candidate tokens. The first rejected token $\hat{x}i$ is resampled from the corrected residual distribution: $$p{\text{adjusted}}(x) = \text{relu}\left(p_{t}(x) - p_{d}(x)\right)$$

This mathematical formulation guarantees that the final output distribution remains provably identical to sampling directly from the target model.

Sampling & Rejection Implementation

import torch
import torch.nn.functional as F

def spec_sample(seq, k, draft_fn, target_fn, temp=0.8):
    """
    Executes speculative decoding with exact target distribution preservation.
    """
    draft_tokens, draft_logits = [], []
    current_seq = list(seq)

    # 1. Draft Step: Generate k candidates sequentially
    for _ in range(k):
        logits = draft_fn(current_seq)
        next_logit = logits[-1]
        draft_logits.append(next_logit)

        token = torch.multinomial(F.softmax(next_logit / temp, dim=-1), 1).item()
        draft_tokens.append(token)
        current_seq.append(token)

    # 2. Verify Step: Parallel validation over k+1 positions
    target_logits = target_fn(seq + draft_tokens)
    prefix_offset = len(seq)
    accepted_count = 0

    # 3. Correction Step: Rejection sampling loop
    for i in range(k):
        pt = F.softmax(target_logits[prefix_offset - 1 + i] / temp, dim=-1)
        pd = F.softmax(draft_logits[i] / temp, dim=-1)
        candidate = draft_tokens[i]

        # Accept condition
        if torch.rand(1).item() < min(1.0, (pt[candidate] / pd[candidate]).item()):
            accepted_count += 1
        else:
            # Reject: Resample candidate from adjusted distribution (pt - pd)+
            residual = torch.clamp(pt - pd, min=0.0)
            residual /= residual.sum()
            resampled_token = torch.multinomial(residual, 1).item()
            return seq + draft_tokens[:accepted_count] + [resampled_token], accepted_count

    # Bonus token sampling if all k drafted tokens are accepted
    bonus_token = torch.multinomial(F.softmax(target_logits[prefix_offset + k - 1] / temp, dim=-1), 1).item()
    return seq + draft_tokens + [bonus_token], k

4. The Speculative-Decoding Family

While all speculative decoding variants rely on the same fundamental parallel verification framework, they differ in their structural draft generation mechanisms:

Approach	Draft Generation Architecture	Target Acceleration	Primary Operational Trade-off
Speculative Sampling	Independent auxiliary Small LM	2.0–3.0×	Requires hosting separate draft model & tokenizer alignment
Medusa	Multi-head prediction heads on target	2.3–3.0×	Requires parameter fine-tuning of prediction heads
EAGLE	Feature-level drafting with tree attention	2.0–3.0×	High draft acceptance rate; higher system complexity
Self-Speculative	Target model internal early-exiting	1.3–1.8×	Zero additional weight hosting; limited by layer alignment
Lookahead	Jacobi iteration / N-gram retrieval	1.8–2.3×	Parameter-free; highly sequence/task dependent

Production Framework Integrations

Hugging Face Transformers (`assisted_generation`)

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
target_model = AutoModelForCausalLM.from_pretrained("gpt2")
assistant_model = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2")

inputs = tokenizer("The future of artificial intelligence is", return_tensors="pt")
outputs = target_model.generate(
    inputs.input_ids,
    assistant_model=assistant_model,
    do_sample=True,
    temperature=0.8,
    max_new_tokens=50
)

vLLM Native Scheduler Integration

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-8B-Instruct \
  --speculative-model meta-llama/Llama-3-1B-Instruct \
  --num_speculative_tokens 5 \
  --max-model-len 4096

5. Benchmarking & Empirical Performance

Mathematical Speedup Condition ($\alpha > c$)

Let $c$ represent the ratio of draft model execution cost relative to target model execution cost per token:

$$c = \frac{\text{Cost}{\text{draft}}}{\text{Cost}{\text{target}}}$$

Let $\alpha$ represent the average token acceptance rate across speculative steps. The theoretical speedup factor $S$ relative to standard autoregressive execution is modeled as:

$$S \approx \frac{1 + k \cdot \alpha}{1 + k \cdot c}$$

To achieve a net speedup ($S > 1$), the pipeline must satisfy the inequality:

$$\alpha > c$$

Measured Experimental Results

Evaluating speculative decoding across unaligned draft configurations demonstrates the real-world operational impact of the $\alpha > c$ constraint:

Draft Configuration	Relative Cost Fraction ($c$)	Acceptance Rate ($\alpha$)	Speedup ($k=1$)	Speedup ($k=3$)	Net Performance Impact
`tiny-gpt2` (10M vs 124M)	~0.08	4.0%	0.68×	0.44×	Performance Penalty
`gpt2` Early-Exit (Layer 3/12)	~0.25	12.0%	0.68×	0.38×	Performance Penalty
`gpt2` Early-Exit (Layer 10/12)	~0.83	29.0%	0.38×	0.35×	Performance Penalty

    Speedup (x Baseline)
    1.2x ┼─────────────────────────────────────────────────── (Break-even: 1.0x)
    1.0x ┼───────────────────────────────────────────────────
    0.8x ┼───── Top Performance (k=1, alpha=4%..12%): ~0.68x
    0.6x ┼───────────────────────────────────────────────────
    0.4x ┼───────────────── Top Performance (k=3): ~0.35x..0.44x
    0.2x ┼───────────────────────────────────────────────────
         └───────┬───────────────┬───────────────┬───────────
               k=1             k=2             k=3

System Failure Modes

Distributional Mismatch ($\alpha < c$): Using an unaligned draft model causes frequent rejection steps, incurring severe draft-overhead penalties.
High Batch Aggregation ($B \gg 1$): At high query volumes, hardware compute pipelines shift from memory-bound to compute-bound states. Under these conditions, speculative verification offers diminishing latency returns while increasing total FLOP utilization.
High-Entropy Generation Tasks: Tasks with high logical complexity (e.g., code generation or complex mathematical reasoning) exhibit lower acceptance rates ($\alpha$), reducing maximum attainable sequence extensions per step.

Conclusion

References

Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. International Conference on Machine Learning (ICML). arXiv:2211.17192
Chen, C., Borgeaud, S., Zhou, G., Steiner, D., & Chen, Z. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv preprint. arXiv:2302.01318
Cai, T., Li, Y., Geng, Z., Peng, L., Li, F., & Xiao, W. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. International Conference on Machine Learning (ICML). arXiv:2401.10774
Li, Y., Wei, F., Zhang, C., & Zhang, H. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. International Conference on Machine Learning (ICML). arXiv:2401.15077
Zhang, J., Wang, J., Huang, H., Chen, Y., & Zhou, W. (2023). Self-Speculative Decoding with Self-Draft and Self-Verification. Empirical Methods in Natural Language Processing (EMNLP). arXiv:2311.08466
Fu, Y., Bailis, P., Stoica, I., & Zhang, H. (2024). Break the Sequential Dependency of LLM Inference Using Lookahead Decoding. arXiv preprint. arXiv:2402.02057

KV-Cache Quantization: How 2–4 Bits Cut LLM Memory by 4–8 Without Killing Quality

Yuvraj singh Bhadoria — Sat, 29 Aug 2026 15:24:17 +0000

KV-Cache Quantization: I Ran the Experiments So You Don't Have To

KV-cache quantization is one of the highest-leverage knobs in modern LLM inference. As context windows
stretch past 128K tokens and batch sizes climb, the Key-Value cache — not the model weights — becomes the
dominant memory consumer, quietly turning serving into a memory-bound problem. Storing that cache in 4-bit
or 2-bit cuts its footprint 4–8× with minimal quality loss, but naive rounding collapses attention rather
than compressing it.

In this deep-dive I don't take that on faith. I load a real model (Qwen2.5-0.5B), capture its actual
KV-cache activations, quantize them four different ways, and measure exactly how each scheme distorts
attention. Then I compare the three production schemes — KIVI, KVQuant, and GEAR — and show
how to wire quantization into a real serving stack (vLLM / HuggingFace).

The Goal

By the end you should be able to answer:

Why the KV cache — not the weights — is the memory wall on long context.
The trap I measured: why naive round() on the cache is a silent quality killer.
Which scheme to reach for: KIVI, KVQuant, or GEAR.
How to wire it into a real serving stack (vLLM / HuggingFace).
How to compute the cache size for your model in 30 lines of Python.

Hardware & Experiment Setup

Before the theory, the actual rig the numbers below come from:

Model: Qwen/Qwen2.5-0.5B-Instruct — 0.5B params, 24 layers, GQA (14 query / 2 KV heads), head_dim 64.
Machine: a CPU-only laptop (no GPU) — the point is you can reproduce this without a cluster.
Prompt: one 77-token real sentence, run through the model's forward pass.
What we captured: the actual Key and Value projections at every layer.
What we measured: relative attention-score error ‖A − Â‖ / ‖A‖ after quantizing the KV cache.

No synthetic tensors, no "trust me" — the activations are the model's own.

Component 1 — The cache is a tax that scales with tokens

Every transformer layer keeps the Key and Value vectors of every token it has seen, so attention doesn't
have to recompute them. That storage is:

KV_bytes = 2 · n_layers · n_kv_heads · head_dim · bytes_per_elem · seq_len · batch

For Llama-3-8B (32 layers, 8 KV heads, head_dim 128) that's 128 KiB per token at fp16, and it
grows linearly with both sequence length and batch. The picture below is the whole problem in one frame.

Context (tokens)	fp16	8-bit	4-bit	2-bit
4,096	0.54 GB	0.27 GB	0.13 GB	0.07 GB
8,192	1.07 GB	0.54 GB	0.27 GB	0.13 GB
32,768	4.29 GB	2.15 GB	1.07 GB	0.54 GB
131,072	17.18 GB	8.59 GB	4.29 GB	2.15 GB

Batch 64 requests at 32K: the fp16 cache alone is ~275 GB — past any single GPU. Even 4-bit leaves
~69 GB. This is why long-context serving is memory-bound: the GPU's compute cores sit idle while the
KV cache is shuffled from main memory into SRAM for every generated token (KIVI, ICML 2024).

[!WARNING]
A smaller cache ≠ automatic speedup. The win is throughput via bigger batches, and only if your
serving stack (paged memory, fused dequant-matmul, correct calibration) supports KV quantization end to
end. Quantize the tensor but not the memory manager and you get neither.

Component 2 — I measured the trap myself

The obvious move is uniform, per-token quantization. That's also the mistake. Here is the experiment, end
to end:

I fake-quantized the captured KV cache at 2-bit and 4-bit under four schemes, recomputed attention
scores, and measured the error. The bar chart is the measured result:

Scheme (bits)	2-bit attn error	4-bit attn error
uniform (K,V per-token) — the naive one	0.786	0.401
uniform (K,V per-channel)	0.623	0.215
K per-channel, V per-token (KIVI)	0.623	0.215
K per-token, V per-channel (wrong)	0.786	0.401

Two things jump out, and they match the literature:

Naive per-token quantization is the worst — highest attention error at both bit-widths. Rounding each token's vector independently lets a few large outlier channels poison every other channel.
Quantizing keys per-channel fixes most of it. Per-channel key (0.623) cuts the error vs per-token key (0.786) — about 1.3× lower here, and KIVI reports a much larger ~5× gap on Llama-2 because its outlier channels are more aggressive. Direction is identical; magnitude scales with the model.

Why? Look at the distribution. Keys have a few fixed outlier channels (same channels, every token);
values have no such pattern but are mixed by attention into the output. So keys want per-channel
quantization, values want per-token. Uniform quantization ignores that and pays for it.

The whole asymmetric scheme in a dozen lines:

def fake_quant(x, bits, dim):
    qmax = 2 ** bits - 1
    scale = x.abs().amax(dim=dim, keepdim=True).clamp_min(1e-9) / qmax
    return torch.round(x / scale).clamp(-qmax - 1, qmax) * scale

# KIVI: keys per-channel (dim = sequence), values per-token (dim = head_dim)
K_q = fake_quant(K, 2, dim=2)   # per-channel over tokens
V_q = fake_quant(V, 2, dim=3)   # per-token over head_dim

Component 3 — Three schemes that actually ship

Three papers define the frontier. They agree on the outlier structure and diverge on how hard they push
the bit-width.

Method	Scheme	Bits	Quality	Memory / throughput
KIVI (ICML'24)	per-channel K, per-token V; recent tokens kept fp16	2	~2% drop on Llama-2/Mistral (GSM8K); Falcon needs 4-bit	2.6× peak mem, 4× larger batch, 2.35–3.47× throughput
KVQuant (NeurIPS'24)	pre-RoPE per-channel K, non-uniform, dense-and-sparse (1% outliers)	3	<0.1 perplexity drop (WikiText-2, C4)	4.8× compression; LLaMA-7B at 1M ctx on 1× A100, 10M on 8× GPU; ~1.7× matvec speedup
GEAR (ICML'24)	quantization + low-rank error + sparse outliers	2	near-lossless; up to 24.4% over SOTA at 2-bit	2.39× peak mem, 2.1–5.07× throughput

KIVI is the pragmatic default. Tuning-free, plug-and-play, ships as a HuggingFace wrapper. The full-precision window for the most recent tokens rescues hard reasoning — without it, fake 2-bit on GSM8K craters.
KVQuant is the "go long" play. Quantizing before RoPE (which otherwise mixes outlier channels) plus sensitivity-weighted non-uniform codebooks and a 1% sparse outlier store reaches 3-bit with sub-0.1 perplexity loss and turns context length into a tunable dial.
GEAR is the "go low" play. Instead of fighting residuals it models them: a low-rank matrix recovers the coherent part of the quantization error, a sparse matrix catches the rest. It layers on top of any base quantizer and is the only one of the three that stays near-lossless at 2-bit on complex generation.

[!WARNING]
These numbers are method-specific. KIVI reports that Falcon-7B (multi-query attention, a single KV
head) needs 4-bit, not 2-bit — MQA is already so compressed there's no redundancy left to trade. If
your model uses MQA, don't assume 2-bit works.

Component 4 — Serving architecture & manifests

This isn't a research curiosity you bolt on by hand — it slots into the paged-memory path every modern
server already runs:

For the long-context frontier, KVQuant and GEAR ship CUDA kernels that fuse dequant into the matmul —
without that fusion, the quantization overhead eats the memory win. In practice you rarely write the kernel
yourself; two drop-in paths:

vLLM — fp8 KV cache (the safe default today):

# vllm serve config
model: meta-llama/Llama-3-8B-Instruct
kv_cache_dtype: fp8_e5m2        # 2x smaller KV, near-lossless
max_model_len: 131072
gpu_memory_utilization: 0.85

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-8B-Instruct \
  --kv-cache-dtype fp8_e5m2 --max-model-len 131072

HuggingFace + KIVI (research-grade 2-bit):

from models.llama_kivi import LlamaForCausalLM_KIVI
config.k_bits = 2; config.v_bits = 2          # 2-bit KV cache
config.group_size = 64
config.residual_length = 64                    # recent tokens kept fp16
model = LlamaForCausalLM_KIVI.from_pretrained("meta-llama/Llama-2-7b-hf", config=config)

Where it still breaks

Small model + long context = danger zone. Perplexity degradation grows as you shrink the model and lengthen the prompt; the headroom that hides quantization error in a 70B model doesn't exist in a 7B one.
Reasoning tasks are unforgiving. GSM8K-style math is where fake quantization fails hardest; the recent-token window (KIVI) or low-rank residual (GEAR) is not optional.
Throughput only follows if the kernel exists. KVQuant and GEAR ship custom CUDA kernels; without fused dequant-matmul the overhead eats the memory win. Production stacks now default to fp8 KV cache — 2× smaller, near-lossless, hardware-accelerated on Hopper/Ada. 4-bit and below still need the research kernels.
Calibration matters at the low end. KVQuant calibrates key scales offline; skip it and 3-bit quality drops sharply.

The memory math (verified)

The tables above aren't benchmark results — they're deterministic accounting that tells you whether the
request fits at all. Here's the whole computation, in 30 lines:

def kv_bytes(seq_len, batch, bits):
    # 2 (K and V) * layers * kv_heads * head_dim * (bits/8) * seq * batch
    return 2 * 32 * 8 * 128 * (bits / 8) * seq_len * batch

print(kv_bytes(131072, 1, 16) / 1e9)   # 17.18 GB  (fp16, 1 req, 128K)
print(kv_bytes(131072, 1, 4)  / 1e9)   # 4.29 GB   (4-bit)

Run it and you get the figures in this post. That's the question KV-cache quantization exists to answer:
does this request fit?

Takeaway

KV-cache quantization is the rare optimization that's both mathematically simple and empirically subtle.
The win is real and large — 4–8× memory, 2–5× throughput — but it's earned by respecting the cache's
outlier structure, not by rounding harder. I measured it directly on a real 0.5B model: naive per-token
quantization distorts attention ~1.3× more than the per-channel key scheme, and the gap is wider on bigger
models. For most teams the right move today is fp8 in the serving stack; for the long-context
frontier, KVQuant and GEAR show 2–3-bit caches are deployable, not research curiosities.

I started this wanting to understand why my 128K requests kept OOMing, and ended up rebuilding the cache
path from the outlier math up. Next up: I'll break down speculative decoding the same way — where the
famous 2–3× claim actually holds, and where it quietly loses.

References

Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (ICML 2024). arXiv:2402.02750 · github.com/jy-yuan/KIVI
Hooper et al., KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (NeurIPS 2024). arXiv:2401.18079 · github.com/SqueezeAILab/KVQuant
Kang et al., GEAR: An Efficient Error Reduction Framework for KV Cache Compression in LLM Inference (ICML 2024). arXiv:2403.05527 · github.com/opengear-project/GEAR
Kwon et al., vLLM / PagedAttention (SOSP 2023) — the system layer KV quantization slots into. arXiv:2309.06180
Experiment code & raw results: real_experiment.py, assets/real_results.json in this project.

I Wrote the Three Attention Kernels That Run LLM Inference — in Portable Triton, on a Free Colab T4

Yuvraj singh Bhadoria — Sat, 29 Aug 2026 15:05:17 +0000

I Wrote the Three Attention Kernels That Run LLM Inference — in Portable Triton, on a Free Colab T4

When people talk about LLM inference kernels, they usually go straight to the hardware-specific fireworks: Hopper TMA, FP8, warp-specialization, FlashAttention-3. But if you actually want to understand why inference is fast or slow on ordinary hardware, you have to start with the foundation: the three attention kernels that every serving engine secretly runs — fused prefill, split-KV decode, and paged KV-cache attention.

I wrote all three from scratch in portable Triton (no CUDA C++), ran them on a free Colab NVIDIA T4, and validated every one against PyTorch SDPA. This is the engineering post: what each kernel does, why it matters, where it wins, and the traps I hit.

What is attention, really?

Inference attention computes softmax(QKᵀ / √d) V. The catch is the middle matrix QKᵀ:

At prefill (processing a prompt), Q is large, so this is compute-bound — but the naive version materializes the full [B, H, N, N] score matrix, which is O(N²) in memory.
At decode (generating one token), Q is a single vector (Q=1), so the matmul is tiny and the cost is dominated by streaming the KV cache from HBM — it is memory-bound, not compute-bound.

That single distinction — compute-bound prefill vs memory-bound decode — is the whole reason serving engines use different kernels for each phase. Ignore it and your "optimization" optimizes the wrong thing.

The memory model (why naive attention dies)

Your GPU has fast compute but it does not know how to stage your exact workload. With custom/Triton kernels you control data movement:

Off-chip (HBM, 16 GB on T4): where Q, K, V and the KV cache live. Huge but slow.
On-chip (shared memory / registers): tiny but fast. You want to pull K, V once and reuse them across the Q dimension.

For prefill, the naive implementation pulls this wrong: it builds the full N×N score matrix in HBM. That is the trap.

My hardware: a free Colab T4

I have no H100. This runs on an NVIDIA T4 (sm75, 16 GB) — the free Colab tier. Ada/Hopper tricks (TMA, FP8) are intentionally out of scope; the point is a portable kernel any laptop-class GPU can run. fp16 throughout.

The trap: naive attention OOMs at 8K

I benchmarked prefill peak memory exactly (peak-memory is unaffected by Colab's compute throttling, so these numbers are trustworthy):

N_CTX	Naive attention	Triton fused
4K	4.46 GB	0.44 GB (10.2× less)
8K	OOM	0.57 GB (fits)

The naive materialized [B,H,N,N] matrix is O(N²). At 8K it blows past 16 GB and dies. Triton's fused kernel is O(N·D) — it never materializes the scores, it streams K,V and reduces online (FlashAttention-2 style). Same math, 10× less memory, and it keeps running where naive crashes.

[!WARNING]
On a free T4, PyTorch SDPA dispatches to vendor-tuned FlashAttention-2 / cuDNN. Do not expect a raw-speed win over SDPA. The honest, demonstrable win here is memory and portability — not latency. Measuring latency against SDPA and calling it a loss misses the point.

The three kernels (evolution, with why)

1. Fused prefill attention (`flash_attention.py`)

The baseline. Online softmax, autotuned, FA2-style. Processes the whole prompt in one pass. Its job is to prove the correctness methodology: every output is checked against F.scaled_dot_product_attention with allclose(atol=1e-2, rtol=1e-2) in fp16, including non-power-of-two and GQA shapes.

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "num_warps": 4}, num_stages=2),
        triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "num_warps": 4}, num_stages=2),
        triton.Config({"BLOCK_M": 32, "BLOCK_N": 32, "num_warps": 4}, num_stages=2),
    ],
    key=["N_CTX"],
)
@triton.jit
def _flash_attention_fwd(Q, K, V, sm_scale, Out, ...):
    start_m = tl.program_id(0)   # one tile of query rows per program
    # online softmax: stream K,V, reduce in registers -- never materialize NxN

2. Split-KV flash-decoding (`flash_decoding.py`)

The decode kernel. With Q=1 there is no query-dimension parallelism, so the only way to parallelize decode is to split the KV dimension across SMs — exactly what vLLM does for long-context decode. Each SM reduces its slice of K,V, then a final cross-SM reduction gives the output.

Why it matters: the tiny Q=1 matmul is memory-bound, so you hide HBM latency by having many SMs pull different KV blocks at once.

split = tl.program_id(2)                  # KV dimension split across SMs
kv_per_split = tl.cdiv(tl.cdiv(N_CTX, BLOCK_N), NUM_SPLITS) * BLOCK_N
start_n = split * kv_per_split            # this SM owns its slice of K,V

# ... after each SM reduces its slice, a final cross-SM combine:
lse_max = partial_lse.max(-1, keepdim=True).values
w = torch.exp(partial_lse - lse_max)
out = (partial_out * w.unsqueeze(-1)).sum(2) / w.sum(-1, keepdim=True)

3. Paged KV-cache attention (`paged_attention.py`)

The inference-systems piece. Instead of one contiguous KV tensor, it reads through a block table (GQA-aware) — the exact memory layout vLLM uses. This is a capability SDPA literally cannot express: PyTorch's built-in attention assumes a dense KV tensor; real serving does not have one.

# paged_attention.py -- read K,V through a block table, not one dense tensor
page = tl.load(BlockTable + b * stride_bb + (p // PAGE_SIZE) * stride_bm)
k = tl.load(KV_Cache + page * stride_kn + off_in_page * stride_kps)

Benchmark method

Shapes: B=4, H=16, D=64, Q=1, fp16.
Decode sweep: N_CTX ∈ {1024, 2048, 4096, 8192, 16384}, num_splits ∈ {1, 8, 16, 32, 64, 128}.
Timing: 20 warmup + 100 measured CUDA-event iterations (async events, like a real benchmark).
Correctness: allclose vs SDPA on every shape before any timing runs.

Results

Best num_splits per context (lowest kernel latency):

N_CTX	best num_splits	Triton (ms)	SDPA (ms)	Triton / SDPA
1024	32	0.405	0.249	1.62×
2048	32	0.297	0.379	0.78×
4096	16	0.528	0.393	1.34×
8192	16	1.018	0.845	1.21×
16384	16	2.012	1.737	1.16×

Step-by-step optimization journey

From single-split to split-KV (the big one). At N=1024, num_splits=1 stalled at 1.14 ms; moving to num_splits=32 dropped it to 0.40 ms — because splitting the KV dimension is the only source of sequence parallelism when Q=1.
Prefill memory: fused Triton went from 4.46 GB → 0.44 GB at 4K (10.2×), and fit at 8K where naive OOMs.
Versus SDPA: within ~1.2–1.6× at 4K–16K, and faster at 2K (0.78×) in this run. SDPA stays the vendor-tuned reference; the benchmark characterizes split-KV scaling, not a speed claim.

Key Tradeoffs & Technical Takeaways

1. More splits is NOT monotonic — launch overhead beats memory reuse.
num_splits=128 is the worst configuration at every context length. The sweet spot (16–32) does not even grow steadily with N. Past the sweet spot, cross-SM launch + reduction overhead dominates the HBM savings. Takeaway: split-KV helps until reduction bookkeeping costs more than the memory you saved.

2. Memory win vs speed win are different wins.
On T4 we beat naive on memory 10×, and we match/beat SDPA on latency at scale — but SDPA still wins at small N because it dispatches to cuDNN/FA2. Takeaway: on constrained hardware, optimize the metric that's actually broken (here, decode memory pressure), not the one the vendor already solved.

3. Portability is a deliberate tradeoff.
No TMA, no FP8, no warp-spec — those need Hopper+. By targeting sm75 Triton, the kernel runs on a free Colab T4 and any datacenter GPU, at the cost of leaving Hopper bandwidth on the table. Takeaway: write the portable kernel first; specialize only when the hardware is guaranteed.

Conclusion

Triton is a new thing I explored in this project, and it's a bit low-level — but watching three kernels go from "naive OOM" to "validated against SDPA on a free T4" made the inference path click in a way no framework toggle ever did. Building the paged kernel especially — the one SDPA can't express — is what showed me why serving engines look the way they do.

Benchmarking on shared Colab GPUs has been tough lately (numbers wander with contention), so I report medians across 2–3 runs. Still, the memory numbers and the split-KV shape are solid.

Thanks for reading till here. Repo: triton-attention-kernels.

DEV Community: Yuvraj singh Bhadoria

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Executive Summary

Visual Overview

1. Model Setup & Architecture

2. The Serial Memory Bottleneck

3. The Speculative Sampling Mechanism

Sampling & Rejection Implementation

4. The Speculative-Decoding Family

Production Framework Integrations

Hugging Face Transformers (assisted_generation)

vLLM Native Scheduler Integration

5. Benchmarking & Empirical Performance

Mathematical Speedup Condition ($\alpha > c$)

Measured Experimental Results

System Failure Modes

Conclusion

References

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Executive Summary

Visual Overview

1. Model Setup & Architecture

2. The Serial Memory Bottleneck

3. The Speculative Sampling Mechanism

Sampling & Rejection Implementation

4. The Speculative-Decoding Family

Production Framework Integrations

Hugging Face Transformers (assisted_generation)

vLLM Native Scheduler Integration

5. Benchmarking & Empirical Performance

Mathematical Speedup Condition ($\alpha > c$)

Measured Experimental Results

System Failure Modes

Conclusion

References

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Demystifying Speculative Decoding: From Architecture to Production Bottlenecks

Executive Summary

Visual Overview

1. Model Setup & Architecture

Experimental Setup Parameters

2. The Serial Memory Bottleneck

3. The Speculative Sampling Mechanism

Sampling & Rejection Implementation

4. The Speculative-Decoding Family

Production Framework Integrations

Hugging Face Transformers (assisted_generation)

vLLM Native Scheduler Integration

5. Benchmarking & Empirical Performance

Mathematical Speedup Condition ($\alpha > c$)

Measured Experimental Results

System Failure Modes

Conclusion

References

KV-Cache Quantization: How 2–4 Bits Cut LLM Memory by 4–8 Without Killing Quality

KV-Cache Quantization: I Ran the Experiments So You Don't Have To

The Goal

Hardware & Experiment Setup

Component 1 — The cache is a tax that scales with tokens

Component 2 — I measured the trap myself

Component 3 — Three schemes that actually ship

Component 4 — Serving architecture & manifests

Where it still breaks

The memory math (verified)

Takeaway

References

I Wrote the Three Attention Kernels That Run LLM Inference — in Portable Triton, on a Free Colab T4

I Wrote the Three Attention Kernels That Run LLM Inference — in Portable Triton, on a Free Colab T4

What is attention, really?

The memory model (why naive attention dies)

My hardware: a free Colab T4

The trap: naive attention OOMs at 8K

The three kernels (evolution, with why)

1. Fused prefill attention (flash_attention.py)

2. Split-KV flash-decoding (flash_decoding.py)

3. Paged KV-cache attention (paged_attention.py)

Benchmark method

Results

Hugging Face Transformers (`assisted_generation`)

Hugging Face Transformers (`assisted_generation`)

Hugging Face Transformers (`assisted_generation`)

1. Fused prefill attention (`flash_attention.py`)

2. Split-KV flash-decoding (`flash_decoding.py`)

3. Paged KV-cache attention (`paged_attention.py`)