Speculation as Compression

This repo explores a new direction on top of DFlash: use speculative decoding not only for faster generation, but also as a signal for online context compression.

The current codebase has two active tracks:

1. DFlash reproduction and diagnostics: run Qwen3 target models with DFlash block-diffusion draft models, measure acceptance length and speedup, and inspect draft-vs-target behavior.
2. Accepted-span compression prototype: after speculative verification succeeds, compress verified spans into learned latent units and continue generation on a compressed context. This path is approximate, not logit-exact.

The older "Draft-Guided Decoding / TTT deep readout" idea is still useful background, but the current implementation focus is DFlash plus compression.

Repository Layout

src/dg_ttt/model/dflash.py          # Qwen3 DFlash draft model and speculative generation
src/dg_ttt/model/compression.py     # accepted-span latent compression prototype
src/dg_ttt/model/utils.py           # datasets, sampling, attention backend helpers
scripts/benchmark.py                # DFlash benchmark, Transformers backend
scripts/diagnostic.py               # draft-vs-target token diagnostic
scripts/train_compressor.py         # offline one-latent compressor training prototype
scripts/benchmark_compressed.py     # compressed DFlash benchmark prototype
ref/dflash/                         # DFlash paper package and reference src
record/                             # phase-0 report and figures

ref/dflash/src contains the upstream reference implementation:

• model.py: PyTorch/Transformers DFlash model and dflash_generate
• model_mlx.py: MLX streaming implementation with hidden-state hooks and cache rollback
• benchmark.py: multi-backend benchmark driver

Our src/dg_ttt/model/dflash.py is derived from the PyTorch reference path.

Environment

Use the repo virtual environment and local package path:

export PYTHONPATH=src
export HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface

The project is managed by uv:

/home/ubuntu/.nix-profile/bin/uv run python --version

Installed stack used for the current smoke tests:

torch 2.11.0+cu130
transformers 5.6.0
flash_attn 2.8.3

Local Model Cache

The Qwen and DFlash models have been downloaded into the repo-local Hugging Face cache:

.cache/huggingface/hub/models--Qwen--Qwen3-4B                 7.6G
.cache/huggingface/hub/models--Qwen--Qwen3-8B                 16G
.cache/huggingface/hub/models--z-lab--Qwen3-4B-DFlash-b16     1.1G
.cache/huggingface/hub/models--z-lab--Qwen3-8B-DFlash-b16     2.0G

If the cache is missing, download models with:

mkdir -p .cache/huggingface/hub

env -u ALL_PROXY -u all_proxy .venv/bin/hf download Qwen/Qwen3-4B \
  --cache-dir .cache/huggingface/hub
env -u ALL_PROXY -u all_proxy .venv/bin/hf download Qwen/Qwen3-8B \
  --cache-dir .cache/huggingface/hub
env -u ALL_PROXY -u all_proxy .venv/bin/hf download z-lab/Qwen3-4B-DFlash-b16 \
  --cache-dir .cache/huggingface/hub
env -u ALL_PROXY -u all_proxy .venv/bin/hf download z-lab/Qwen3-8B-DFlash-b16 \
  --cache-dir .cache/huggingface/hub

The env -u ALL_PROXY -u all_proxy prefix avoids the local httpx[socks] dependency issue when a SOCKS proxy is configured.

Attention Backend

The local transformers 5.6.0 flash-attention wrapper has a compatibility bug for Qwen3:

s_aux=s_aux.to(query.dtype)

Qwen3 does not pass s_aux, so unpatched flash_attention_2 raises:

AttributeError: 'NoneType' object has no attribute 'to'

This repo patches the registered Transformers flash-attention function at startup via:

patch_flash_attention_s_aux()
resolve_attn_implementation(...)

The patch is local to the running Python process; it does not edit site-packages.

All project scripts accept:

--attn-implementation {auto,sdpa,eager,flash_attention_2}

Use flash_attention_2 for performance after the patch:

--attn-implementation flash_attention_2

Use sdpa as a conservative fallback.

DFlash Benchmark

Small 4B smoke:

HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface \
PYTHONPATH=src \
/home/ubuntu/.nix-profile/bin/uv run python scripts/benchmark.py \
  --model-name-or-path Qwen/Qwen3-4B \
  --draft-name-or-path z-lab/Qwen3-4B-DFlash-b16 \
  --dataset gsm8k \
  --max-samples 1 \
  --max-new-tokens 32 \
  --temperature 0.0 \
  --attn-implementation flash_attention_2

Observed result:

Decoding speedup: 4.84
Average Acceptance length: 5.50

Larger 4B smoke:

HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface \
PYTHONPATH=src \
/home/ubuntu/.nix-profile/bin/uv run python scripts/benchmark.py \
  --model-name-or-path Qwen/Qwen3-4B \
  --draft-name-or-path z-lab/Qwen3-4B-DFlash-b16 \
  --dataset gsm8k \
  --max-samples 8 \
  --max-new-tokens 256 \
  --temperature 0.0 \
  --attn-implementation flash_attention_2

Observed result:

Decoding speedup: 5.41
Average Acceptance length: 6.84
Acceptance length histogram:
['0.0%', '10.4%', '15.3%', '11.0%', '13.0%',
 '6.2%', '4.5%', '6.8%', '3.6%', '3.2%',
 '3.2%', '1.6%', '3.6%', '2.9%', '2.3%',
 '1.0%', '11.4%']

8B smoke candidate:

HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface \
PYTHONPATH=src \
/home/ubuntu/.nix-profile/bin/uv run python scripts/benchmark.py \
  --model-name-or-path Qwen/Qwen3-8B \
  --draft-name-or-path z-lab/Qwen3-8B-DFlash-b16 \
  --dataset gsm8k \
  --max-samples 4 \
  --max-new-tokens 256 \
  --temperature 0.0 \
  --attn-implementation flash_attention_2

How DFlash Works

DFlash uses speculative decoding with a block-diffusion draft model:

target prefill
  -> sample first token
  -> extract target hidden states from layers [1, 9, 17, 25, 33]
  -> draft predicts a block in one forward pass
  -> target verifies the block in one forward pass
  -> accept matching prefix plus one target bonus token

The draft model is a small Qwen3-style transformer. Its attention is modified so every draft layer attends to target context features:

Q = q_proj(draft block hidden)
K = concat(k_proj(target_hidden), k_proj(draft block hidden))
V = concat(v_proj(target_hidden), v_proj(draft block hidden))

The draft shares the target model's token embedding and lm_head. Speculative verification still comes from the target model, so standard DFlash decoding is lossless with respect to the target.

Diagnostics

Run draft-vs-target token diagnostics:

HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface \
PYTHONPATH=src \
/home/ubuntu/.nix-profile/bin/uv run python scripts/diagnostic.py \
  --model Qwen/Qwen3-4B \
  --draft z-lab/Qwen3-4B-DFlash-b16 \
  --dataset gsm8k \
  --mode mask \
  --max-samples 5 \
  --output results/phase0/diagnostic_gsm8k_mask.json \
  --attn-implementation flash_attention_2

The existing phase-0 report is in:

record/phase0_diagnostic_report.md

That report found that vanilla DFlash draft logits are not a reliable guided-decoding improvement over target logits on the small GSM8K diagnostic, although intermediate target layers do carry real signal.

Accepted-Span Compression Prototype

The compression path is experimental and approximate.

Current implementation:

• AcceptedSpanCompressor reads a verified span and emits one continuous latent.
• compressed_dflash_generate replaces eligible verified spans with one latent in the target cache.
• The output text is still the verified speculative output, but the future target context becomes compressed. Therefore this is no longer an exact acceleration of the original target model.

This code should be treated as a baseline, not a faithful C3 implementation. A more faithful C3-style variant would use a separate llm1 encoder, latent query placeholders, and a projector into the decoder hidden space.

Train a prototype compressor:

HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface \
PYTHONPATH=src \
/home/ubuntu/.nix-profile/bin/uv run python scripts/train_compressor.py \
  --model-name-or-path Qwen/Qwen3-4B \
  --dataset alpaca \
  --output-dir results/compressor_qwen3_4b \
  --attn-implementation flash_attention_2

Run compressed DFlash:

HF_HOME=/mnt/public2/ubuntu/flashSpeculation/.cache/huggingface \
PYTHONPATH=src \
/home/ubuntu/.nix-profile/bin/uv run python scripts/benchmark_compressed.py \
  --model-name-or-path Qwen/Qwen3-4B \
  --draft-name-or-path z-lab/Qwen3-4B-DFlash-b16 \
  --compressor results/compressor_qwen3_4b/compressor.pt \
  --dataset gsm8k \
  --attn-implementation flash_attention_2

References

• DFlash: block-diffusion speculative draft model with target hidden-state KV injection.
• ref/dflash: local DFlash paper package and reference source.
• ref/C3-Context-Cascade-Compression: reference for latent context compression; useful for the next compressor design, but not yet faithfully implemented here.