Repository navigation
llama + CUDA: flash_attn_ext_rows (fix unified kv for multi-seq) - #29510
Conversation
ea25637 to
7908e9e
Compare
| // same as ggml_flash_attn_ext, but each q slice attends to its own list of K/V rows | ||
| // q: [n_embd_k, n_batch, n_head, ne3] | ||
| // k: [n_embd_k, n_rows, n_head_kv, 1 ] shared by all q slices | ||
| // v: [n_embd_v, n_rows, n_head_kv, 1 ] | ||
| // mask: [n_kv, n_batch, ne32, ne3] | ||
| // kv_rows: [n_kv, ne3] I32, mask column i of slice i3 uses row kv_rows[i, i3] of k and v | ||
| // negative rows are padding, their mask entries must be -inf | ||
| GGML_API struct ggml_tensor * ggml_flash_attn_ext_rows( |
There was a problem hiding this comment.
Wouldn't it be better to keep the q, k, v and mask the same as they are without changes and just pass an extra kv_rows in addition to what already exists?
The Metal (and I believe the Vulkan too) implementation is already quite well optimized so there is little to no benefit from the new kv_rows and at the same time the restructuring of the mask does not allow to use the existing kernels which have a fast-path for skipping -INF blocks.
|
Hi @am17an , I have also been spending some time on kvu performance. ynankani@2e06853 covers indexed rollback, compact attention views, and CUDA masked-block skipping. Your The indexed rollback change uses the existing per-sequence position index to find cells for |
|
@ynankani you can probably benchmark and see, my intention with this PR is first fix the unified kv cache and secondly introduce the concept of indirection - ie rows are pages of size 1 at the moment, but that can change. I believe it will help mature the multi-seq setup quite a bit. |
|
@am17an I verified your PR closes the largely closes the gap between unified and non-unified KV. Tested on RTX 5090 and N1X. I had another PR that improved creating the attention masks (#29769) , could you confirm that I should keep #29769 closed in favor of your PR? Just not sure if you're moving this PR (#29510) from draft to production and what are your next steps. Let me know if you'd like to me to help on any more testing! |
|
@joelagnel that PR feels like a band-aid rather than a fix. I would much prefer some shape of this PR to land. |
|
@am17an Thanks for reply, Aman. I don't know if it's a bandaid in the sense it does mask setting faster if only fewer cells need update (it's only one commit of another larger branch I had doing kvu optimization). So it's an algorithmic optimization. I agree your approach is better and eliminates the need for that optimization. Do you have any thoughts on next steps for your PR and when you're expecting it to land? I was concerned your PR is in draft? Happy to help if you want me to work on some part of this. |
|
@joelagnel yes I mean the mask as it exists today will not be sufficient to close the gap. I'm ready to make this PR as ready to review, I don't have much to add here. The only reason it's draft because I wanted to get feedback from other maintainers and especially @ggerganov and @JohannesGaessler as to what we should do. I've been using this PR for agentic coding with pi subagents with kvu and it seems to be noticeably faster after ~30k context. |
|
It's a good direction, but I wanted to achieve some benefit in the Metal backend from this. So far haven't been able to do so - #29568. That's why I suggested to keep the mask as it is: #29510 (comment) I'll need to look more into this. It would be nice to have an efficient unified cache so the PR is definitely of interest. |
|
For metal I think we need to optimize bs=2..8 mmv kernels to see a big benefit in decode. For prefill I would think there is already sufficient mask elision. For vulkan I can see a benefit in decode even if it has a better mask algorithm. |
|
Yes, I meant the new approach should not at least make it worse, even if there is no benefit. My last tests showed that using kv_rows is slower compared to -INF skipping on the Metal backend. So this is irrelevant of the 2..8 mmv kernels. But I need to spend more time to see if I can make it work. |
7908e9e to
4de4f8f
Compare
…ggml-org#29510) Port of 1b6935b onto temp-11739. Fabric changes: qwen4exp QSA keeps the dense mask, no kv_rows in training, Metal merged into fabric's FA host code, Hexagon rejects kv_rows, dense-only masks are reused like before. (cherry picked from commit 1b6935b) Assisted-by: Claude Opus 5.5
…argument (ggml-org#31550) The kv_rows argument added to the FA kernels in ggml-org#29510 changes their mangled names, so the known spills were no longer ignored. Fabric's list has more entries than upstream's; all flash_attn_ext_f16/vec and flash_attn_tile entries are renamed. (cherry picked from commit 74101b3) Assisted-by: Claude Opus 5.5
…w SYCL PRs Move the stack from dde879bc1 to upstream fac3a89 (43 commits, merged into this branch by the previous commit). Fork, PR, extracted and local patches (40), the 13 candidates and the 4 Vulkan patches apply in script order with plain git apply, and the three unapply scripts return the tree to upstream. GitHub's pull/N.patch URLs return 403 through this environment's proxy, so the PR files were made from pull/N/head fetched with git: git format-patch for unpinned carries, git diff for re-anchored ones. Retired, merged upstream on 2026-11-02: SYCL PRs 30956 and 31325. New SYCL carries: - 31536: oneMKL FA with a mask broadcast over heads or batches. - 31542: attention sinks (gpt-oss, Flint) may take the oneMKL FA path. Both are re-anchored and pinned: their test-backend-ops additions meet upstream's new kv_rows FA test (ggml-org#29510) and PR 30606's GQA loop. Refreshed: SYCL PR 30606 (3 commits) and PR 31285 (reworked upstream to a BMG-only check; re-anchored after 30606's docs row and pinned), Vulkan PR 29357 (shader and host accuracy fixes, prefill FA SLM capped at 48 KiB). Skipped: 31540 (CI only). 31139 stays skipped: SYCL reports no allocation limit, so it changes nothing there. Re-anchored: fork 0002 and 0008 onto merged 31325's dmmv launcher templates, candidate C51 onto the CPU FA loop ggml-org#29510 changed. New opt-in candidates from SYCL drafts: - C62, PR 31574: oneDNN fused SDPA at head_dim 64 on Alchemist. - C70, PR 31379: the gated delta net reads its state through the GET_ROWS ids (re-anchored on PR 31355 and local 0015). - C75, PR 31313: lightning-indexer vector and XMX kernels, without the PR's fused-gemm.cpp dedup, which local 0010/0017 and C20/C21 edit. - C76, local: 256 registers for C75's vector kernel on ACM, where it spilled 1.6-1.8 KB per thread at 128. Upstream b74298c writes delta-net conv-state windows with one SET_ROWS, so carried PR 31355's CPY matcher no longer fires on Qwen3.5/Qwen3-Next; recorded in the README and dev log. Checked here: the apply/unapply round trip above, and acm-g10 AOT builds of lightning-indexer.cpp (C75+C76), gated_delta_net.cpp (C70), fattn-mkl.cpp (31536, 31542) and dmmv.cpp (0002, 0008), none spilling. CPU OpenCL device test results follow in the next commit. Nothing ran on an Intel GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018jZpMGd3MavhKeWoY4gTBL
…ort of ggml-org#29510) (#467) * llama : flash_attn_ext_rows - more efficient unified kv for multi-seq (ggml-org#29510) Port of 1b6935b onto temp-11739. Fabric changes: qwen4exp QSA keeps the dense mask, no kv_rows in training, Metal merged into fabric's FA host code, Hexagon rejects kv_rows, dense-only masks are reused like before. (cherry picked from commit 1b6935b) Assisted-by: Claude Opus 5.5 * ci : update the HIP VGPR spill ignore list for the kv_rows FA kernel argument (ggml-org#31550) The kv_rows argument added to the FA kernels in ggml-org#29510 changes their mangled names, so the known spills were no longer ignored. Fabric's list has more entries than upstream's; all flash_attn_ext_f16/vec and flash_attn_tile entries are renamed. (cherry picked from commit 74101b3) Assisted-by: Claude Opus 5.5 * metal : fix the kv_rows gather for head sizes 40/72, reject unpadded row lists The row gather copied chunks of 4 half4 over a flattened (head, value) index. With DK = 40 or 72 (10 or 18 half4 per head) a chunk clamped at the end of a head skipped the first values of the next head. The chunks now stay within a head; for head sizes that are a multiple of 16 the work split is unchanged. supports_op now rejects row lists whose length is not a multiple of the KV block sizes, so the scheduler falls back instead of asserting in the op. The KV cache pads the lists to 256, so llama.cpp inference is not affected. Adds test-backend-ops cases for both. Both issues are in upstream master too. Assisted-by: Claude Opus 5.5
…ml-org#31356) * Optimizing Gemma4 graph for prompt processing TODO: - Currently complete microbatch needs to contain tokens which does not require the logits; need to optimize this further - There are a few more operations that are performed after the kv cache is calculated, need to safely optimize them. * Not filtering the rows when they are unmasked * Cropping out the attention mask too. * gemma-4: improved pruning, mask cropping and PLE scaling Prune before the first KV-shared layer with the single-stream, causal, non-embedding safety gates. Use cropped masks only during shared-layer attention. Preserve requested layer-input embeddings in the KV prefix. Remove the PLE contiguous materialization and fold equivalent PLE scaling into RMSNorm parameters/weights. * spec-mtp: Enable masked h_nextn extraction only when the draft context shares the target KV cache This lets Gemma 4 prune unused prompt rows before its shared layers so that we keep the full perf P.P gains with MTP. Avoid extracting target hidden states for output-free prompt chunks in that shared-MTP path. When a batch has a selected final token, extract only that row unless the whole verification batch is selected. Keep ordinary MTP and classic speculative decoding behavior unchanged: they still use unmasked h_nextn because the draft context must catch up on all target rows. speculative-simple needs an update to select the final prompt row only for shared-KV MTP, so the optimization works there without changing other speculative modes. And it looks like the server already selects its final prompt row for first-token sampling. * spec-mtp: make sure to not drop first draft round * gemma4: enable pruning based on static cparams.kv_unified/n_seq_max=1 * gemma4-mtp: simplify prunning logic and avoid graph branches/re-allocs * gemma4: update pruning gates to use only n_seq_max==1, kv_unified can now be used for multi-seq * spec-simple: revert spec-simple changes for now * gemma4: remove causal_attn from pruning gate (last source of reallocs) * cont : adapt to llm_attn_mask changes * cont : add TODO --------- Co-authored-by: Meet Amin <meetamin378@gmail.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com> (cherry picked from commit 6cb3e67; the shared-layer mask crop uses self_kq_mask_cnv / self_kq_mask_swa_cnv, since this branch predates the llm_attn_mask change of ggml-org#29510)
The fused quantized-KV flash-attention kernel assumed the KV rows it needs are always contiguous. Upstream's flash_attn_ext_rows work (ggml-org#29510) broke that assumption: with kv-unified + multi-slot serving, each request's rows are gathered through a row-index list, and the fused kernel ignored the list and read consecutive rows instead. Symptom: in multi-slot serving the kernel attended to other sequences' cache rows - logits looked plausible but were wrong, producing repetition loops and premature stops (4 of 12 long-generation runs on Qwen3.6-35B, np4; 0 of 12 with the fused kernel off). Single-slot traffic was safe by luck: there the row list happens to be contiguous, which is why all earlier np1 tests matched master bit-for-bit. Fix: gather K/V rows through the index list exactly like upstream's unfused loader (clamping negative padding to zero), at all three load sites. Verified bit-exact vs the unfused kernel across 8 op shapes, and 12/12 multi-slot requests clean end-to-end. Memory win kept (660 vs 1140 MiB).
The fused quantized-KV flash-attention kernel assumed the KV rows it needs are always contiguous. Upstream's flash_attn_ext_rows work (ggml-org#29510) broke that assumption: with kv-unified + multi-slot serving, each request's rows are gathered through a row-index list, and the fused kernel ignored the list and read consecutive rows instead. Symptom: in multi-slot serving the kernel attended to other sequences' cache rows - logits looked plausible but were wrong, producing repetition loops and premature stops (4 of 12 long-generation runs on Qwen3.6-35B, np4; 0 of 12 with the fused kernel off). Single-slot traffic was safe by luck: there the row list happens to be contiguous, which is why all earlier np1 tests matched master bit-for-bit. Fix: gather K/V rows through the index list exactly like upstream's unfused loader (clamping negative padding to zero), at all three load sites. Verified bit-exact vs the unfused kernel across 8 op shapes, and 12/12 multi-slot requests clean end-to-end. Memory win kept (660 vs 1140 MiB).
…06) into the P100 fork Conflicts and what was kept: - KQ mask: upstream's llm_attn_mask / kv_rows refactor (ggml-org#29510). Our compact I32 causal mask is handled in llama_kv_cache_context::set_input_attn_mask; the single-sequence fast mask fill only runs when the mask columns are the cells (no kv_rows). kv_rows attention is MMA-only, so on sm_60 the probe reports it unsupported and the dense/compact mask path is used as before. - MoE layout: took upstream's linearized MMQ/MMF MoE data layout (ggml-org#31047) in mmid.cu, mmf.cuh, mmq.cuh. It replaces our repeated-expert-id handling there (63924ad); upstream keeps those test cases disabled and top-k routers never repeat an expert. The cuBLAS sort path keeps its fix. The opt-in fp16 MoE GEMM (GGML_CUDA_P60_MOE_GEMM=1) now builds its jobs from the helper's tile list; the math is unchanged. - mmq.cu: upstream applies the tokens-per-expert tile sizing to every arch (ours did it on sm_60 only). - scale.cu: upstream's strided rewrite plus our scale == 0 -> bias rule. - common_state_data (upstream) is an alias of our common_ckpt_buffer, so state buffers keep the pooled, huge-page, no-zero-fill allocator. Checkpoints are ref-counted upstream now (ggml-org#30137): our keep-the-just-restored-checkpoint path claims it through the const pointer, except a slot-file one. - MTP: upstream's mem-shared (gemma4) verify rows merged with our verify_pos0/pending_pos bookkeeping; upstream's own MTP get_state/set_state dropped in favour of ours (position + row, resets stale state). - fattn: the q4p kernel takes upstream's new KV_rows argument (unused). qwen35 MTP K/V-only path uses the per-layer K/V rotations (ggml-org#31073). Math: with upstream's softplus log1pf change (74bca39) reverted as a diagnostic, logits match the pre-merge build exactly (27B Q6_K -c 4096 x 8: KLD -0.000006, same top 100 %; -ub 1: KLD 0). Kept: it moves the Qwen3.8 GDN gate (KLD 0.001-0.003) and reads no worse against the fp32 reference (Q5_K_M -ub 1: KLD 0.001713 vs 0.001779). Upstream's q8_1 correction-sum fix (8fc0376) moves only the MTP head, whose weights are Q4_0: greedy-draft MTP text and acceptance are identical to the base, the draft probabilities move, and the quantized sum is the more exact one (simulated Q4_0 x q8_1 dot: RMS error 0.58 vs 1.44). Kept. OPTLOG 329. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01AaVsfEiPALnEHpNH9UZbjk
- sm_60 MUL_MAT_ID kernels: keep the per-expert-bounds ids helper (mmid-bounds.cu) that upstream ggml-org#31047 replaced by tile specs - deepseek41: the graph on upstream ggml-org#31733 (model member, no_build base ctor), the llm_attn_mask raw mask, build_attn_mha's kv_rows argument - llama-context: graph_params initialiser; the DFlash device-copy reader event (ggml-org#31506) is recorded in decode_finish (a launched-ahead decode finishes in llama_decode_commit) - kv-cells: seq_rm_all clears the sequence's seq_bits (ggml-org#29510) - graph slots debug: llm_attn_mask; model saver: DEEPSEEK41 in the DSV4 condition Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… derived kq mask Re-based onto a691b23. Upstream ggml-org#29510 (flash_attn_ext_rows, the kv_rows unified-KV form) and the llm_attn_mask refactor meet this block in 9 files / 48 hunks. Resolutions worth knowing: - the kernel signatures carry both parameters (upstream's `KV_rows` then our `kq_blocks`); the loaders take upstream's kv_rows `indices` with our native-KV arguments; the C-tile sizing keeps the delivery's fp32-aware expression under upstream's renamed `AMD_WMMA_256B_AVAILABLE` (256B == RDNA3). - the derived kq mask MOVED ABI: upstream's kv_rows took src[5], which the derived form had been using, so cell_pos/tok_lo/tok_hi are now src[6], src[7], src[8] (ggml.c, ggml-cpu/ops.cpp, the CUDA launcher, and both fattn.cu predicates). Before the move the derived probe resolved to "not supported" and every batch silently fell back to the packed mask. - the rotation constants are no longer filled per ubatch: upstream ggml-org#31073 keeps them on the GPU and consumes them at graph-build time (get_k_rot()), so the block's set_input_k_rot/v_rot calls and their declarations (removed from the base by upstream) are dropped. - input classes use upstream's llm_attn_mask (raw/rows) instead of the separate kq_mask/kq_mask_cnv members; the derived registration hangs off the dense branch of build_attn_inp_attn_mask, which is now a graph-context member so it can add inputs. The two compact forms are mutually exclusive and the flash-attention build emits one op (kv_rows XOR derived XOR plain). Validated: build green; derived kq mask enabled ("derived kq mask flash attention enabled"); derived vs LLAMA_KQ_MASK_DERIVED=0 byte-identical at a 4096-ctx prefill (149 chars, sha 30c7fc8705a6); q8_0 KV runs with flash attention.
…r, QSA/indexer) Re-based onto a691b23: 9 files / 34 hunks (the probe predicted 7 files; the cascades were llama-memory-recurrent.cpp and tests/CMakeLists.txt). Decision D3 was measured, not assumed -- see below; this commit keeps the DELIVERY's implementation. Base-API drift this rung had to re-derive (all compiler-driven, no behaviour change): - llama-kv-cache: the ctor keeps BOTH indexer flags -- upstream's `is_indexer` (indexer RoPE params) and the delivery's `v_enabled` (keys-only caches) -- and drops the delivery's `n_pad` ctor parameter and member: upstream ggml-org#31200 removed them and hardcodes the 256-cell floor in the size math, so only the derived `n_pad_prefill` (LLAMA_KV_N_PAD_PREFILL / _MIN / _FROM) survives. `get_n_kv()` now REQUIRES n_tokens (no silent default: a 0 would turn the prefill pad off at that call site) and llama-kv-cache-dsv4 passes it. - llama-memory-hybrid-idx: the keys-only indexer cache is built with upstream's post-ggml-org#31141 ctor (no_alloc, no n_pad) plus is_indexer/v_enabled; the pool body is the delivery's. - llama-memory-recurrent: R04 wins (upstream's rs_max bound, the tracker stays dropped) and the block's two real additions are kept -- the getenv-once `rs_rollback_debug()` helper (issue ggml-org#273) and the `has_recr` gate that stops the qwen4exp MTP context (whose recurrent filter matches nothing) spamming the non-consecutive-position warning. - models.h / qwen4exp.cpp: ggml-org#31733 made `model` a member of llm_graph_context, so the graph and graph_mtp ctors lose their `const llama_model &` argument (and the class its shadowing `const llama_model & model` member); ggml-org#31073 turned the rotation constants into per-layer vectors (`self_k_rot[il]`) and ggml-org#29510/ggml-org#31683 replaced `get_kq_mask()` with `self_attn_mask.raw`; build_attn_mha grew `n_kv_max`. - indexer-score.cu: `ggml_rope_yarn_corr_dims` gained the truncate flag. - tests/CMakeLists.txt: R03 had already re-placed every registration this block adds (ggml-org#30880 moved the generated-model fixtures), so only `test-lazy-reader` is new; test-llama-archs keeps upstream's get_model/get_model_and_ctx split with the block's `lazy_buf_size` handed over as LLAMA_LAZY_BUF_MB, and test-backend-ops keeps BOTH sides' Flash-Attention classes (upstream's kv_rows and the delivery's QSA) and both perf blocks. D3 (qwen4exp / indexer / sparse selection) -- MEASURED, decision: take the delivery's implementation. Upstream now carries its own qwen4exp (1514 lines, by upstream contributors: ggml-org#30965, ggml-org#31243, ggml-org#31443, ggml-org#30927, ggml-org#31086) with its own HC mix/combine (plain ops), QSA selection (build_qsa_sel, an additive mask) and indexer (plain ops); the delivery's is 3197 lines with fused HC ops (HC_MIX/HC_COMBINE, HC16 planar layout), the QSA store+top-k selection (no mask), INDEXER_SCORE/TOPK/FILL and a keys-only indexer cache. A/B on the same GGUFs, same build: Flash-Next GSQ-IQ3_XXS, 3 GPU -sm tensor, -c 8192, 5256-token prompt upstream pp 1399.7 t/s tg 53.5 t/s text 7a8734a9f1a6 delivery pp 1432.1 t/s tg 70.1 t/s text 7a8734a9f1a6 <-- identical Qwen3.8-27B UD-Q4_K_M, 1 GPU, same prompt upstream pp 1311.1 t/s tg 32.4 t/s text b905c9763a58 delivery pp 1303.2 t/s tg 32.4 t/s text b21e94a5b6e1 (neutral; text diverges, needs the KLD gate to judge) So on the config this repo actually validates the delivery's implementation is BIT-IDENTICAL to upstream's and +31 % decode (+2.3 % prefill); on the 27B it is throughput-neutral and its greedy text diverges at 16 tokens (not a purity claim either way -- the chain differs, and the R17 KLD gate is the oracle). Keeping the delivery's implementation also keeps r235's qwen4exp validation and the two delivery gates (gate-hc-fusions.sh, gate-qwen4exp-quant-coherence.sh) meaningful. Validated: build green, no `<<<<<<<`; dense 4B (1472ebdadd46) and Qwen3-Coder-30B-A3B (013d09f165a5) unchanged; Flash-Next and 27B run to completion on the VMM build; gate-hc-fusions.sh PASS (5/5 hard arms byte-identical to the chain, 449 chars d62340fbf566); gate-qwen4exp-quant-coherence.sh PASS (IQ3_XXS / IQ4_NL / Q4_K_M: gather on == off, no '////'). Known drift to record at R17: on this tree the two REPORTED dsv4_pre/post arms no longer agree with the chain on the BF16-hc class (chain 449 chars d62340fbf566, dsv4 arms 6a102dbc62d6 = the value the chain produced at r235), i.e. the HC chain itself moved with the re-base while the fused pre/post ops still reproduce the old one. The gate reports (does not fail) them, and its header says byte-identity IS expected on this class -- the soft-arm expectation needs a decision.
Overview
During multi-seq inference, unified kv-cache pays a big penalty due to having to filter out other sequence tokens from the mask. This PR introduces an extension to flash attention (
flash_attn_ext_rows) which when supported by a backend allows llama.cpp to send an optimized representation of the query. Queries shared the kv rows and their mask + a new data structkv_rowsenables an indirect look-up for a compact representation of the required tokens per seq.This PR is approximately on par with the non-unified kv cache (
-nkvu) for performance while keeping the benefits of the unified kv-cache.Tested on Qwen3.8-27B and Gemma-26BA4B. On RTX Pro 6000:
llama-batched-bench -m ~/workspace/gemma-4-26B-A4B-it-Q8_0.gguf -ngl 99 -fa on -c 400000 -b 2048 -ub 2048 -npp 99936 -ntg 64 -npl 4 -kvuMaster:
PR:
llama-batched-bench -m ~/workspace/Qwen3.8-27B-Q8_0.gguf -ngl 99 -fa on -c 400000 -b 2048 -ub 2048 -npp 99936 -ntg 64 -npl 4 -kvuMaster:
PR:
Additional information
Requirements