From a4e0ee7dbea56fd58769ba7604a5055f01fa766c Mon Sep 17 00:00:00 2001 From: Masahito Suzuki Date: Tue, 8 Sep 2026 00:17:00 +0900 Subject: [PATCH] qwen4exp: drop the 1/r scale on the QSA block mean Its only consumer is the RMS norm below it, and RMS norm is scale invariant: rms(x*s) = x*s / sqrt(mean(x^2)*s^2 + eps) = x / sqrt(mean(x^2) + eps/s^2) so dropping the divide only moves the effective epsilon from eps to r^2*eps -- 1e-6 to 1.6e-5 against a mean square of order 1. The scale was a full read and write of [idx_dim, n_blocks] f32 per layer per ubatch: 2.25 ms of a 97 ms decode step at 131k context. This is the one patch in the series that is not bit-exact by construction. Measured on Strix Halo (gfx1151, Vulkan), ctx 262144, f16 KV, MTP n_max=3: 131072 goes 28.26 -> 32.46 t/s cumulative with the previous patch. --- src/models/qwen4exp.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index e9f23af6ac9b..03d73506ea3b 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -713,8 +713,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( members->nb[2], members->nb[3], i*members->nb[1]); fresh = fresh ? ggml_add(ctx0, fresh, slice) : slice; } - fresh = ggml_scale(ctx0, fresh, 1.0f/(float) r); - cb(fresh, "indexer_k_pooled", il); + // no ggml_scale by 1/r here: the only consumer is the RMS norm below, and RMS norm is + // scale invariant. rms(x*s) = x*s / sqrt(mean(x^2)*s^2 + eps) = x / sqrt(mean(x^2) + eps/s^2), + // so dropping the divide only moves the epsilon from eps to r^2*eps -- 1e-6 to 1.6e-5 against + // a mean square of order 1. the scale was a full read and write of [idx_dim, n_blocks] f32 per + // layer per ubatch: 2.25 ms of a 97 ms decode step at 131k context. + cb(fresh, "indexer_k_sum", il); // count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536 fresh = ggml_reshape_3d(ctx0, fresh, idx_dim, n_recomp*n_stream, 1);