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.

[TheTom]

Merged: findings and verification summary

Full writeup of what changed and what was checked before merging this.

Build fixes (the CI blockers)

Two real -Werror failures in common/speculative.cpp, both in the DFlash deferred-injection path:

1. for (llama_seq_id seq_id = 0; seq_id < n_seq; ++seq_id) was comparing a signed llama_seq_id against unsigned n_seq with no cast, tripping -Werror=sign-compare on CUDA and ubuntu-arm64. Every other identical loop in the file (22 of them) already casts to (llama_seq_id); this one was the outlier. Added the cast to match.
2. sg.data.assign(inp_g + ..., inp_g + ...) on a freshly-constructed empty vector triggered a GCC/libstdc++ -Wnonnull false positive on the pre-allocation memmove path (the compiler can't prove the destination is non-null across the internal reallocation, even though GGML_ASSERT(inp_g && ...) already guarantees the source is valid). Replaced with resize() + std::copy(), which takes a different code path GCC doesn't flag. Same runtime result, no behavior change.

Both confirmed fixed on two independent platforms: a local build matching CI's exact -DLLAMA_FATAL_WARNINGS=ON flags, and CI's own ubuntu-arm64 job (failed before the fix, passed after).

Blast-radius verification

This changelist does more than add DFlash: it also moves the K/V f16 scratch-buffer allocation for the TILE/WMMA/MMA flash-attention paths from space reserved inside dst's own buffer to pool allocations (raw cudaMalloc/cudaFree on HIP, citing upstream issue ggml-org#22107 about pool retention on RDNA), and rewrites dequantize_V_turbo4_0's dispatch. That's shared infrastructure every turbo-quantized KV call goes through, not just the DFlash path, so I verified it didn't regress anything else on an RTX 5090:

• turbo3, turbo4, q8_0+turbo3 asym, q8_0+turbo4 asym, and plain f16 PPL are all bit-exact to pre-PR values
• fused turbo3 decode speed and the turbo2-hd128 fused-MMA gate (from PR turbo KV: correct Lloyd-Max centroids (4.125 bpw), PDL + fused-MMA decode for turbo4/3/2, perplexity 32K fix #197) are unchanged
• the -ctk q8_0 -ctv iq4_nl mixed-KV fallback (the crash fix from fix: cherry-pick 16 upstream P0/high-P1 fixes (CUDA FA crashes + overflows, server DoS, OOB reads, tool-call parsing) #202) still completes cleanly

Also found and cleaned up one real leftover: get_alloc_size() was still computing and reserving space for the old in-buffer K/V scheme even though nothing wrote there anymore (launch_fattn already uses the new pool-based allocation). That was pure wasted VRAM, not a correctness bug, but nvcc's own -Werror flagged the now-fully-orphaned struct and helper function as unused once I stopped calling them, so removed both and simplified get_alloc_size() to just return ggml_nbytes(dst). Re-verified bit-exact PPL after that cleanup too.

One pre-existing, unrelated item I ruled out along the way: test-backend-ops -o SET_ROWS_TURBO4 shows 4 "not supported" cases on CUDA. Confirmed this is identical on the tip commit with none of this PR's changes, so it's a separate, already-existing gap, not something introduced here.

Remaining CI reds (cuda, hip, musa, hip-quality-check) all trace to the same UI-assets CDN download flake (dist.tar.gz HTTP error) that's shown up on other unrelated PRs this cycle, not a code issue.

---

Thanks @giveen, this is a strong contribution, both the DFlash port itself and the deferred KV injection optimization (with the Gemma-vs-Qwen acceptance-rate analysis and the model-type gate) were well benchmarked and landed cleanly once rebased.

One thing worth a follow-up if you have the cycles: your very first benchmark comment on #219 showed a consistent, large prompt processing (prefill) regression with DFlash enabled across every model tested, independent of the generation-side work, e.g. Qwen3.6-27B 203.3 -> 100.9 t/s, Gemma4-31B 385.4 -> 139.0 t/s, Gemma4-31B-QAT 487.8 -> 142.7 t/s. The deferred-injection patch addressed decode; prefill still takes that hit whenever -md/--dflash is active. Would be great to see you dig into that the same way you did the decode side, since it looks like a real, well-quantified gap.