qwen4exp: Qwen3.8-Flash-Next long-context decode, +66% at 131k

[firelzrd]

Six patches on release/v0.7.4-staging (ea35c5066) that cut work the QSA indexer
was repeating, plus two dead allocations in the caches around it. No shader is
touched: every change is in how the graph is assembled, or in what the memory module
keeps.

Why

Qwen3.8-Flash-Next slows down badly with context. On a Strix Halo it does 47 t/s at
8k and 25 t/s at 131k. Profiling the decode step with GGML_VK_PERF_LOGGER=1 at two
depths and taking the increment says where it goes:

	share of the 2,048 → 32,768 increment
FLASH_ATTN_EXT	29%
CONT	15%
TOPK_QSA GET_ROWS	13%
GET_ROWS	11%
RMS_NORM_MUL	6%
ROPE	5%
ADD	5%

Attention itself is 29%. The other 55% is the plumbing that feeds the indexer, and
most of it is recomputed from scratch every ubatch even though nothing it depends on
has changed.

Result

Qwen3.8-Flash-Next UD-IQ4_XS, -ctk f16 -ctv f16, MTP n_max=3, ctx 262144,
llama-server with prompt cache reuse, temp=0, n_predict=160, gfx1151 / Vulkan:

depth	base	with these patches
2,048	33.69	27.26
8,192	47.29	49.70
32,768	35.22	42.82
131,072	25.05	41.59

+66% at 131k. The gain scales with context because what is removed is
proportional to n_kv, so there is little to win at 2,048 and a lot at 131,072.

The 2,048 row moves around between runs rather than tracking the patches: with MTP
on, a run that happens to accept one more draft token per step reads several percent
faster, and at that depth the patches themselves are worth almost nothing. The deep
rows are the ones that say what the series does.

Also measured on 2x RTX 3090 (CUDA, IQ1_S, -ctk q8_0 -ctv q8_0, no MTP, ctx
154624): 131,072: 11.97 → 13.52 t/s, +13.0%. What is removed is memory traffic
and dispatches proportional to n_kv, and a 3090 has several times a Strix Halo's
bandwidth, so removing the same bytes buys less. The two V narrowings gave back
700 MiB there, which moved that host's usable context ceiling from 157696 to 180224.

The patches

1. share the QSA input set per compress ratio — set_input_qsa is an O(n_kv)
   scan and it ran once per QSA layer on byte-identical data: twelve scans where one
   would do, about 41 ms per step at 131k, on the CPU and invisible to a GPU
   profiler. Not new work: this is upstream ggml-org's code, from the commit
   that added qwen4exp (6c84c7d5d, model: add Qwen3.8-Flash-Next (qwen4exp) ggml-org/llama.cpp#27742). This fork's
   independent port did not carry it, so it is picked back out and applied on its
   own; authorship on the commit is Daniel Han's.
2. drop the ggml_cont on the block-mean slices — ggml_add has no contiguity
   requirement on either the Vulkan or the CUDA backend (both gate it on type alone
   and index through the nb strides), and ggml_dup_tensor already gives the sum a
   contiguous home. The copies were about 34 MB per ubatch at 33k over 12 layers.
   Addition order is unchanged, so the arithmetic is identical.
3. cache the pooled indexer keys, recompute only the tail — a full block never
   changes again: its cells are written once, and pooling, normalisation and rotation
   are all position-determined. Rebuilding every block every ubatch was about a
   quarter of a 97 ms decode step at 131k. The commit message documents four things
   this has to get right, each of which cost a real bug during development.
4. narrow the indexer cache's V — nothing reads it. build_qsa_top_k only ever
   calls cpy_k and get_k on that cache; the indexer scores blocks against a key
   and has no value side. llama_kv_cache allocated one anyway at the model's own
   n_embd_head_v of 256. Worth 1.6 GiB at ctx 262144 with an f16 cache.
5. narrow the pooled-key cache's V as well — the same dead allocation in the
   cache patch 3 adds. KV buffer size 768.01 → 387.01 MiB, and GTT after load moves
   by the same 381 MiB with every other buffer unchanged to the byte.
6. drop the 1/r scale — its only consumer is an RMS norm, which is scale
   invariant: rms(x*s) = x / sqrt(mean(x^2) + eps/s^2). This is the one patch
   that is not bit-exact by construction: the effective epsilon moves from eps to
   r^2*eps, 1e-6 to 1.6e-5 against a mean square of order 1. It is worth 2.25 ms of
   a 97 ms decode step at 131k.

It is last on purpose. Everything below it is bit-exact, so if that trade is not
wanted, dropping the top commit leaves a series that changes no output at all.

The per-commit numbers were taken in the order the patches were originally
developed, which had 6 in third position. The endpoints are unaffected; only the
intermediate cumulative column would shift.

Verifying

Speed alone will not catch the failure modes here. Patch 3 in particular can be
completely inert — recomputing everything every step — and still produce correct
output, so a correctness suite that only walks forward passes while the cache does
nothing. What was used:

• Output hashes against the build without the patch, at four depths. Everything
  except patch 6 is bit-exact.
• Ascending, then descending, then a changed prefix, then back. Descending
  shrinks the cache through seq_rm; the changed prefix goes through the server's
  prompt cache and its restore path, which never calls seq_rm at all. The same span
  must give the same output however it was reached.
• A profile, not just a timer. GGML_VK_PERF_LOGGER=1 aggregated over 40+ decode
  steps. This is what caught an invalidation bug: end to end the build was only 5%
  slow, but GET_ROWS had not moved at all, which meant the incremental path was
  never running. Speculative decoding calls seq_rm on every single step, and
  treating that as "forget everything" makes the cache correct and useless at once.

One caveat on that profiler: it inserts a sync per op, so an op's time there is not
its contribution to the step. Two other candidates were dropped after checking that
distinction end to end.

Not included

• Extending the Vulkan FA gather union to GQA. Dropped after profiling at 131k:
  FLASH_ATTN_EXT is 4.5 ms there against 4.6 ms at 32k, so the gather already caps
  it. Ranking on the 32k profile alone would have overvalued this.
• Replacing the pooling gather with a view of the cache. Needs the cell layout to
  be the identity, and checking that is an O(n_kv) CPU scan at graph build time — the
  cost patch 1 removes, reintroduced in another form.