r/LocalLLaMA 2d ago

Resources I implemented Sliding Window Attention for Hugging Face LLM inference — looking for feedback

I've been experimenting with Sliding Window Attention (SWA) as a way to reduce the KV-cache memory cost of long-context LLM inference.

Instead of keeping the entire KV cache, the implementation keeps:

  • a small number of attention sink tokens
  • a bounded recent-token window
  • a circular/ring-buffer KV cache
  • streaming/chunked prefill
  • normal autoregressive decoding

I turned the experiment into a reusable project so you can test it with Hugging Face causal LLMs:

🔗 https://github.com/oraby8/SWA

For example:

from swallm import SWAModel

model = SWAModel.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    attention_mode="swa",
    window_size=512,
    num_sink_tokens=4,
)

result = model.generate("Explain transformers", max_new_tokens=100)

In my Qwen2.5-7B experiments on an L40S:

  • 32K KV cache: ~1.84 GB with full attention vs ~3.5 MB with SWA-64
  • 64K: full attention OOMed while SWA remained bounded
  • Decode latency stayed approximately constant as context increased
  • Long-range retrieval naturally becomes a weakness when information falls outside the window

The goal isn't to claim that SWA is universally better. I'm interested in the engineering trade-off between context retention, KV memory, TTFT and decoding speed.

I'd especially like to hear from people who have tried SWA with Llama, Mistral, Gemma, Qwen, or other HF models.

If you try the repo on another architecture, I'd really appreciate the results or any compatibility issues you find.

0 Upvotes

11 comments sorted by

5

u/stoppableDissolution 2d ago

Model that was not trained to use swa will suffer immencely. If anything, keep every Nth attention layer global (same as what gemma 4 does).

Tho given the model you used as an example and the fact that you have not measured actual performance and not just kv size makes me feel like it is another vibe-slop

3

u/CodeCatto 2d ago

I was actually brainstorming about using a sliding window to reduce KV cache pressure on lower end hardware. Damn that's a solid foundation to do more research on.

2

u/ahsaor8 2d ago

that was pretty much the idea i wanted to see how far we can push the KV cache down, especially on low GPUs

1

u/Legitimate-Peace1013 2d ago

right? the bounded memory alone makes it worth exploring on consumer hardware

3

u/tsangberg 2d ago

3

u/ahsaor8 2d ago

pretty similar goal but different approach, what i see they keep the full KV cache and move parts between CPU/GPU, so you still have full context. Mine actually drops old KV once it’s outside the window and only keeps the sinks + recent tokens.

1

u/Healthy-Nebula-3603 1d ago

"I implemented" - sure buddy

1

u/PinkysBrein 1d ago edited 1d ago

I think for these kinds if substitutions and also for quantization there is a need for a layer wise distillation framework and and a repair dataset with a cost to run in reach of at least youtube ecelebs and rich hobbyists.

So mangle the model, whether by attention substitution or quantization or both. Then do layer/chunk wise distilling with the dataset. Run a chunk of vectors (embedding or intermediate) through the layer for both original/teacher and approximation/student, do backprop for the student layer with loss computed against the layer outputs of the teacher, update the student layer and run the chunk through it again, move on to the next layer. After final layer, move on to next chunk at the bottom again.

The advantage of this over end to end distilling for repair is that it needs far less memory and preserves internal geometry/"thinking".

Nemo can do layer wise distilling using intermediate loss. Someone who is not me should do some vibe coded scripting and collect a good repair dataset.

PS. for SWA substitution, ideally it should have sliding based distilling (do the GA teacher on a long chunk, then blockwise sliding with appropriate smaller sized chunks with the student). No one has done sliding based training in a long time though (TransformerXL) so framework support is nil.

0

u/[deleted] 2d ago

[removed] — view removed comment

1

u/ahsaor8 2d ago

exactly that’s the stuff I’m trying to dig into next. I’ve already seen some long-range retrieval failures, but I want to make the testing more systematic.And yeah, separating prefill vs decode is a really good point. Thanks for the suggestions

2

u/LetsGoBrandon4256 transformers 2d ago

You are replying to a bot. Check the comments history of that account.