From 7c79795196a49b9551cb5a3938e1eb59e0acd8a2 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Sat, 26 Sep 2026 07:53:39 -0300 Subject: [PATCH] zaya: correct CCA conv input with several sequences in a ubatch The conv input was built by transposing QK over all the ubatch's tokens and then reshaping to [T, n_qk, S]; after the transpose a channel's row runs across every sequence, so with n_seqs > 1 the reshape mixed channels and sequences. Split the sequences first, then transpose (as the Mamba graphs do). Seen as llama-perplexity at its default batch (4 x 512-token sequences per ubatch): ZAYA1-8B F16 on wikitext-2, 4 chunks, gave PPL 58,259 against 27.89 with one sequence per batch; now both give 27.8871 on the CPU (Vulkan 27.54 / 27.62). llama-server with parallel slots batches requests the same way. Co-Authored-By: Claude Opus 5.5 --- src/models/zaya.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/models/zaya.cpp b/src/models/zaya.cpp index e6deacffca9e..58e5bf1c9769 100644 --- a/src/models/zaya.cpp +++ b/src/models/zaya.cpp @@ -291,8 +291,10 @@ llama_model_zaya::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * qk_mean_k = ggml_scale(ctx0, ggml_add(ctx0, Qmean, Kpre), 0.5f); cb(qk_mean_k, "qk_mean_k", il); - ggml_tensor * QKraw_t = ggml_cont(ctx0, ggml_transpose(ctx0, QKraw)); - QKraw_t = ggml_reshape_3d(ctx0, QKraw_t, n_seq_tokens, n_qk, n_seqs); + // [n_qk, T, S] -> [T, n_qk, S]: split the sequences before transposing, or with more than + // one sequence in the ubatch a channel's row would run across all of them + ggml_tensor * QKraw_t = ggml_reshape_3d(ctx0, QKraw, n_qk, n_seq_tokens, n_seqs); + QKraw_t = ggml_cont(ctx0, ggml_transpose(ctx0, QKraw_t)); ggml_tensor * conv_input = ggml_concat(ctx0, conv_state, QKraw_t, 0); cb(conv_input, "cca_conv_input", il);