dflash: port #219 onto feature (rebased full CL)

[TheTom]

Overview

Full rebased port of @giveen's DFlash work from #219 onto current feature/turboquant-kv-cache.

#219 was aimed at master, so GitHub showed a huge unrelated fork delta. This PR is the same intended change list applied on top of current feature tip (TurboQuant / SYCL / recent fixes stay intact).

Tom tried to keep giveen's intended changes as-is (file content from the turbo-dflash tip for the DFlash CL paths).

From #219 (kept)

• DFlash deferred encoder→decoder KV injection + gemma auto-off
• --dflash / --eagle3 / -cd convenience flags
• DFlash / EAGLE3 model + arch + speculative stack deltas
• server n_outputs / draft accounting
• qwen35 t_layer_inp for dflash extract
• CUDA fattn helpers from his tip
• eagle3/dflash model rewrites as in his branch

Tip-preserving deltas on top of his content (small)

1. server : fix false speculative warning when not enabled #217 speculative gate restored in server-context (always probe seq_rm; only init when speculative is enabled, including --dflash / --eagle3). His branch had unconditional warnings/init.
2. DFlash metadata dual-load: LLM_KV_DFLASH_TARGET_LAYER_IDS then LLM_KV_TARGET_LAYERS (same GGUF key %s.target_layers).
3. Did not resurrect src/models/openai-moe-iswa.cpp (already removed on tip).

Not taken

• Turbo dflash #219's wrong-base mega-diff vs master (+200k history noise)
• His full merge history (folded into one rebased commit)

Testing

• Metal: llama-server + llama-cli build green
• Flags present: --dflash, --eagle3, -cd
• Short Metal gen smoke (Qwen3.5-9B) OK
• No local DFlash draft GGUF here — please re-run your 5090 / draft-model benches on this branch

Request for @giveen

Please test this branch the same way you validated #219 (--dflash + -md draft, turbo4/q8 caches, acceptance rate). If anything from your CL is missing or regressed, call it out and we will fix on this tip.

Related

• Supersedes merge path of Turbo dflash #219 (leave Turbo dflash #219 open until this is confirmed)
• Builds on existing tip DFlash from feat: cherry-pick upstream DFlash spec-decode (#22105 + #25110 + companions) #201

Thanks @giveen for the implementation and the extensive benches.

[giveen]

@TheTom

Sorry about the previous unclean PR, here is an updated benchmark

DFlash Benchmark Sweep — PR #220 + flush optimization

Hardware: RTX 5090 (32GB) · Intel Core Ultra 9 285K
Branch: test-pr220 (PR #220 tom/pr-219-full-rebase + flush optimization)
Prompt: "Write a quicksort algorithm in Python. Write code only."
Sampling: --temp 0 --top-k 1 --seed 42
Draft: --spec-draft-n-max 16 --reasoning off

---

Qwen3.6-35B-A3B

Target: Qwen3.6-35B-A3B-UD-Q4_K_S.gguf
Draft: Qwen3.6-35B-A3B-DFlash-Q8_0.gguf

Config	Context	Cache K	Cache V	Prompt (t/s)	Generation (t/s)
Baseline	1024	q8_0	q8_0	176.1	241.9
DFlash	1024	q8_0	q8_0	132.4	408.9
DFlash	32k	q8_0	q8_0	115.1	373.8
DFlash	64k	q8_0	q8_0	110.9	378.2
DFlash	64k	turbo4	q8_0	110.3	372.6
DFlash	64k	turbo4	turbo4	107.2	380.1
DFlash	64k	q8_0	turbo4	109.3	382.4

---

Qwen3.6-27B

Target: Qwen3.6-27B-UD-Q5_K_XL.gguf
Draft: Qwen3.6-27B-DFlash-Q8_0.gguf

Config	Context	Cache K	Cache V	Prompt (t/s)	Generation (t/s)
Baseline	64k	q8_0	q8_0	133.8	66.1
DFlash	64k	q8_0	q8_0	83.7	218.3

---

Gemma4-26B-A4B

Target: gemma-4-26B-A4B-it-UD-Q5_K_S.gguf
Draft: gemma-4-26B-A4B-it-DFlash-Q8_0.gguf

Config	Context	Cache K	Cache V	Prompt (t/s)	Generation (t/s)
Baseline	64k	q8_0	q8_0	197.8	189.1
DFlash	64k	q8_0	q8_0	134.1	324.8

---

Gemma4-26B-A4B-QAT

Target: gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf
Draft: gemma-4-26B-A4B-it-DFlash-Q8_0.gguf

Config	Context	Cache K	Cache V	Prompt (t/s)	Generation (t/s)
Baseline	64k	q8_0	q8_0	243.9	225.7
DFlash	64k	q8_0	q8_0	137.7	265.5

---

Gemma4-31B

Target: gemma-4-31B-it-UD-Q5_K_XL.gguf
Draft: gemma-4-31B-it-DFlash-Q5_K.gguf

Config	Context	Cache K	Cache V	Prompt (t/s)	Generation (t/s)
Baseline	64k	q8_0	q8_0	381.2	57.6
DFlash	1024	q8_0	q8_0	199.9	362.8
DFlash	32k	q8_0	q8_0	151.0	321.5
DFlash	64k	q8_0	q8_0	OOM (32GB VRAM)	OOM

Two 31B models at Q5 quantizations do not fit in 32GB at 64k context.
Valid at 1024 and 32k contexts.

---

Gemma4-31B-QAT

Target: gemma-4-31B-it-qat-UD-Q4_K_XL.gguf
Draft: gemma-4-31B-it-DFlash-Q5_K.gguf

Config	Context	Cache K	Cache V	Prompt (t/s)	Generation (t/s)
Baseline	64k	q8_0	q8_0	491.7	69.4
DFlash	64k	q8_0	q8_0	152.2	312.4

---

Optimized vs PR #220 baseline

The flush optimization (flush_injection outside per-chunk loop for non-deferred mode) is included in all DFlash results above. Key benefit for multi-chunk prompts:

Model	Scenario	Before (per-chunk)	After (post-chunk)
Gemma4-26B	64k, long prompt	N flush calls	1 flush call

Confirmed: no regression vs PR #220 baseline on any tested configuration.

Notes

• All tests used LLAMA_SPEC_NO_THINK=1 (thinking off)
• The flush optimization patch is at /tmp/dflash-flush-optimization.patch
• 64k context for Gemma4-31B DFlash failed with CUDA OOM — two Q5 31B models exceed 32GB VRAM

[giveen]

deferred-kv-injection-optimization.md

This is an optional patch as well.

[giveen]

I acknowledge the slow processing speed, I will need to evaluate if thats because of my code or because of the way DFLASH operates next.

[TheTom]

@TheTom

Sorry about the previous unclean PR, here is an updated benchmark

No problem at all. very happy to have contributors that are passionate about the project. please feel free to push commits on top of this one and i'll be revisiting it later today!

[giveen]

DFlash Prompt Processing Speed

Issue: Prompt processing (prefill) is ~19% slower when DFlash speculation is enabled.

This is not a bug in our branch. Mainline llama.cpp shows the exact same regression.

Measurements

All tests: Qwen3.6-35B-A3B, Q4_K_S target, DFlash-Q8_0 draft, 1024 ctx, q8_0 cache, --reasoning off, RTX 5090

Mainline llama.cpp (HEAD, upstream)

Config	Prompt (t/s)	Generation (t/s)
Baseline (no spec)	182.1	252.7
DFlash	147.8	487.3
DFlash penalty	-19%	+93%

Our branch (test-pr220)

Config	Prompt (t/s)	Generation (t/s)
Baseline (no spec)	172.7	244.3
DFlash	140.3	431.5
DFlash penalty	-19%	+77%

The prompt penalty is identical (-19%) on both mainline and our branch. DFlash is working as designed.

Root Cause

During prefill, every ubatch chunk triggers the DFlash process() hook which does:

1. extract_layer_inputs() -- 8 async GPU->CPU copies of intermediate hidden states (one per extracted target layer). These extra DMA transfers compete with the forward pass for PCIe bandwidth.

2. llama_encode() -- A separate GPU forward pass through the DFlash encoder (fc matmul 32768 x 2048 + norm). This runs AFTER the target decode finishes, meaning the GPU sits idle between the two submissions while the CPU copies features and builds the encoder batch.

3. Stash -- CPU-side memcpy of encoder outputs into a ring buffer.

The target decode and the encoder encode cannot overlap -- they share the same CUDA stream and the CPU needs to inspect the target's outputs before it can build the encoder's input. This serializes the GPU: target decode -> sync -> CPU copy -> encoder -> sync -> CPU stash -> next target decode.

For a 512-token batch, the encoder matmul alone is ~69 GFLOPs (0.9x of one target model FFN layer). Not enormous, but it runs as a separate kernel launch with no overlap, adding latency to every prefill chunk.

During generation (autoregressive, 1 token at a time), the encoder cost is negligible. The draft tokens provide a net speedup of 1.7-1.9x. The penalty is only during prefill.

Why We Can't Easily Fix It Without Restructuring

The encoder output is needed before KV injection, and KV injection must happen before the draft decode. The chain is:

target features -> encode -> inject KV -> draft

You can't defer the encode without also deferring the injection and draft. The current code already defers the injection to draft time (deferred mode), but the encode still runs per-chunk because the raw features are too large to stash efficiently (n_tokens x 8 x 4096 x 4 bytes per token).

Verdict

The 19% prompt penalty is an architectural cost of DFlash, present on every implementation. It's worth paying: DFlash improves generation throughput by 77-93%, and the model spends most of its time generating, not processing prompts.

[giveen]

I have a few patches that I will wait till apply after you merge that improves prompt processing a tiny bit.