Feature Request: --swa-extra parameter needed to restore speculative decode function with SWA

[steampunque]

Prerequisites

• I am running the latest code. Mention the version if possible as well.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

A new flag --swa-extra which enables specifying an amount to extend the sliding window length beyond the training length in models with SWA.

Motivation

Use of this parameter will allow spec decode to be used without losing any information while not needing --swa-full. Spec decode gives around 2X speedup and there is no reason to give up this performance gain on SWA models without requiring --swa-full to be used (i.e. the high memory of SWA and high speed of spec decode can both be achieved at once)

Possible Implementation

Example use: Here the SWA is extended by 8 tokens to accommodate a maximum 8 token speculation.

--swa-extra 8

when this flag is used the extra tokens must be masked out in the attention computation similar to the way they are now masked out with --swa-full and context > sliding window length.

[ggerganov]

After #13746 is merged, I will make this change. For simplicity, might keep the extra cells hardcoded to 32 or 256 instead of adding a parameter. Nvm, there is even a simpler solution.

[ggerganov]

@steampunque Could you try commit 8323e23 and see if speculative decoding works as expected with a SWA model? In general the performance should also be slightly better thanks to the smaller SWA cache that is used now. Let me know if you notice any issues.

[steampunque]

@steampunque Could you try commit 8323e23 and see if speculative decoding works as expected with a SWA model? In general the performance should also be slightly better thanks to the smaller SWA cache that is used now. Let me know if you notice any issues.

@ggerganov I did a quick run with Gemma 3 12B, speculated by Gemma 3 1B. Hw is a 4070 fully offloaded, spec is my custom downstream algo with fixed spec block size 4.

Run 1, with the new patch:

FATTN=1
SWA=1
NKV=65792
QKV=f16

Stats results of a prompt:
LOGPROB=-101.82892608642578 PERPLEXITY=1.2238988253775662
PN=504 PP=899.9069274342174 TG=54.42962665055953 DN=648 DA=342

Run 2, with the new patch and --swa-full used:

FATTN=1
SWA=0
NKV=11904
QKV=f16
LOGPROB=-101.82892608642578 PERPLEXITY=1.2238988253775662
PN=504 PP=478.5771181134883 TG=62.3601845069586 DN=644 DA=343

Defs:
FATTN flash attention
SWA sliding window attenction (SWA=0 means --swa-full)
NKV size of KV buffer in tokens
QKV KV quant type
LOGPROB : cumulative logprob of generation
PERPLEXITY: perplexity of generation
PN : number generated tokens
PP : prompt processing speed in t/s
TG : token gen speed in t/s
DN : number of drafted tokens total
DA : number of accepted draft tokens

Observations:

The identical generation was achieved, but there are differences in the underlying speculation stats which indicates a problem somewhere. (one less full block was speculated when --swa-full was given yet 1 more draft token got accepted, the divergence event(s) had to happen at very end of generation or probs and gen would not be identical). The results should be 100% identical on both gen results and spec stats with or without --swa-full flag used.

If I do the same prompt without spec:

Run1, no spec.

LOGPROB=-77.95978546142578 PERPLEXITY=1.1787707026323244
PN=474 PP=800.451761816258 TG=48.11971698297089

Run2, no spec, --swa-full used:

LOGPROB=-77.95978546142578 PERPLEXITY=1.1787707026323244
PN=474 PP=757.7212312451021 TG=48.28126341676723

Identical generations are given with spec turned off. Note gemma2 never speculated well and gemma3 appears not to be efficient either, but there is some spec speedup (less when SWA=1 for some reason).

[ggerganov]

Thanks for giving this a try.

What is the meaning of NKV and why is it less when --swa-full is used?

The PP with --swa-full seems unusually smaller (478.5 vs 899.9) - I wonder why that is.

The identical generation was achieved, but there are differences in the underlying speculation stats which indicates a problem somewhere.

The results (both for the draft and target) will not be bit-for-bit identical with/without --swa-full because this parameter changes how the KV cache data is ordered in memory and this makes some of the floating point reduction operations in the attention to produce slightly different results depending on where the KV cache ends up being placed (even if it is the same data). So occasionally small differences can be expected on master when changing the --swa-full flag. And these differences will be observed even more often with the new changes from 8323e23, because the cache buffer is even smaller (n_swa + n_ubatch).

[steampunque]

Thanks for giving this a try.

What is the meaning of NKV and why is it less when --swa-full is used?

Its the size in tokens of the KV cache. For this model the cache gets about 5.5x bigger due to the memory efficiency of SWA.

The PP with --swa-full seems unusually smaller (478.5 vs 899.9) - I wonder why that is.

I'm not sure what is going on here either, first time I ran a test like this.

The identical generation was achieved, but there are differences in the underlying speculation stats which indicates a problem somewhere.

The results (both for the draft and target) will not be bit-for-bit identical with/without --swa-full because this parameter changes how the KV cache data is ordered in memory and this makes some of the floating point reduction operations in the attention to produce slightly different results depending on where the KV cache ends up being placed (even if it is the same data). So occasionally small differences can be expected on master when changing the --swa-full flag. And these differences will be observed even more often with the new changes from 8323e23, because the cache buffer is even smaller (n_swa + n_ubatch).

My results showed the logits in my particular generation had to be exactly identical over ~500 tokens. At very end of gen something must have changed very slightly though. Good to know that there is a known source of possible difference though. I'm guessing for longer gens the small differences will start to kick in. 500 token probably not enough to even push out of the SWA window of this model.

One problem I did find with this model (Gemma 12B instruct) is when I set q8_0 KV cache, the computations start running on CPU instead of GPU and it becomes unusably slow. I'm not sure if this is expected or a bug. q8_0 opens up NKV to well over 100k tokens but its unusable due to the slow CPU processing. Qwen3 8B runs fine with q8_0 KV at full speed in master, but Qwen does not have SWA so I think something related to SWA is causing this problem. I can open a separate issue on this if needed, unless its a simple pilot error on my part.

[ggerganov]

@steampunque PR #13833 is the latest version that supports speculative decoding with SWA. If things go well with the pending changes, I plan to eventually merge that, so any feedback about that branch is appreciated.

[steampunque]

@steampunque PR #13833 is the latest version that supports speculative decoding with SWA. If things go well with the pending changes, I plan to eventually merge that, so any feedback about that branch is appreciated.

@ggerganov Merge appears to be working as before.

MODEL_NAME="gemma-3-12b-it"
MODEL_QUANT="Q4_K_H"
MODEL_DRAFT="gemma-3-1b-it"
CACHE=0
FATTN=1
SWA=1
QKV=f16
VERSION="llama.cpp version: 5556 "

NKV=65792

SPEC: fixed 4 tokens
Stats results of a prompt:
LOGPROB=-101.82892608642578 PERPLEXITY=1.2238988253775662
PN=504 PP=497.7634737071787 TG=63.471593942996456 DN=644 DA=343

NO SPEC:
LOGPROB=-77.95978546142578 PERPLEXITY=1.1787707026323244
PN=474 PP=1142.2223265347634 TG=47.903020032760004

Same tests with SWA=0 :

NKV=11904

SPEC: fixed 4 tokens
Stats results of a prompt:
LOGPROB=-101.82892608642578 PERPLEXITY=1.2238988253775662
PN=504 PP=939.0879269312408 TG=63.518197333395165 DN=644 DA=343

NO SPEC:
LOGPROB=-77.95978546142578 PERPLEXITY=1.1787707026323244
PN=474 PP=1161.5325865580448 TG=48.08895604659271

Testing with SWA=1 and QKV=q8_0:

NKV=109952
QKV=q8_0

SPEC: fixed 4 tokens
Stats results of a prompt:
LOGPROB=-90.51702117919922 PERPLEXITY=1.2019889925011844
PN=492 PP=290.36697612626585 TG=30.324223266901303 DN=628 DA=335

NO SPEC:
LOGPROB=-83.11764526367188 PERPLEXITY=1.1770081742194152
PN=510 PP=482.51223631202015 TG=21.904122819766695

Same huge slowdown as before with q8_0 kv, there is a lot of CPU processing going on when q8_0 is selected, happens with both SWA=1 and SWA=0 on this model.