qwen4exp: sum the indexer heads by slices (#28023)

* qwen4exp: sum the indexer heads by slices

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.

* qwen4exp: drop the redundant cont on the indexer query

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
This commit is contained in:
Pascal
2026-09-01 06:23:59 +02:00
committed by GitHub
parent 458681e1d5
commit 09412af38a
+11 -4
View File
@@ -576,12 +576,19 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
// rectify each head dot product before the sum, as in the DeepSeek lightning indexer
// mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s
ggml_tensor * score = ggml_mul_mat(ctx0, pooled,
ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream));
ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h*n_tps, n_stream));
score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);
score = ggml_relu(ctx0, score);
score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
score = ggml_sum_rows(ctx0, score);
score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream);
// 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);
}
score = summed;
cb(score, "indexer_score", il);
// one value per block, so it is cheaper to bias here than after the cells are expanded