Turbo dflash

[giveen]

Hardware: RTX 5090 (32GB) · Intel Core Ultra 9 285K
Target model: Qwen3.6-35B-A3B-UD-Q4_K_S.gguf
Draft model: Qwen3.6-35B-A3B-DFlash-Q8_0.gguf
Prompt: "Write a quicksort algorithm in Python. Write code only."
Sampling: --temp 0 --top-k 1 --seed 42

---

Overview

Port DFlash speculative decoding (block-diffusion drafting with KV cache injection, from (#103)) onto feature/turboquant-kv-cache. DFlash uses a lightweight draft model to propose blocks of tokens by encoding target model hidden states and conditioning on noise latents, then validates against the target model.

Key changes

• Adds llama_model_dflash and llama_model_eagle3 model architectures (encoder/decoder graph builders, tensor loader, hparams)
• Registers --dflash flag in CLI/server, with speculative type auto-detection
• Wires ctx_other for target model embedding/output sharing in DFlash decoder
• Adds target layer extraction infrastructure (embeddings_layer_inp + t_layer_inp population) needed by the DFlash encoder
• Fixes output_reserve sizing for target contexts when draft flags are active (server_n_outputs_max now accounts for --dflash/--eagle3)
• Adds KV/tensor enum entries and template instantiations for DFlash/EAGLE3 metadata keys

---

Results

Baseline (no speculation)

Command	Prompt t/s	Generation t/s
q8_0 cache, 1024 ctx	252.5	237.4

LLAMA_SPEC_NO_THINK=1 timeout 120 ./build/bin/llama-cli \
  -m /mnt/storage/models/qwen3.6/35B/Qwen3.6-35B-A3B-UD-Q4_K_S.gguf \
  -p "Write a quicksort algorithm in Python. Write code only." \
  -n 256 -c 1024 --temp 0 --top-k 1 --seed 42 -ngl 99 \
  --cache-type-k q8_0 --cache-type-v q8_0 --reasoning off

DFlash On

Cache K	Cache V	Context	Prompt t/s	Generation t/s
q8_0	q8_0	1024	184.8	340.6
q8_0	q8_0	32k	155.9	317.3
q8_0	q8_0	64k	147.7	319.0
q8_0	turbo4	64k	146.5	316.0
turbo4	turbo4	64k	144.1	315.4
turbo4	q8_0	64k	144.5	396.2

LLAMA_SPEC_NO_THINK=1 timeout 120 ./build/bin/llama-cli \
  -m /mnt/storage/models/qwen3.6/35B/Qwen3.6-35B-A3B-UD-Q4_K_S.gguf \
  -md /mnt/storage/models/qwen3.6/35B/Qwen3.6-35B-A3B-DFlash-Q8_0.gguf \
  --dflash \
  -p "Write a quicksort algorithm in Python. Write code only." \
  -n 128 -c <CONTEXT> --spec-draft-n-max 16 \
  --temp 0 --top-k 1 --seed 42 -ngl 99 -ngld 99 \
  --cache-type-k <K> --cache-type-v <V> --reasoning off

---

Summary

• Best result: 396.2 t/s generation (DFlash, turbo4 K / q8_0 V, 64k context) vs 237.4 t/s baseline — 1.67x speedup
• DFlash maintains high throughput even at 64k context (319.0 t/s with q8_0 both)
• turbo4 KV cache works with DFlash and can improve generation speed in mixed configurations
• Acceptance rate on this model pair is ~7% (model training quality, not a code issue)

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, to resolve merge conflicts and to fix my coding mistakes.

[giveen]

Speculative Decoding (Draft Model Enabled)

./build/bin/llama-cli \
  -m /mnt/storage/models/qwen3.6/27B/Qwen3.6-27B-UD-Q5_K_XL.gguf \
  -md /mnt/storage/models/qwen3.6/27B/Qwen3.6-27B-DFlash-Q8_0.gguf \
  --dflash \
  -p "Write a quicksort algorithm in Python. Write code only." \
  -c 64000 \
  --spec-draft-n-max 16 \
  -ngl 99 \
  --cache-type-k q8_0 \
  --cache-type-v q8_0 \
  --reasoning off \
  --temp 0.1

[ Prompt: 100.9 t/s | Generation: 164.8 t/s ]

---

Standard Decoding (No Draft Model)

./build/bin/llama-cli \
  -m /mnt/storage/models/qwen3.6/27B/Qwen3.6-27B-UD-Q5_K_XL.gguf \
  -p "Write a quicksort algorithm in Python. Write code only." \
  -c 64000 \
  -ngl 99 \
  --cache-type-k q8_0 \
  --cache-type-v q8_0 \
  --reasoning off \
  --temp 0.1

[ Prompt: 203.3 t/s | Generation: 63.5 t/s ]

[giveen]

Okay got server up and running

Prompt Eval (Prefill) : 104.34 t/s (24 tokens in 230.03 ms)
Generation (Decode)   : 156.27 t/s (96 tokens in 614.32 ms)
Draft Acceptance Rate : 86.67% (91 accepted / 105 generated)

[giveen]

Gemma 4 Speculative Decoding Benchmarks

A performance comparison of Gemma 4 (26B and 31B parameter models, including QAT variants) running baseline vs. speculative decoding via DFlash.

Summary of Results

Model	Speculative (DFlash)	Prompt Eval (t/s)	Generation (t/s)	Speedup (Gen)
Gemma4-26B (Q5_K_S)	❌ Baseline	236.7	182.3	1.00x
Gemma4-26B (Q5_K_S)	DFlash (Q8_0)	142.8	312.4	1.71x
Gemma4-26B-QAT (Q4_K_XL)	❌ Baseline	229.7	215.1	1.00x
Gemma4-26B-QAT (Q4_K_XL)	DFlash (Q8_0)	132.1	255.9	1.19x
Gemma4-31B (Q5_K_XL)	❌ Baseline	385.4	55.7	1.00x
Gemma4-31B (Q5_K_XL)	DFlash (Q5_K)	139.0	310.5	5.57x
Gemma4-31B-QAT (Q4_K_XL)	❌ Baseline	487.8	66.8	1.00x
Gemma4-31B-QAT (Q4_K_XL)	DFlash (Q5_K)	142.7	308.9	4.62x

---

Detailed Benchmark Logs

1. Gemma4-26B

Baseline (Standard Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/26B/gemma-4-26B-A4B-it-UD-Q5_K_S.gguf
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 236.7 t/s
• Generation: 182.3 t/s

DFlash (Speculative Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/26B/gemma-4-26B-A4B-it-UD-Q5_K_S.gguf
-md /mnt/storage/models/gemma4/26B/gemma-4-26B-A4B-it-DFlash-Q8_0.gguf
--dflash
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
--spec-draft-n-max 16
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 142.8 t/s
• Generation: 312.4 t/s

---

2. Gemma4-26B-QAT

Baseline (Standard Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/26B/gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 229.7 t/s
• Generation: 215.1 t/s

DFlash (Speculative Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/26B/gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf
-md /mnt/storage/models/gemma4/26B/gemma-4-26B-A4B-it-DFlash-Q8_0.gguf
--dflash
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
--spec-draft-n-max 16
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 132.1 t/s
• Generation: 255.9 t/s

---

3. Gemma4-31B

Baseline (Standard Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/31B/gemma-4-31B-it-UD-Q5_K_XL.gguf
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 385.4 t/s
• Generation: 55.7 t/s

DFlash (Speculative Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/31B/gemma-4-31B-it-UD-Q5_K_XL.gguf
-md /mnt/storage/models/gemma4/31B/gemma-4-31B-it-DFlash-Q5_K.gguf
--dflash
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
--spec-draft-n-max 16
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 139.0 t/s
• Generation: 310.5 t/s

---

4. Gemma4-31B-QAT

Baseline (Standard Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/31B/gemma-4-31B-it-qat-UD-Q4_K_XL.gguf
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 487.8 t/s
• Generation: 66.8 t/s

DFlash (Speculative Decoding)

Command:
./build/bin/llama-cli
-m /mnt/storage/models/gemma4/31B/gemma-4-31B-it-qat-UD-Q4_K_XL.gguf
-md /mnt/storage/models/gemma4/31B/gemma-4-31B-it-DFlash-Q5_K.gguf
--dflash
-p "Write a quicksort algorithm in Python. Write code only."
-c 64000
--spec-draft-n-max 16
-ngl 99
--cache-type-k q8_0
--cache-type-v q8_0
--reasoning off
--temp 0.1

Metrics:

• Prompt: 142.7 t/s
• Generation: 308.9 t/s

[giveen]

Deferred KV Injection Optimization for DFLASH Speculative Decoding

The Problem

DFLASH is a block-diffusion drafting method for speculative decoding. It uses a small "draft" model alongside the big "target" model. The draft model has two parts:

1. An encoder that reads hidden states from the target model's layers
2. A decoder with its own KV cache that generates draft tokens

The connection between these two: every time the target model processes a token, DFLASH needs to encode that token through the draft encoder, then inject the encoder output into the draft decoder's KV cache so the decoder can attend to it.

In the original implementation, this injection happened immediately inside the process() callback -- one llama_decode call per ubatch of tokens (typically 16 tokens). For a verification batch of 16 draft tokens, the original code would:

• Call process() → encode 16 tokens → immediately decode (inject) 16 tokens
• Generate 3 accepted tokens
• Target processes the next token at the new position → process() fires → encode & inject
• Repeat

Each injection decode is cheap -- it's just K/V projections (Wk, Wv matmuls) plus a cache write -- but it still takes GPU time, and it's serialized with the target model's own work.

The Optimization: Deferred Injection

The insight: we don't actually need to inject each encoder output into the KV cache right away. The draft decoder only reads the cache when it's time to generate drafts (the draft() call). Between process() calls, the cache isn't accessed.

So instead of injecting immediately, we stash the encoder output and wait. When the orchestrator calls draft(), we batch all the stashed entries together into a single llama_decode call.

This turns:

  tgt: [decode 16] → [encode&inject 16]  [decode 16] → [encode&inject 16]  [decode 16] → [encode&inject 16]  ... draft()

Into:

  tgt: [decode 16] → [encode 16]          [decode 16] → [encode 16]         [decode 16] → [encode 16]        ... [inject all 48][draft()]

The total work is the same, but batching the injection improves GPU utilization (larger matmuls) and removes serialization overhead between the target and draft GPU operations.

A safety cap (MAX_STASH = 64) prevents unbounded memory growth.

Benchmark Results

Two model families told opposite stories:

Qwen3.6 models (big improvement)

Model	Before	After	Change
Qwen3.6-27B dense	164.8 t/s	212.6 t/s	+29%
Qwen3.6-35B MoE	319.0 t/s	364-380 t/s	+14-19%

These models have near-100% draft acceptance rate. Every stashed encoder output gets used by the decoder, and batching the injection is a pure win.

Gemma4 models (regression)

Model	Before	After	Change
Gemma4-26B	312.4 t/s	289.6-293.2 t/s	-6%
Gemma4-31B	310.5 t/s	267.8-278.5 t/s	-10%
Gemma4-31B-QAT	308.9 t/s	237.2 t/s	-23%

Acceptance rate is the root cause. Gemma4 accepts only ~43% of draft tokens, versus Qwen's near-100%. These low rates mean most tokens the target verifies are rejected, but they still trigger process() calls that stash encoder output. Every rejection adds to the stash at a position the sequence will skip past. By the time draft() fires, flush_injection() has to process a large batch of mostly-wasted entries — concentrating that dead work inside draft() makes it the bottleneck. High acceptance (Qwen) avoids this entirely: the stash stays small, the flush is quick, and batching is pure win.

The Gate

Rather than applying the optimization universally, we detect the target model type and disable deferred injection for Gemma4 models. When disabled, the code falls back to the original per-chunk immediate injection behavior.

The detection is a simple substring check on llama_model_desc(), which returns strings like "gemma4 Q5_K_XL":

if (strstr(model_desc, "gemma")) {
    m_use_deferred = false;
}

When m_use_deferred = false:

• process() stashes each chunk and flushes immediately -- same as original code
• draft() skips the flush loop (stash is always empty)

When m_use_deferred = true (the default):

• process() accumulates a stash, auto-flushing only at 64 entries
• draft() flushes all stashed entries before generating

Verified After Gating

Model	Gen t/s	Note
Gemma4-31B DFlash	314.3	Back to PR baseline (~310)
Qwen3.6-35B DFlash	364.4	Still showing 14-19% improvement

What's in the Code

• ~70 lines added to common/speculative.cpp in the common_speculative_impl_draft_dflash struct
• Covers: stash type (StashedG), stash array per-sequence, flush_injection() method, deferred/immediate branch in process(), deferred guard in draft(), and model-type detection in the constructor
• One #include <cstring> for strstr

[TheTom]

@giveen — thanks again for the DFlash work and the benches on #219.

We could not land #219 as-is: base was master, so the GitHub diff was mostly fork history, and a straight merge would also undo some tip fixes (e.g. speculative/server gate from #217).

Tom rebased your intended CL onto current feature/turboquant-kv-cache as #220 (branch tom/pr-219-full-rebase).

He tried to keep your intended changes (content from your turbo-dflash tip for the DFlash paths): deferred KV injection, --dflash / --eagle3 / -cd, model/arch/spec stack, qwen35 extract, CUDA fattn bits, etc.

Only small tip-preserving tweaks on top:

• restore server : fix false speculative warning when not enabled #217 seq_rm / speculative enable gate in server-context
• dual-load DFlash target-layer metadata (your key + tip TARGET_LAYERS; same GGUF key)
• did not resurrect deleted openai-moe-iswa.cpp

We do not have your DFlash draft GGUFs here. Could you please re-test that PR branch the same way you did on #219 (5090, -md draft, turbo/q8 caches, acceptance)? If anything from your CL is missing or wrong, say so and we will fix on tip.

Once that looks good we can merge the rebased PR and close this one as superseded. Appreciate you driving this.