spec: enable backend sampling for both dflash & dspark

[ruixiang63]

Overview

Similar to #24655, this PR adds backend sampling support for DFlash. After #25532 got merged, we can enable multi-token sampling for DFlash/DSpark now.

This PR provides a simpler, more robust implementation to replace #26842.

Additional information

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, paired with cursor

[pich]

I ran #26958 against the specialized #26842 implementation on my Glimmer setup.

Same box, same config:

RTX PRO 4000 Blackwell SFF 24 GB / 70 W
Debian 13.6, kernel 6.12.38
driver 610.57.04, CUDA 12.9
Muse Glimmer 30B Q5_K_M, 256K ctx, Q8 KV, DFlash F16 KV, n_max=15, mmproj

10 runs, 1024 tokens each:

#26842: 80.05 tok/s avg
#26958: 80.69 tok/s avg

Pretty much identical performance, #26958 was actually a tiny bit faster here.

Acceptance was identical too: 38.67%, average accepted length 6.77. CPU usage was basically unchanged at around one core.

The one difference I noticed is memory. #26842 loaded at 22,852 MiB, while #26958 used 23,052 MiB, so roughly 200 MiB extra.

I reran the mixed workload as well and got 37.51 tok/s with 14.41% acceptance, very close to my previous numbers. So the big difference between coding and mixed workloads still seems to be mostly about how predictable the output is.

Also tested with the KV actually filled to 262,116 tokens. No OOM, peak was 23,120 / 24,467 MiB.

I like the direction of #26958 more. We get basically the same speed without needing the DFlash-specific argmax path, and the same mechanism can work for DFlash and DSpark.

Maybe one thing worth looking at is that extra ~200 MiB. With greedy top_k(1) I wonder if we could keep only the sampled IDs and avoid some of the full logits/probs/candidate buffers.

I wrote up the original Glimmer/DFlash experiments here if any of the numbers are useful for comparison:
https://piszczek.pl/blog/dflash-changes-what-tokens-per-second-means

[ruixiang63]

I ran #26958 against the specialized #26842 implementation on my Glimmer setup.

Same box, same config:

RTX PRO 4000 Blackwell SFF 24 GB / 70 W Debian 13.6, kernel 6.12.38 driver 610.57.04, CUDA 12.9 Muse Glimmer 30B Q5_K_M, 256K ctx, Q8 KV, DFlash F16 KV, n_max=15, mmproj

10 runs, 1024 tokens each:

#26842: 80.05 tok/s avg #26958: 80.69 tok/s avg

Pretty much identical performance, #26958 was actually a tiny bit faster here.

Acceptance was identical too: 38.67%, average accepted length 6.77. CPU usage was basically unchanged at around one core.

The one difference I noticed is memory. #26842 loaded at 22,852 MiB, while #26958 used 23,052 MiB, so roughly 200 MiB extra.

I reran the mixed workload as well and got 37.51 tok/s with 14.41% acceptance, very close to my previous numbers. So the big difference between coding and mixed workloads still seems to be mostly about how predictable the output is.

Also tested with the KV actually filled to 262,116 tokens. No OOM, peak was 23,120 / 24,467 MiB.

I like the direction of #26958 more. We get basically the same speed without needing the DFlash-specific argmax path, and the same mechanism can work for DFlash and DSpark.

Maybe one thing worth looking at is that extra ~200 MiB. With greedy top_k(1) I wonder if we could keep only the sampled IDs and avoid some of the full logits/probs/candidate buffers.

I wrote up the original Glimmer/DFlash experiments here if any of the numbers are useful for comparison: https://piszczek.pl/blog/dflash-changes-what-tokens-per-second-means

Thanks for testing and sharing the results. Very useful! @pich

[am17an]

This is still not going to work with -sm tensor right?

[pcuenca]

cc @ngxson @CISC

[pcuenca]

It works according to my tests and @pich's 🚀

[ruixiang63]

This is still not going to work with -sm tensor right?

I think -sm tensor falls back to CPU sampling when the draft model itself is tensor-split, GPU sampling still kicks in when the draft runs on a single GPU.