Bound DFlash SWA K/V cache to the sliding window

[SamuelOliveirads]

Following #2052, this PR gives DFlash SWA layers a K/V layout bounded by min(cross_ctx, n_swa). Full-attention layers retain the complete cross_ctx history.

This avoids exposing the wrapped full-context ring to SWA attention and makes the FA layout valid without changing model behavior.

I Validated with Gemma 4 26B IQ4_NL and Q8 DFlash draft. Each result uses 2 warmups followed by 5 measured 512-token runs:

• 512: main 111.94 ± 1.44 t/s, PR 112.58 ± 0.61
• 2048: main 115.20 ± 0.52 t/s, PR 115.82 ± 0.42
• 4096: main 114.23 ± 0.71 t/s, PR 115.12 ± 0.33
• 8176: main 112.92 ± 1.12 t/s, PR 113.98 ± 1.14

Everything else like output, tokens count, acceptance rate match it.

@ikawrakow I didn’t notice any uplift per se, so I’m not sure if I needed more context or if it’s because the Draft model has a lower cost.

[ikawrakow]

What was the cross context in these tests?

I was running this query (which I tested because that was the query used to claim 4X performance improvement in the beellama repo):

Write a complete Python 3 module implementing a doubly-linked list with the following methods:
append, prepend, insert_at, remove_at, find, reverse, to_list, length, is_empty, iter. Include comprehensive
docstrings, type hints, and pytest unit tests for every method. Return only the code, no commentary.

This produces in the range of 10k tokens. Acceptance rate is significantly higher when using a cross context of 8192 (compared to 4096 or less). In that case drafting takes up to 30% of the total time, and FA of the DFlash model is quite expensive because one cannot take advantage of the SWA to speed it up. Hence I was thinking that one will get a non-negligible performance gain by having a contiguous range for the attention mask in the DFlash SWA layers.

[SamuelOliveirads]

cross_ctx is set to the value of each bullet point: 512, 2048, etc. I used a synthetic prompt for testing, but while I was at it, I also tested the prompt that you passed with max_tokens=10000.

Results:

• cross_ctx=4096
  • main: 121.8 t/s, 82.27s, accept 88.55%, draft_n=5303, draft_accepted=4696
  • PR: 121.1 t/s, 82.72s, accept 88.55%, draft_n=5303, draft_accepted=4696
  • delta: PR is about 0.55% slower
• cross_ctx=8192
  • main: 116.7 t/s, 85.84s, accept 88.73%, draft_n=5298, draft_accepted=4701
  • PR: 118.8 t/s, 84.34s, accept 88.62%, draft_n=5301, draft_accepted=4698
  • delta: PR is about 1.79% faster

The final output was identical in both cases, although here I only ran a simple test for each, considering the inherent noise it is similar to the results from my first prompt.

[ikawrakow]

What is n_max? I used 4 and from memory acceptance was around 75% for 8192 and around 60% with 2048.

[SamuelOliveirads]

Only a n_max = 1, I can test it later with a larger size but initially it was --spec-type dflash:n_max=1,p_min=0.0,cross_ctx=8192

[SamuelOliveirads]

I ran a new benchmark using --spec-type dflash:n_max=4,p_min=0.0,cross_ctx=<ctx> for gemma-4-26B-A4B, using the shared prompt with max_tokens set to 10,000 and 7 runs each. The result was:

cross_ctx	main	PR	delta
512	153.93 ± 1.41 t/s	153.19 ± 0.28 t/s	-0.48%
2048	160.59 ± 1.90 t/s	161.39 ± 0.42 t/s	+0.50%
4096	162.43 ± 0.72 t/s	162.66 ± 0.42 t/s	+0.15%
8176	159.64 ± 0.29 t/s	162.01 ± 0.13 t/s	+1.49%

Also everything matched except cross_ctx 8176 where measured runs had main hash ffecab6525c9c84e vs PR hash 446db0e888bbf7c5, with main draft_n=10541 / accepted=7363 / accept=69.85% and PR draft_n=10528 / accepted=7367 / accept=69.98%.

[ikawrakow]

So, then, if you don't see a noticeable performance gain still merge?

I did a few quick tests and don't see real performance gains either. On the main branch I had profiled it, so that's how I saw the nearly 30% being taken by DFlash. This was the combination of the actual DFlash compute graph evaluation plus the time it takes to construct the cross context. I didn't do that here, but my guess is that all the gain in the DFlash evaluation is being traded against a slower cross context construction.

[SamuelOliveirads]

I've made some minor architectural changes to align with what's in vLLM, but essentially I don't see why we should merge it considering there may still be undiscovered bugs. That said, if you see value in the code because it aligns with the FA and you can envision future improvements it might be worth it.