Repository navigation
gemma4: graph optimizations (massive gains on pp for E2B and E4B) - #31356
Conversation
|
@ngxson @ggerganov |
| // KV-shared layers do not write to the KV cache, so prune non-output token rows before the first KV-shared layer | ||
| bool prune_shared = crop_before_nextn(inp_out_ids) && n_outputs < n_tokens && cparams.causal_attn && | ||
| hparams.n_layer_kv_from_start > 0 && hparams.n_layer_kv_from_start < n_layer && | ||
| !cparams.embeddings && inp_attn->get_kq_mask()->ne[3] == 1; | ||
| for (int il = (int) hparams.n_layer_kv_from_start; prune_shared && il <= n_layer; ++il) { | ||
| prune_shared = !cparams.embeddings_layer_inp[il]; | ||
| } |
There was a problem hiding this comment.
When the MTP is enabled, there will be no benefit because we still need to extract the full embeddings.
Also, this will create branches in the compute graph when using split KV cache. Can we avoid the inp_attn->get_kq_mask()->ne[3] == 1; condition somehow?
There was a problem hiding this comment.
Ah. That's one of the use-cases I missed.
Will play with that and see if we can avoid branches.
|
Oops. Clearly my coding skills are a bit better than my searching skills :). Not sure how I missed that. |
c2f9ab8 to
f81e718
Compare
|
@ggerganov the MTP issue should be resolved. @am17an can you please check it out. Here are some numbers on S26+ MTP No-MTP And the M4 MTP No-MTP Working on the split-KV issue next |
|
@ggerganov @am17an @ngxson Basically, we get full pruning gains for E2B and E4B with and without MTP with the unified-kv (no graph branches/reallocs). With split-KV we disable the pruning. Please let me know how this looks. Feel free to push updates directly if you want to tweak something. Here is a little test report on M4 Pro with Metal backend (tested hexagon as well)Tested short prompts and my usual ~16K sample prompt (13,789 tokens) with No graph reallocs across all 27 prompt ubatches (n_tokens = 512, n_outputs = 0 for the first 26 ubatches and n_outputs = 1 on the final ubatch) and the transition into t.g
|
|
@ggerganov @am17an sorry for bugging you guys. Just wanted to bump this in your feeds while it's still fresh. |
Picked from ggml-org#31356 (open at the time of picking), head df6b479, net diff against its merge base 792fcf4: 75e417d Optimizing Gemma4 graph for prompt processing 8485adb Not filtering the rows when they are unmasked 4f70ec7 Cropping out the attention mask too. 1ef40ac gemma-4: improved pruning, mask cropping and PLE scaling f81e718 spec-mtp: Enable masked h_nextn extraction only when the draft context shares the target K 99444d3 spec-mtp: make sure to not drop first draft round 7dadb47 gemma4: enable pruning based on static cparams.kv_unified/n_seq_max=1 df6b479 gemma4-mtp: simplify prunning logic and avoid graph branches/re-allocs Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com> Merge notes: common/speculative.cpp keeps the fork's draft-context setup (nextn embeddings only when n_max != 1), moved below is_mem_shared like the PR's target-context call. Assisted-by: Claude (Opus 5.5), conflict resolution
ggml-org#31356 (picked above) returns the permuted per-layer input tensor without ggml_cont(). Each layer then reads a strided 2D slice whose byte span covers nearly the whole tensor. When the layer runs on a different backend than the tensor (e.g. -sm layer over two GPUs, or the projection staying on the CPU input layer), the scheduler copies the whole view, so each such layer transfers the inputs of all layers: about n_layer times the traffic and a compute buffer about n_layer times larger for these copies. Unified-memory backends, where the PR was measured, don't copy and so don't see it. Restore the ggml_cont() that master has. The slice helper from the PR works on the contiguous tensor unchanged, and the row pruning and PLE scale folding are not affected. Assisted-by: Claude (Opus 5.5)
df6b479 to
1ed6351
Compare
|
I just realized that latest master broke this branch. |
There was a problem hiding this comment.
changes here seem unrelated, should it be a follow-up?
There was a problem hiding this comment.
Technically needed to make speculative-simple behave like spec-common in this case but I yeah the change is kind of hacky (I didn't like adding that include for sure). Reverting for now.
|
Fixed compiler errors and pruning gates after #29510 |
|
metal failures with |
am17an
left a comment
am17an
left a comment
There was a problem hiding this comment.
I think deepseek 4.1 would have a similar issue, will take a look.
max-krasnyansky
added
the
merge ready
|
@ggerganov please review/merge when you get a chance |
|
Will take a look today |
|
@max-krasnyansky Please rebase now that #31570 is merged. |
b61ab2b to
7fff406
Compare
|
#31558 changed the api again and broke this :) |
|
I think need to directly use the |
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.
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.
…t 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.
… now be used for multi-seq
219ada5 to
0971a7a
Compare
| // KV-shared layers do not write to the KV cache, so prune non-output token rows before the first KV-shared layer | ||
| bool prune_shared = crop_before_nextn(inp_out_ids) && | ||
| hparams.n_layer_kv_from_start > 0 && hparams.n_layer_kv_from_start < n_layer && | ||
| !cparams.embeddings && cparams.n_seq_max == 1; |
There was a problem hiding this comment.
@max-krasnyansky You mentioned that the approach works with unified KV cache, but here you actually have a stronger restriction: cparams.n_seq_max == 1. So for example, the default server config uses 4 sequences with unified KV cache and it does not benefit from the pruning because of this check.
There was a problem hiding this comment.
@ggerganov
The earlier check was n_seq_max == 1 || unified_kv == true but it stopped working after #29510 got merged.
It probably worked earlier just because unified_kv was not used witn n_seq_max > 1, or perhaps I misread the code.
As things are right now it only works with the Single-seq configurations.
Are you OK merging with this restriction or should we try to make it work with Multi-seq first?
It seems pretty tricky :) and I'm not sure there is a good use-case for Multi-seq with these smaller models. It'd be very useful for 26B but all layers feed the KV$ in 26B so we can't prune anything.
There was a problem hiding this comment.
It's not very difficult to implement - I added a TODO in the code. If you prefer, you can merge it like this and resolve it later.
There was a problem hiding this comment.
Ok. I'd say let's merge to enable the default case, and I'll do another pass to resolve the TODO.
Approvals expired. You'd have to merge or approve :)
There was a problem hiding this comment.
Hi @max-krasnyansky ,
If you are not planning to implement that change in this PR, I would like to take it up. I was trying to understand a few changes that you made on top of my commits.
There was a problem hiding this comment.
Hi @max-krasnyansky , If you are not planning to implement that change in this PR, I would like to take it up. I was trying to understand a few changes that you made on top of my commits.
Go for it!
I think we're in a very good shape for the default use cases: single-seq with/without MTP.
Multi-seq support is what's missing. If you have cycles please give it a go.
Please tag me in the next PR just so that it shows up in my feed (thank you thank you).
[no ci]
…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>
…a61c46) The -X theirs replay of 190 commits onto master 162e2bb dropped upstream's side where both changed a hunk. This commit sets the 9 files that differ to the one-time merge resolution (73fa61c46), so the tip's tree equals it: ggml-sycl/fattn-mkl.cpp and fattn-onednn.cpp (ours: upstream's ggml-org#30761 fixes the same strides and per-sequence offsets as 8e35c17 / a4cc5c9; the replay applied both offsets, ib * nb3 on top of K_batch), ggml-sycl/dmmv.cpp (upstream's templated launchers, ggml-org#31325, without the three reorder launchers they replace), ggml-sycl/ggml-sycl.cpp (upstream's comm_exchange pointer, ggml-org#30956, kept next to our host exchange, which returns first for one context), ggml-sycl/CMakeLists.txt (upstream's pkg-config lookup, ggml-org#31076, our include line in its fallback), common/speculative.cpp (upstream's DFlash device features, ggml-org#31506, next to our pos_inject; the MTP nextn setters after is_mem_shared, ggml-org#31356), src/llama-context.cpp (ggml-org#31506's wait before our ubatch loop), src/models/dflash.cpp, tests/test-backend-ops.cpp. Assisted-by: Claude Opus 5.5
|
Nice! 
|
Yep. The gains are quite awesome. |
|
@angt what is this dashboard? |
|
@am17an a small dashboard I maintain to track ROCm perf: https://huggingface.co/spaces/llama-benches/amd |
…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)
Overview
While optimizing
ggml-hexagonwe realized that gemma-4 E2B and E4B still runs compute for all the layers that do not contribute to the KV$. This PR fixes that.Original work by @Meet91721 and @huangwei858
Drop unused prompt tokens before KV-shared layers. Since the upper layers (!hparams.has_kv(il)) reuse earlier layers' KV cache and write nothing to the cache, we prune non- output token rows (inpL, inp_pos, kq_mask, and inp_this_layer) right before the first KV-shared layer instead of waiting until the last layer.
During prefill, those layers become a no-op on intermediate ubatches (n_outputs == 0) and only run 1 token row (n_outputs == 1) on the final ubatch.
The layer trimming is support only for the unified-KV.
Simplify PLE (per-layer embedding) projection.
Fold the two PLE scaling constants into the RMSNorm eps, the 256-element norm weight, and tok_embd_scale, and slice the permuted 3D PLE tensor directly using its strides instead of making a contiguous copy (ggml_cont).
Additional information
The gains are quite huge for the E2B, a bit less impressive for the E4B and none for 26B (all layers feed KV$).
I tested all the common use-cases with this models: text-to-text, image-to-text, audio-to-text and the outputs are identical to the master branch.
Also tested 16K and 32K prompts with
Needle in a Haystackcontent and the outputs are good.Here are some perf numbers vs master with the official Google QAT Q4_0 models.
Requirements