openpangu: fused latent attention op (GGML_OP_LATENT_ATTN)

[joelfarthing]

This is the GGML_OP_LATENT_ATTN split requested in #2159. The ablation there put nearly all of the TG improvement on this op, so this PR carries exactly that piece, rebuilt as one commit on current main, with applicability checking cut down to a single per-layer capability gate. Any other #2159 ops are deferred to one-by-one follow-ups under the openPangu + DS4 usefulness bar. This op fuses openPangu's latent-attention chain, scores through joint softmax over the learned param_sink prefix rows plus the cache through value contraction, into one node.

GGML_OP_LATENT_ATTN is:

• ggml_latent_attn_prefix_ext(q, cache, prefix_k, prefix_v, mask, ...): dense mode. q is [Dk, T, H] F32; cache is [Dk, N] in F32, F16, or Q8_0 through a windowed row view. F32/F16 rows are read in place; Q8_0 is dequantized per call into a pool-allocated F16 buffer (packed rows required). The prefix rows are always visible at bias 0 and jointly softmaxed with the cache scores; the mask applies to the cache segment only. op_params carry dv, dv_off (MLA value-slice layout, so DeepSeek-style rope-first packing is expressible), scale, and max_bias (asserted zero; the op has no ALiBi and the gate refuses it rather than silently ignoring it).
• ggml_latent_attn_indexed_ext(q, cache, prefix_k, prefix_v, indices, ...): sparse mode for DSA decode. Same contract, but attends only the top-k cache rows selected by the existing indexer (indices [topk, T]). The CUDA path gathers the selected rows into bounded in-op tiles rather than materializing a graph-level gathered cache tensor, so the graph carries no ggml_get_rows result and the tile memory is pool-reused. openPangu's indexed sites pass no mask (selection is already causal-masked upstream); the indexed-mode mask path exists on both backends but has no in-tree user yet.
• CUDA implementation (ggml/src/ggml-cuda/latent_attn.cu) plus a scalar CPU reference forward in ggml.c that pins the op's semantics; the CPU backend truthfully reports support, and not adopting the op on CPU-resident layers is performance policy in the openPangu gate below. Internal compute precision: for an F16 or Q8_0 cache the CUDA path runs the score and value GEMMs with F16 inputs (the F32 query range-scaled to F16, softmax weights stored F16) and F32 accumulation; an F32 cache keeps the whole chain F32. The fused/unfused perplexity comparison below bounds the effect. The op's op_params join the CUDA-graph reuse fingerprint (same treatment as GGML_OP_SCALE).

Routing in build_openpangu.cpp: the dense/SWA/MTP full-span sites and the DSA decode site build the op when a per-layer capability gate passes; otherwise they build exactly main's unfused chains, which remain intact as the fallback. The gate is one check per site: a probe node tested with ggml_backend_supports_op against the scheduled backend of the layer's attention projection, refusing the CPU backend by policy, and requiring the latent cache resident on that same backend, so --no-kv-offload keeps the unfused chain instead of re-uploading the cache view every graph. For an A/B against the unfused path, make openpangu_fused_attn_enabled() return false, the same condition flip used for the #2159 ablation; there is no build switch and no env var. The chunked-prefill score buffering is retired on op-served paths.

On the ggml_flash_attn_set_prefix alternative sketched in #2159: I read the DSA-in-FA path before deciding. The API shape would be natural (the src[4] sink / src[5] indexer attachment precedent), but the cost is in the kernels, and two things block it: (1) the prefix rows are learned weights with independent K and V, so they can neither be represented as cache rows nor ride the sink mechanism, whose sinks are per-head scalar logits; serving them through FA means second-source score and value accumulation in every kernel variant openPangu touches, and FA exposes no LSE for a post-hoc merge. (2) The measured risk runs the wrong way: my earlier openPangu FA experiments traded TG away for PP, and the #2159 ablation shows this op carries the TG win. fattn has no F32 K/V path, while the op reads F32/F16/Q8_0 in place, which only matters for -ctk f32.

DS4's FA serves it correctly by construction (scalar sinks, top-k as cur->src[5], F16 cache), so neither blocker exists there and nothing asks DS4 to move; the dedicated op exists for the one corner FA cannot express, independently learned prefix K/V rows attending jointly with the latent cache. The op carries dv_off so a DS4 trial stays expressible later. The indexed mode is close in shape to the dsa_attn gather, and a later convergence of openPangu decode onto that machinery with prefix-row support could collapse the two implementations.

CPU-only inference: as noted in #2159, this PR does basically nothing there, by design. The builder gate adopts the op only on a non-CPU backend, so CPU-only builds keep the vectorized unfused chain and the expected delta is zero.

Validation (RTX 4070 12 GB + i7-11700K, 64 GB RAM, CUDA build, openPangu-2.0-Flash Q4_K_M, -ngl 999 -ot exps=CPU -fa off -ctk q8_0 -ictk q8_0 -fidx -dsatk 64; every item below was run at this PR's head, raw artifacts retained):

• Graph placement: a full GGML_SCHED_DEBUG dump shows every LATENT_ATTN node (184/184 at -c 2048) scheduled on CUDA0 with zero CPU placements, sources verified as q, the raw latent-cache view, both param_sink tensors, and the mask (dense mask on dense layers, SWA-windowed view on SWA layers). The -fidx indexer top-k nodes from indexer_topk: fix quantized q8_1 scratch sizing on CUDA #2158 compose underneath unchanged.
• Gate refusals, verified by sched dump: -ngl 0 (CPU-resident layers) and --no-kv-offload both build zero LATENT_ATTN nodes and keep the exact unfused chain.
• Coherence: greedy completions at -c 8192 are clean and coherent, and byte-stable across repeated fused runs. The fused continuation diverges from the unfused one at a near-tie token, which is the expected effect of the F16 internal compute disclosed above; the perplexity comparison below is the quantitative bound.
• Flag coverage: a default-cache run (no -ctk/-ictk, F16 latent cache) and a run without -fidx (unfused indexer scoring feeding the same selection path) both produce clean, factually correct greedy output, and the cache-type flags verifiably change the load (K f16 103.50 MiB vs q8_0 54.98 MiB at -c 2048).
• Long context: a 32,541-token wikitext haystack with a planted needle at depth, -c 34816; greedy decode returns the needle verbatim as the first generated tokens through the fused path.
• Perplexity, fused vs unfused at the benchmarked config (IQ4_NL model, q8_0 caches, full wiki.test at n_ctx 2048, 150 chunks): fused 7.4359 +/- 0.0597 vs unfused 7.4089 +/- 0.0593, a +0.36% difference, within the reported error bars. Greedy outputs can fork at near-tie tokens accordingly; both continuations stay clean, and the fused output is byte-stable across repeated runs.
• Load-log sanity at -c 34816 (llama-cli defaults, -ub 512): KV 1009.71 MiB (q8_0) + 72.25 MiB indexer, CUDA0 compute buffer 301.00 MiB, 9989 graph nodes / 90 splits.
• bulgaria / quicksort / sayap / double-link-list all generate as expected.

Performance (warm A/B/A/B vs main e5357286c, same binary pair, per-N_KV medians of same-arm pairs, quiet host):

Model: openPangu-2.0-Flash IQ4_NL imatrix (ji-farthing/openPangu-2.0-Flash-ik-llama-GGUF), -ngl 999 -ot exps=CPU -fa 0 -ctk q8_0 -ictk q8_0 -fidx -dsatk 64 -t 8 -c 34816 -b 2048 -ub 2048 -n 512, llama-sweep-bench. Worst spread between same-arm pairs was 0.27% on S_PP and 2.02% on S_TG.

N_KV	main S_PP	PR S_PP	∆PP	main S_TG	PR S_TG	∆TG
0	268.40	292.06	+8.8%	15.32	17.78	+16.1%
2048	272.07	297.50	+9.3%	13.62	17.70	+29.9%
4096	267.41	292.25	+9.3%	13.60	17.64	+29.7%
6144	266.46	291.24	+9.3%	13.48	17.46	+29.5%
8192	263.77	287.72	+9.1%	13.43	17.36	+29.3%
10240	258.90	282.31	+9.0%	13.33	17.16	+28.7%
12288	257.54	280.62	+9.0%	13.35	17.21	+28.9%
14336	255.35	278.50	+9.1%	13.32	17.18	+29.0%
16384	251.82	274.24	+8.9%	13.25	17.13	+29.2%
18432	250.92	272.86	+8.7%	13.20	17.04	+29.1%
20480	247.49	268.73	+8.6%	13.17	17.00	+29.1%
22528	244.09	264.52	+8.4%	13.16	16.95	+28.7%
24576	242.58	262.31	+8.1%	13.12	16.84	+28.4%
26624	240.44	259.89	+8.1%	13.05	16.81	+28.8%
28672	237.98	258.12	+8.5%	13.00	16.73	+28.7%
30720	236.13	254.72	+7.9%	12.96	16.69	+28.7%
32768	232.52	250.74	+7.8%	12.94	16.62	+28.4%

Buffer accounting at this configuration, from the two load logs: the KV layout is identical in both builds (CUDA0 KV buffer 1009.71 MiB q8_0 plus the 72.25 MiB indexer), while the CUDA0 compute buffer drops from 2808.27 MiB on main to 1204.00 MiB with the op (graph nodes 14177 to 10565, splits unchanged at 90). Against that reduction, the op's transient pool scratch peaks near 135 MiB here: the capped 64 MiB score tile, its 32 MiB F16 weights companion, and a 38 MiB Q8_0 dequant window, all reused across layers by the pool allocator.

• I have read the contributing guidelines
• Self-reported review complexity:
  • Low
  • Medium
  • High

[ikawrakow]

On the ggml_flash_attn_set_prefix alternative sketched in #2159: I read the DSA-in-FA path before deciding. The API shape would be natural (the src[4] sink / src[5] indexer attachment precedent), but the cost is in the kernels,

I didn't say that you need to try to squeeze the prefixed FA into the regular DSA path. The main purpose of handling it via the existing GGML_OP_FLASH_ATTN_EXT is to avoid the noise added by a new op. But once you arrive at the call site, you can of course call a fully independent implementation.

But OK, let's do it that way.