hipfire
/learn · attention · memory

QSA sparse attention and the VMM KV backend

Qwen3.8 Flash-Next accepts up to 262,144 tokens of context. Two things make that usable on a 32 GB card or an APU: the model attends to a selected 2,048 tokens per query instead of all of them, and hipfire's VMM backend only commits memory for the context a request actually touches.

01 Where attention lives in Flash-Next

the idea

Flash-Next has 48 layers. Every fourth one is a Qwen Sparse Attention (QSA) layer; the other 36 are Gated DeltaNet layers, which carry a fixed-size recurrent state instead of a per-token cache. So only 12 layers keep per-token K/V at all, and each of those reads only a selected subset of it.

fig · layer pattern, 48 layers
  L0  L1  L2  L3   L4  L5  L6  L7   …   L44 L45 L46 L47
  GDN GDN GDN QSA  GDN GDN GDN QSA   …   GDN GDN GDN QSA
  └── fixed-size recurrent state ──┘    QSA: per-token K/V + index keys, sparse reads

02 How QSA chooses what to read

the idea

Besides its ordinary K/V, a QSA layer writes a small index key (one 128-wide head) for every token. Every 4 consecutive index keys are averaged into one pooled block key. For each query row, four small index-query heads score every pooled block (the sum of their ReLU’d dot products, scaled by 1/√128), and the top-scoring blocks within a 2,048-token budget (512 blocks) are kept, plus the incomplete block at the end. Full attention then runs over only those rows.

fig · schematic: one query row's selection over pooled blocks
context, 4 tokens per pooled block → ■ selected (top blocks within the budget) ■ incomplete tail block, always kept attention reads only the selected rows: ≤ 2,048 + 3 tokens per query, at any context length

The cost of the attention itself stops growing with context; what still grows is the cheap scoring pass over pooled blocks (one per 4 tokens) and the memory to hold the cache.

03 How hipfire runs it on RDNA

how hipfire does it
  • Read each selected row once. Flash-Next has 24 query heads over 2 KV heads. The grouped kernel gives one workgroup to each KV group of 12 query heads, so a selected K/V row is read once instead of three times.
  • Score, then select. Block scoring has a few-row variant (≤ 8 rows: decode and verify) and a 16-row variant (prefill), followed by a radix select from the scores. This series folds in the bit-exact QSA work of #774 by @fivetide.
  • Gathered F16 WMMA for prefill. On Strix Halo and the R9700, prefill chunks of ≥ 512 rows convert the layer's cache to F16 K plus a block-transposed V, then one workgroup per (row, KV head) walks that row's selection in 128-entry tiles on WMMA. Its scratch (2 KiB per context token, 512 MiB at 262K) is reserved at load, so it never grows under a captured graph.
  • PeaceMaker builds. The gathered attention (both arches) and the Strix Halo selector run from certified PM modules, byte-identical to the hipcc kernels.

State formats

QSA K/V statedefault onnotes
f32Strix Halo (gfx1151)exact reference; 4,736 bytes per token per layer
fp8 (E4M3 K/V + f16 scale per head and token, BF16 index keys)R9700 (gfx1201)gfx12 only
q8 (Q8_0 K/V, BF16 index keys)opt-in (--kv-mode q8)1,408 bytes per token per layer; selection keys unchanged

04 VMM: reserve big, commit what you touch

the idea

At F32, 12 QSA layers at the full 262,144-token context are 14,208 MiB of context state. A server has to be able to reach that length, but most requests never do. The legacy backend allocates the whole reservation as one physically contiguous block up front. The VMM backend reserves a virtual address range of that size and maps physical pages into it only as the context grows.

fig · legacy vs VMM for one QSA arena
  legacy   [██████████████████████████████████████████████]  physical for all of max_seq, at load
  VMM      [████████▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒▒]  virtual range for max_seq
            └ mapped ┘ ← grows as the context is written
  both keep one stable base pointer, so captured graphs and replays stay valid
how hipfire does it

In 0.4.1 the Flash-Next QSA context and workspace are demand-mapped automatically on certified gfx1201 (R9700) as well as gfx1151 (Strix Halo); committed pages grow with touched context. Unsupported combinations fall back to legacy once per load with a logged reason.

One sharp edge shaped the design: on ROCm 10.0, a virtual address can keep translating to its first physical backing after it is unmapped and remapped to a different one, so kernels read stale pages (every page in a repro on gfx1201, 84 % on gfx1100, none on gfx1151). hipfire therefore never hands a released virtual range back to the driver; it retires it for the life of the process, which is cheap (Qwen3.8-27B at max_seq 4096 retires 192 MiB per load).

05 What it buys you

Flash-Next cold prefill on real prompt files · 0.4.1 · medians of three fresh processes
2,082.6tok/s
R9700, 8K prompt
range 2,078–2,087.8
2,051.3tok/s
R9700, 32K prompt
Strix Halo: 1,776.5 tok/s
1,814.6tok/s
R9700, 128K prompt
87 % of the 8K rate at 16× the context
Prompt files are the Ciru prompt set. R9700 figures were measured on an earlier build of the 0.4.1 release cycle and were not re-run on the final build. On the R9700 VMM backend, a 64K-token prompt at max_seq 262,144 decodes coherently.

06 Limits and switches

exactness
  • Gathered F16 WMMA attention is not bit-exact: it rounds operands to F16 and is admitted against an f64 reference with a BF16-storage error bar. A near-tied argmax at short context can flip. Opt-out: HIPFIRE_QWEN4_QSA_WMMA_GATHER=0.
  • q8 QSA state is opt-in, not a default.
  • Past 15,360 pooled blocks (61,440 tokens) the selector keeps its score rows in global memory instead of LDS.
  • memory.kv_backend = "legacy" / HIPFIRE_KV_BACKEND=legacy opts out of VMM.

07 Sources