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.
Instrucciones de origen · Vista previa de solo lectura
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 modeldefextract_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
return
Step 2: Design KV injection mechanism
Inject target model context directly into draft model's Key-Value cache at every decoder layer, replacing standard feature fusion.
# KV injection into draft model cachedefinject_target_conditioning_to_kv(
draft_kv_cache, target_conditioning, layer_idx,
injection_strength=0.5):
"""
Inject target model hidden states directly into KV cache.
Enables acceptance scaling with depth.
"""
k, v = draft_kv_cache[layer_idx]
# Target features as additional context in value
target_feat = target_conditioning[layer_idx]
# Blend target features with draft KV
v_augmented = v + injection_strength * target_feat
draft_kv_cache[layer_idx] = (k, v_augmented)
return draft_kv_cache
Step 3: Implement block diffusion with anchor sampling
Initialize diffusion process with random anchor tokens at block boundaries. Use scheduled denoising to generate multiple tokens in parallel.
Step 4: Apply loss weighting favoring early positions
During training, weight loss to prioritize early tokens in block, which are most predictable.
# Loss weighting for block diffusion trainingdefweighted_diffusion_loss(
predictions, targets, position_weights=None):
"""
Weight tokens early in block more heavily.
Early tokens are easier to predict; later tokens need guidance.
"""if position_weights isNone:
# Default: exponential decay favoring early positions
block_size = targets.shape[-1]
position_weights = torch.exp(
-torch.arange(block_size, dtype=torch.float32) * 0.3
)
position_weights = position_weights / position_weights.sum()
# Per-position loss weighting
loss = F.cross_entropy(
predictions.reshape(-1, predictions.shape[-1]),
targets.reshape(-1),
reduction='none'
)
loss = loss.reshape(targets.shape)
# Apply position weighting
weighted_loss = (loss * position_weights).mean()
return weighted_loss
Step 5: Integrate speculative decoding with verification
Generate drafts in parallel blocks; verify and accept valid tokens; fall back to target model only when needed.
# Speculative decoding with DFlashdefflash_speculative_decode(
target_model, draft_model, input_ids,
max_new_tokens=100, block_size=4):
"""
Generate tokens using DFlash speculative decoding.
"""
generated = input_ids.clone()
target_conditioning = extract_target_conditioning(
target_model, input_ids, layer_indices=[12, 24]
)
while generated.shape[1] < input_ids.shape[1] + max_new_tokens:
# Draft: generate block in parallel
draft_block = block_diffusion_generate(
draft_model, generated, block_size=block_size
)
# Verify against target modelwith torch.no_grad():
target_logits = target_model(
torch.cat([generated, draft_block], dim=1)
).logits
# Acceptance checking: compare probabilities
accepted = verify_tokens(
draft_block,
target_logits[:, -block_size:],
temperature=1.0
)
# Add verified tokens; fall back to target if none accepted
num_accepted = accepted.sum().item()
if num_accepted > 0:
generated = torch.cat([
generated,
draft_block[:, :num_accepted]
], dim=1)
else:
# Fall back to target model
next_token = target_model(generated).logits[:, -1].argmax(dim=-1)
generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)
return generated
Practical Guidance
When to use: High-throughput inference settings (serving, batch processing) where latency is critical. Less beneficial for interactive single-token-at-a-time scenarios.
Hyperparameters:
Block size (2-8): larger blocks → more parallelism but lower acceptance rates; 4 is sweet spot
Diffusion steps (10-20): more steps → better quality but slower; 15 typical
Injection strength (0.3-0.7): balance target guidance with draft autonomy
Layer indices for conditioning: use middle-to-later layers (12-24 in 32-layer models)
Common pitfalls:
Target model hidden state dimension must align with draft model; use projection layers
Block boundaries matter: random initialization worse than structured patterns; try position-aware initialization
Acceptance rate drops if draft model too weak; consider larger draft (7B-13B for 70B target)
Memory overhead from storing target conditioning; cache selectively for long sequences
Scaling: Speedup scales with block size and acceptance rate. Typical 4-6x speedup on standard benchmarks. Works best with strong target models and sufficient compute for parallel drafting.
Reference
Paper: https://arxiv.org/abs/2602.06036
Code: Available at author's repository
Related work: Speculative decoding, EAGLE, diffusion-based generation, fast inference
