qwen4exp: sum the indexer heads by slices - #28023
Conversation
The head reduction went through a transpose and a sum_rows over ne[1], which left sum_rows with ne0 = 4, one block per row for a four element reduction, and the transpose copied the whole block by token surface twice on the way in. The heads are adjacent on ne[1], so each one is a strided view and the sum is a short chain of adds. RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, 55k context, warm runs on top of ggml-org#28011: prompt processing 2170 -> 2366 t/s Generation is unaffected. The removed work scales with n_blocks by n_tokens, so the gain grows with context and with ubatch size.
9180166 to
ead00ae
Compare
|
(rewording commit) |
| // the heads sit side by side on ne[1] and there are only a few of them | ||
| ggml_tensor * summed = nullptr; | ||
| for (int64_t h = 0; h < n_idx_h; ++h) { | ||
| ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream, | ||
| score->nb[2], score->nb[3], h*score->nb[1]); | ||
| summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice); | ||
| } |
There was a problem hiding this comment.
instead of this thing, can we add an operation in the backends which are able to sum along as axis, similar to pytorch sum(axis=), it would be make a lot of permute->cont->sum_rows operations better supported. cc @ggerganov
There was a problem hiding this comment.
Not required for this PR though, just a suggestion
There was a problem hiding this comment.
Sounds good, I've added ggml ops before and I'd be happy to do this one as a follow up, with the tests and the backends. I'm getting the M5 shortly for Metal, and the CI should cover me for AMD.
It would clean up the permute -> cont -> sum_rows pattern in several models here. And reducing along ne1 with a large ne0 is the good shape on GPU anyway, so it should beat the transpose on top of removing the copy.
There was a problem hiding this comment.
Before adding an op, take a look if the graph is optimally constructed. Often cases, having to do such permute->cont->sum_rows means that earlier in the graph the data wasn't arranged properly. I'm not sure if this is the case here, but I would first look for that.
There was a problem hiding this comment.
Had a look: the layout is already fine, the matmul puts the heads on ne[1] with n_blocks as the fast axis so the slices are plain views, and what forces the reduction is the relu between the matmul and the sum.
I also tried dropping the redundant ggml_cont on q before the matmul, since rope already returns it contiguous. It works, but alternated A/B in both orders shows no measurable difference on prefill, the run to run drift is larger than the gap, so I left it out.
On the op: with only 4 heads a chain of adds does the job here, so it would mostly be for readability. The cases worth looking at are the models where the reduced axis is big enough that a chain of adds stops being reasonable, since that is where we are forced to transpose today.
There was a problem hiding this comment.
I also tried dropping the redundant ggml_cont on q before the matmul, since rope already returns it contiguous. It works, but alternated A/B in both orders shows no measurable difference on prefill, the run to run drift is larger than the gap, so I left it out.
It's better to remove the cont - it is a redundant op. Even if it is not measurable, there is no reason to have redundant nodes in the graph.
There was a problem hiding this comment.
Yes, same goes for this one also, since ggml_add allocates a contiguous result anyway:
- summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice);
+ summed = summed ? ggml_add(ctx0, summed, slice) : slice;
Greedy output is unchanged. It only holds because n_idx_h is 4 here, so the loop always runs at least one add: with a single head summed would stay a view, and the paths downstream would be the ones materialising it.
rope returns a freshly allocated, contiguous tensor, so the reshape that feeds the matmul does not need a copy. ggml_reshape_3d asserts contiguity, so a layout that would need the cont cannot slip through silently. Greedy output is unchanged token for token. Address review from @ggerganov
|
@ggerganov could you give guidance on how the |
|
For what it's worth, with 4 heads the chain of adds is already close to optimal, so I would not expect the op to be measurably faster in this particular spot, it would mostly be for readability. The gain should show up where the reduced axis is big enough that a chain of adds is not an option and we have to transpose today. I have a list of the models using sum_rows and I plan to go through them to find a case that actually justifies it before writing the op. |
|
@ServeurpersoCom it will show up on fast GPUs for sure. At least 1% |
|
We can add |
|
@am17an Can confirm: the more bottlenecks you remove, the more the remaining optimisations are worth, and the faster the GPU, the more any small bottleneck shows up as a penalty. I still have a few optimisations in stock to PR, split up as much as possible so there is time to test and review each one properly. |
That works, and keeping ggml_sum_rows as the dim == 0 case means nothing existing changes. Agreed on disallowing dim != 0 for quantized types: reducing along ne0 stays inside the quantization blocks, anything else crosses them and would need per element dequantization. All the permute -> cont -> sum_rows sites I found are f32 anyway. Q. One thing to confirm: the reduced dimension stays at 1 like ggml_sum_rows does, rather than being dropped? Answered myself: ggml_sum_rows sets ne[0] = 1 and keeps the rest, and since a tensor always has GGML_MAX_DIMS dims there is nothing to drop anyway. So ne[dim] = 1 with the rest unchanged, which makes dim == 0 identical to ggml_sum_rows. |
Yes, one can always reshape if they need to drop the reduced dimension. |
|
I did a POC on DeepSeek-V4 to see what the op actually buys: ServeurpersoCom/llama.cpp@7da3e10. It drops two conts per site, but prefill and decode are unchanged on my RTX PRO 6000, and perplexity is the same within reassociation noise. The same pattern shows up in the glm-dsa, deepseek32 and dots3note indexers, though those sit in the else branch of the fused lightning indexer and never run on the backends that have it. So on the model side it is mostly a readability win. Where it looks more promising is audio models, which reduce along the time axis all the time and pay a transpose in and out every time. |
4 participants
Overview
Split out of #27977 as requested, one PR per change. This one is prompt processing only, and it grows with context and ubatch size.
Summing the indexer heads went through a transpose and a reduction over a dimension of 4, which is about the worst shape for that kernel, and the transpose was copying a large surface twice for nothing. The heads sit next to each other in memory, so a short chain of adds over strided views does the same thing without moving anything.
Additional information
The head reduction went through a transpose and a sum_rows over ne[1], which left sum_rows with ne0 = 4, one block per row for a four element reduction, and the transpose copied the whole block by token surface twice on the way in.
The heads are adjacent on ne[1], so each one is a strided view and the sum is a short chain of adds.
RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, 55k context, warm runs on top of #28011:
prompt processing 2170 -> 2366 t/s
Generation is unaffected. The removed work scales with n_blocks by n_tokens, so the gain grows with context and with ubatch size.
Requirements