spec : add DFlash2 support (local convolution + candidate selector)

[SubSir]

Overview

DFlash2 adds two more modules compared to DFlash, grouped dynamic depthwise convolution and candidate selector.

The formula is for reference. out[i,c] = Σ_t (base[t,c] + δ[i,t,g(c)]) · x[i−t,c]. base is the static kernel, δ is the dynamic part predicted from the input.x is divided into groups g(c) and δ is shared across a group. Another one is edge(p→c) = ⟨A[p] ⊙ project(h), B[c]⟩ + unary[c]. A/B are the codebooks for the predecessor and the candidate. This position's hidden is projected by project(h)to the same rank to multiply A. And unary[c] is the draft's own score for c.

DFlash2 is enabled when the checkpoint is DFlash2; no need to use extra flag.

We evaled Qwen3.8-27B at Q4_K_Mon Apple M5 Pro 64 GB. And the test data is first 8 problems of GSM8K.
Concurrency is set to 1 and temperature is 1.0, top-p 0.95, top-k 20. We use default xhigh reasoning and set
2,048 maximum new tokens.

Method	Draft quant	Block	Decode TPS	Decode speedup	Acceptance	Normal stops
Autoregressive	—	—	10.42	1.00×	—	8/8
DFlash2	BF16	8	19.31	1.85×	4.92	8/8
DFlash2	Q8_0	8	18.43	1.77×	5.08	8/8
DFlash2	Q4_K_M	8	18.89	1.81×	5.03	8/8

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. Using AI to inspect and analyze code in the llama.cpp project, and using AI to update the code.

[ggml-gh-bot]

Hi @SubSir, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[SubSir]

I've revised the PR description. It's fully hand-written now.

[treo]

@Shamish There is already a fix for the reported problem at z-lab#1

As for the Performance on a single RTX3090, it seems to be not significantly better than MTP:

MTP:

$ ./build/bin/llama-server -m ~/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/f1bfb127c64f7072bdd2cad55f258b9c8b2910fe/Qwen3.8-27B-UD-Q4_K_XL.gguf  --fit off -np 1  -fa on --spec-type draft-mtp --spec-draft-n-max 3 --reasoning-preserve -c 32000 -ctk q8_0 -ctv q8_0

$ python speed_bench.py --url localhost:8080 --model qwen3.8-27b --concurrency 1 --limit 2
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
speed_bench: loaded 22 samples from bench=qualitative category=all
speed_bench: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 22/22 [11:57<00:00, 32.60s/sample]

Summary (elapsed=717.20s)
category       samples  avg_prompt_t/s  avg_pred_t/s  avg_latency  accept_rate
-------------  -------  --------------  ------------  -----------  -----------
coding         2        454.24          59.84         69.160s      0.6047     
humanities     2        508.66          57.24         56.283s      0.6075     
math           2        27.86           55.76         17.593s      0.5927     
qa             2        85.46           54.37         18.066s      0.5589     
rag            2        704.60          59.45         40.403s      0.6777     
reasoning      2        60.51           54.47         18.138s      0.5954     
stem           2        28.90           54.39         17.867s      0.5981     
writing        2        612.05          61.74         42.117s      0.6609     
multilingual   2        164.13          65.08         4.339s       0.7778     
summarization  2        131.48          54.36         4.965s       0.5473     
roleplay       2        397.24          58.95         69.665s      0.6814     
overall        22       288.65          57.79         32.600s      0.6317  

DFlash2 (using maxn=4 because 7 was slower):

$ ./build/bin/llama-server -m ~/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/f1bfb127c64f7072bdd2cad55f258b9c8b2910fe/Qwen3.8-27B-UD-Q4_K_XL.gguf  --fit off -np 1  -fa on --spec-type draft-dflash --spec-draft-n-max 4 --reasoning-preserve -c 32000 -ctk q8_0 -ctv q8_0 -md ~/.cache/huggingface/hub/models--incoai--Qwen3.8-27B-DFlash2-GGUF/snapshots/6cb5872e2cee6b4e780a8414922350be8e42d65c/Qwen3.8-27B-DFlash2-Q4_K_M.gguf

$ python speed_bench.py --url localhost:8080 --model qwen3.8-27b --concurrency 1 --limit 2
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
speed_bench: loaded 22 samples from bench=qualitative category=all
speed_bench: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 22/22 [11:40<00:00, 31.85s/sample]

Summary (elapsed=700.65s)
category       samples  avg_prompt_t/s  avg_pred_t/s  avg_latency  accept_rate
-------------  -------  --------------  ------------  -----------  -----------
coding         2        271.99          68.00         48.928s      0.6691     
humanities     2        470.70          59.87         62.873s      0.5610     
math           2        28.45           60.14         29.131s      0.5827     
qa             2        82.64           53.74         13.088s      0.5039     
rag            2        701.27          66.42         37.589s      0.6378     
reasoning      2        58.79           60.55         29.046s      0.5827     
stem           2        29.38           60.20         28.899s      0.5827     
writing        2        607.63          69.16         38.621s      0.6552     
multilingual   2        167.16          62.26         5.469s       0.5316     
summarization  2        136.16          59.60         4.324s       0.5690     
roleplay       2        375.62          64.43         52.359s      0.6758     
overall        22       266.35          62.21         31.848s      0.6153  

[ruixiang63]

Thanks for the PR. It would be great if you could run the benchmark on SpeedBench using the corresponding checkpoints and report the performance numbers. @SubSir

[ngxson]

llama.cpp code looks ok, but need an approval for CUDA part

regarding vision input bug: that will be a follow-up hot fix. proper fix will come after #24669

comments regarding this bug will be removed

[ggerganov]

Move this to llama-ext.h

[SubSir]

Done in e91bca5.

[ggerganov]

Suggested change

	MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base",
	MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj",
	MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base",
	MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj",
	MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor",
	MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor",
	MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden",
	MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base",
	MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj",
	MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base",
	MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj",
	MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor",
	MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor",
	MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden",

[SubSir]

Done in 64f765f

[ggerganov]

Suggested change

	ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false);
	hparams.f_final_logit_softcapping = 0.0f;
	ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
	ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false);
	ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false);
	hparams.f_final_logit_softcapping = 0.0f;
	ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
	ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false);

[SubSir]

Done in 64f765f

[ggerganov]

Move this to a separate block - no need to break the indentation here.

[SubSir]

Done in 64f765f. Thanks!

[am17an]

cc: @fairydreaming @ORippler for changes to the top-k kernel

Temporary dismiss my review because the m-rope bug is quite déjà-vu, saw smth similar during muse glimmer development. I will have a look

[fairydreaming]

cc: @fairydreaming @ORippler for changes to the top-k kernel

I tested the performance of the included TOP_K kernel on RTX PRO 6000 Max-Q - compared it to the CUB argsort implementation and upcoming DeviceBatchedTopK (we don't use it yet, still waiting for the next CCCL release).

Used the following test cases and compared times per one test case run (k=64):

for (auto nrows : {2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096 }) {
    for (auto cols : {128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65000, 200000}) {
        test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 64));
     }
}

But noticed that this implementation runtime increases with k, so I also ran slightly different test set with k=1024:

for (auto nrows : {2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096 }) {
    for (auto cols : {2048, 4096, 8192, 16384, 32768, 65000, 200000}) {
        test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 1024));
     }
}

Results (k=64)

As you can see this kernel is generally slower than DeviceBatchedTopK.

But for small k it's faster than the argsort-based implementation. For 200000 cols the implementation falls back to argsort-based implementation.

Results (k=1024)

For large k the new kernel is used only for low number of cols but it's clearly much slower.

Verdict

This kernel is faster only in some cases (small k values). My opinion is that it's better to wait for DeviceBatchedTopK (CUDA) and rocprim::segmented_topk (ROCm) that should be universally faster instead of adding multiple implementations for specialized cases.

[ORippler]

And I added a topk kernel. I found that the original topk implementation is relatively slow in a small k and large number of candidates, which is just the Dflash2 selector's topk case. It overrides ggml_cuda_op_top_k when the row is wider than 1024 and k is small enough, and in my test it's 17x faster than the original topk function on H200 and 15x faster on A10G in Dflash2's setting. Since I override ggml_cuda_op_top_k, it also takes over for the original backend sampler. But in my test, it's faster than the original kernel. If you guys want me to gate the topk kernel only for Dflash2 to achieve a minimal change in this pr, just tell me and I don't mind.

How do these per-kernel times transfer to E2E times on consumer-grade HW?

[ngxson]

@SubSir could you please apply this patch to fix the image input (mrope bug)? https://gist.github.com/ngxson/b5ef18797c1bdcb1814cf09afadfd7eb

that should make it aligned with muse glimmer rope handling

[ngxson]

GGUF will need to be reconverted, and vision input should work with new GGUF (tested on my side)

Please also fix the failing CI

[SubSir]

@ORippler, I tested on RTX PRO 6000 Blackwell using Qwen3.8-27B. It's about 2-5% end to end throughput impact.

	old top-k, t/s	new top-k, t/s	gain
BF16	59.77	61.33	+2.62%
Q4_K_M	60.14	62.14	+3.33%
Q4_K_M, MMVQ capped	92.69	97.08	+4.74%

(Q4_K_M MMVQ capped means I use GEMM kernel for an 8-token batch. It's another issue, not directly related to Dflash2. I provided the data here because otherwise it's a bit strange that BF16's throughput is nearly the same as Q4_K_M).

So if it's not that much to you, I'll follow @fairydreaming's opinion and use the original kernel. But the cost of Dflash2 is a bit higher on llama.cpp than sglang or vllm.

@ngxson @sswtodo, the ci fail is fixed in f7aadef, and I just uploaded the converted GGUF to z-lab/Qwen3.8-27B-DFlash2-GGUF. Later we'll update incoai/Qwen3.8-27B-DFlash2-GGUF. Thanks!

[ORippler]

So if it's not that much to you, I'll follow @fairydreaming's opinion and use the original kernel. But the cost of Dflash2 is a bit higher on llama.cpp than sglang or vllm.

For 2-4% E2E I'm fine with taking it, at least until CCCL3.5 has released. This was Linux? Or Windows

[ruixiang63]

Quick question: It looks the size of dflash_selector is not small as described in blog +2.0M, e.g. with rank = 256 and n_vocab = 248320(qwen3.8 27B) that comes out to 2 × 256 × 248320 ≈ 127M parameters. Is this expected?

[SubSir]

We didn't explain it clearly in the blog. Actually predecessor_codebook and successor_codebook are projected from the target embedding and for Qwen3-4B, plus the hidden projection is 3 * 2560 * 256 = 1.97M. That's what we said 2M. And for implementation convenience, we materialized the projection into codebooks since 77.8M or 127M is not that much.

[ruixiang63]

I feel this part we could optimize with some fusions. Need suggestions from @ggerganov

[ruixiang63]

Nits: the 3/4 suffixes here encode tensor rank rather than meaning, and it gets hard to follow once hid3 / hid / hid4 (and unary / unary3 / un / un4) all coexist across nested scopes. Something like cand_blk / unary_blk / gate_blk would read better.

[SubSir]

Done in 2f3923b. Thanks!

[ruixiang63]

why is 32? can you add comment to explain it?

[SubSir]

Sorry, this is from the old version where selector is implemented by for loop. The node number shouldn't be calculated from the num of token in the batch in this version, and I missed it.

I counted on different number of layers and found that nodes per tensor is about 8.8, a bit higher than 8 nodes per tensor in the original code. So I add a special case for Dflash2 to lift the node per tensor to 12. Just in case that when the number of layers is larger than 8, graph nodes is larger than 1024 and it will throw an error.

layers	tensors	graph nodes	nodes / tensor
4	66	582	8.8
6	96	840	8.8
8	126	1098	8.7
12	186	1614	8.7
16	246	2130	8.7

[art-den]

ggml\src\ggml-backend-meta.cpp:543: GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) failed

On startup with --split-mode tensor for my x2 RTX 5060 Ti.

PS:
DFlash2 file: https://huggingface.co/z-lab/Qwen3.8-27B-DFlash2-GGUF/blob/main/Qwen3.8-27B-DFlash2-Q4_K_M.gguf
It works with --split-mode layer

[SubSir]

@ORippler It's on Linux. Sorry I don't have resources on Windows.