Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Speculative Decoding: Fast LLM Inference Without Quality Loss › Parallel Verification with Causal Masking

Michael BrenndoerferJanuary 16, 20263 min read

Part 42 of 67Inference Optimization

Chapter 323Speculative Decoding: Fast LLM Inference Without Quality Loss › Parallel Verification with Causal Masking

Reading section 4 of 14 · Parallel Verification with Causal Masking · 3 minSection 4 of 14 · 3 min

Section 4 of 14 · 3 min read

Speculative Decoding: Fast LLM Inference Without Quality Loss › Parallel Verification with Causal MaskingOne section · Sections in this chapter4 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

How does the target model verify multiple tokens in one forward pass? The answer lies in the causal attention mask we covered in Part XIII, Chapter 4. This mechanism, originally designed to enable efficient training on entire sequences, turns out to be exactly what we need for parallel verification during inference.

When we feed the sequence $[\text{prompt}, \text{draft}_1, \text{draft}_2, \text{draft}_3, \text{draft}_4]$ to the target model, the causal mask ensures each position only attends to preceding tokens. This creates a natural structure where each position independently computes the probability of its token given everything that came before it, even though all positions are processed simultaneously in one forward pass.

At position $\text{draft}_1$ , the model computes $p(\text{draft}_1 \mid \text{prompt})$ . At position $\text{draft}_2$ , it computes $p(\text{draft}_2 \mid \text{prompt}, \text{draft}_1)$ . Each position's output gives us the probability the target model assigns to the draft token at that position, conditioned on everything before it. These computations happen simultaneously in a single forward pass, not sequentially as they would in standard autoregressive generation. The causal mask handles the dependency structure automatically.

The key insight here is that the causal mask is doing double duty. During training, it prevents the model from "cheating" by looking at future tokens. During speculative decoding verification, it enables something remarkable: the model independently verifies each draft token in its correct conditional context without needing to run $K$ separate forward passes. The same architectural decision that enables efficient training also enables efficient verification. This is not a coincidence: both applications require computing $p(x_t \mid x_{<t})$ for all positions $t$ simultaneously.

In[8]:

Code

import torch
import torch.nn.functional as F


def parallel_verification(target_model, input_ids, draft_tokens):
    """
    Verify K draft tokens in one forward pass.
    Returns probability of each draft token under target model.
    """
    # Concatenate input with draft tokens
    full_sequence = torch.cat([input_ids, draft_tokens], dim=-1)

    # Single forward pass with causal masking
    with torch.no_grad():
        logits = target_model(full_sequence).logits

    # Extract logits at positions where we predict draft tokens
    # Position i predicts token i+1, so we need positions just before each draft token
    num_input = input_ids.shape[-1]
    num_draft = draft_tokens.shape[-1]

    verification_logits = logits[
        :, num_input - 1 : num_input + num_draft - 1, :
    ]

    # Get probability of each draft token
    probs = F.softmax(verification_logits, dim=-1)

    draft_probs = []
    for i in range(num_draft):
        token_id = draft_tokens[0, i].item()
        draft_probs.append(probs[0, i, token_id].item())

    return draft_probs, probs

This parallel verification is what makes speculative decoding efficient. Without it, verifying $K$ tokens would require $K$ forward passes through the target model, eliminating any speedup. With it, we amortize the memory bandwidth cost across $K$ potential tokens. The causal mask ensures that even though we process all positions simultaneously, each position's output depends only on previous positions, maintaining the autoregressive property that makes language model outputs coherent and consistent.

The target model also produces a prediction at the position just after the last draft token. This is a "bonus" token that comes for free from the forward pass. When all $K$ draft tokens are accepted, we can also accept this bonus token, giving us $K+1$ tokens from a single round. This bonus token has zero marginal cost and is always distributed exactly according to the target model, since it is sampled directly from the target's output distribution without any draft or acceptance logic.