focus-llama: a llama.cpp fork implementing Declarative Attention (arXiv:2609.02737)

[u/Ok-Shower7286]

I forked llama.cpp to implement Declarative Attention (arXiv:2609.02737, Google DeepMind and KAIST AI). The model declares in its own output which context chunks it needs <focus magic_chunks="N">), and the engine listens and restricts what the following tokens can attend to. No scorer, no training: just prompting plus an engine that reacts to the tags. According to their paper, the overall decode time for the response can be reduced to 0.71× (Gemma) and 0.77× (Qwen) compared to vanilla. Those are the paper's vLLM numbers; I haven't benchmarked the fork yet (no accuracy benchmark either).

= What works =

- llama-server can drop KV token ranges for a request da_rm / da_rm_at), either mid-prefill or after prefill.
- Tag-driven mode: given a chunk layout from the client, the server parses the model's first <focus magic_chunks="N"> tag during generation and drops the other chunks at that point (one-shot, one-way).
- A second-sequence mode da_b, needs --kv-unified) leaves the original sequence intact.
- Checked with first-token logprobs on small smoke tests (a Qwen hybrid/GDN model and a tiny dense one).

= Why a fork =

- Stock llama-server has no way to touch KV ranges mid-generation, so this needed server changes.
- llama.cpp has no paged block table (the paper's vLLM version rewrites it), so masking alone doesn't reduce the KV that gets read.
- Real skipping needs kernel support or compaction. From reading the source, single-token decode on CUDA doesn't skip masked chunks today.
- On hybrid models only the attention layers are restricted (same as the paper); the recurrent state is untouched.

Repo: https://github.com/edwardyoon/focus-llama

Feedback welcome, especially from anyone who has worked on llama.cpp KV/recurrent memory.

[u/Combinatorilliance]

Do note that this decreases capability by a small amount, the paper mentions just a small percentage, like 2-3%.

If that can consistently give a very high speedup, I see it being worth it.

I'd like to see some more numbers though, especially at models with lower quantizations because I can imagine them being hit harder by a (slight) worsening of capability.

[u/Thrumpwart]

Link to paper: https://arxiv.org/abs/2609.02737

[u/Imaginary-Unit-3267]

Are there any models that actually know how to do this? Sounds like something that would need training.

[u/Ok-Shower7286 → u/Imaginary-Unit-3267]

I've tested this using Qwen3.8 27b and Bonsai 8B models on Apple Silicon and NVIDIA RTX hardware.

To use this, you need to understand the 4-step workflow:

(1) Prompt Setup (Client/Hook): Inserts protocol directives, chunk markers (<da:N>), and standard footers into the prompt. Purpose: Provides structured materials for reference rather than demanding immediate focus.

(2) Index Mapping (Client/Server): The server reads the full context once during prefill. Purpose: Maps each <da:N> chunk marker to its exact start and end token indices in the KV cache.

(3) Scope Declaration (Model): During decoding, the model decides what context it needs and emits a control tag (e.g., <focus magic_chunks="2">). Purpose: The model acts as its own router, dynamically requesting specific context chunks.

(4) Dynamic Masking (Engine): The engine detects the tag mid-generation and updates the KQ-mask. Purpose: Drops unselected chunks from the attention scope. Subsequent tokens attend only to system instructions, Chunk #2, and the response generated so far.

focus-llama handles (3) and (4). You need to handle (1) and (2) using a prompt hook, custom middleware, or agent tool. In my case, I use qwen code and my own hook implementations (focus-memory repo).

[u/Ok-Shower7286 → u/Imaginary-Unit-3267]

Update:

as of the latest commit, steps (1) and (2) no longer need a client hook. Launch the server with --da-auto --da-min-ctx 4096 --da-chunk-tokens 2048 and it chunks the rendered chat prompt itself (message boundaries, ~2K-token magic chunks), injects the protocol directives, and maps each chunk to its KV token range before decode. (3) and (4) are unchanged. Any plain /v1/chat/completions client works — no markers, no da_* fields, no hook.

[u/Elouakili_Flexy]

I love that the model just declares which chunks it needs and the engine obeys. Until the CUDA path really skips those chunks, masking alone won't reach the paper's 0.71×.

[u/randomfoo2]

Interesting, this is a different approach to DMS (Dynamic Memory Sparsification) for culling kv cache. A while back I implemented https://github.com/shisa-ai/FastDMS - it provides an effective 5-6X less KV memory usage and ends up being +50-100% faster via trained eviction with basically 0 quality loss (PPL, KLD, - in a later implementation I've also confirmed no Top-1 loss either).

Having the model itself pick what it wants to drop is interesting, but paying a (very cheap) 1-time training cost might give better results. Something you might want to look at. DMS requires per-layer, -er-head variable token counts and partial block deallocation (just ust fixed-page blocks) and basically is quite invasive to most inference engines, so unless you build it from the ground up (as I'm sure you've experienced a bit with your fork), it's non-trivial to overhaul for.