dflash-block-diffusion-speculative-decoding

name	dflash-block-diffusion-speculative-decoding
title	DFlash: Block Diffusion for Flash Speculative Decoding
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2602.06036
keywords	["Speculative Decoding","Diffusion Models","Parallel Generation","Inference Acceleration","Target Model Conditioning"]
description	Accelerate LLM inference 6x by using block diffusion for parallel token drafting with tight coupling to the target model's hidden representations, achieving higher speedups than existing speculative methods without quality loss.

DFlash: Block Diffusion for Flash Speculative Decoding

Problem Context

Autoregressive LLMs generate tokens sequentially, creating inference bottlenecks. Speculative decoding helps but is limited by reliance on autoregressive drafting—inherently sequential processes that cap speedups around 2-3x. Diffusion models enable parallel generation but typically underperform in quality. DFlash bridges this gap: parallel drafting with quality matching the target model.

Core Concept

DFlash combines [parallel block diffusion, target model conditioning, KV injection] to enable fast high-quality token drafting. The insight is "the target knows best"—leveraging frozen target model representations substantially improves draft quality without needing massive draft models.

Architecture Overview

• Draft model: Small diffusion model generating multiple tokens in parallel
• Target conditioning: Extract hidden representations from frozen target model; inject into draft KV cache per layer
• KV injection design: Directly insert target context features into cache rather than standard fusion
• Training: Random anchor sampling for block initialization; loss weighting prioritizing early positions
• Verification: Standard acceptance checking; discard rejected tokens

Implementation

Step 1: Extract target model representations for conditioning

Freeze the target LLM and extract intermediate layer hidden states to guide draft generation. Cache these for efficiency.

# Extract conditioning from target model
def extract_target_conditioning(target_model, input_ids, layer_indices):
    """
    Forward target model; extract hidden states at specified layers.
    Returns conditioning features for each layer.
    """
    conditioning = {}
    with torch.no_grad():
        output = target_model(
            input_ids,
            output_hidden_states=True,
            return_dict=True
        )
        for layer_idx in layer_indices:
            # Extract and compress hidden state
            h = output.hidden_states[layer_idx]
            # Project to draft model's key-value dimension
            conditioning[layer_idx] = project_to_kv(h)
     conditioning