From 4fbd0c441b57a607df8ba8d85c3027d5cc83f60f Mon Sep 17 00:00:00 2001 From: Gearstickle <32626086+Turbomen008@users.noreply.github.com> Date: Tue, 26 May 2026 06:16:49 -0700 Subject: [PATCH] fa: preserve early-termination, fix multi-slot correctness via union of masks (#1880) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fa: fix FlashQKV early-termination causing S=0 assertion with --parallel N>1 The backward-scan optimization in compute_helper/compute_helper_q checks only one mask position per k_step block on the last query row (q_step-1) to find where valid KV entries end. When q_step > 1 and different query rows have non-overlapping valid KV regions (multi-slot / --parallel N>1), the scan on the last row's mask can miss blocks that contain valid entries for earlier rows. This causes those rows to accumulate S=0, triggering the GGML_ASSERT(S > 0) in normalize_and_store_1row. Fix: remove the early-termination scan at all 4 sites and iterate all nk1/k_step blocks unconditionally. The mask already handles correctness: fully-masked blocks produce smax=-inf and skip V accumulation, so the performance cost is minimal for TG (small nq1) and acceptable for PP. Fixes #809 * fa: refactor multi-slot mask fix into mask_effective_nk1() helper Replace 4× inlined early-termination scans with a shared helper that computes the effective K boundary by scanning ALL query mask rows (union-of-masks). This is the minimal fix for multi-slot parallel inference where different slots have different sequence lengths. The helper returns the k_step-aligned boundary covering the longest active sequence across all rows, preserving single-slot performance (single row = same boundary as before). Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Turbomen008 Co-authored-by: Claude Opus 4.6 --- ggml/src/iqk/fa/iqk_fa_templates.h | 46 +++++++++++++++++------------- 1 file changed, 26 insertions(+), 20 deletions(-) diff --git a/ggml/src/iqk/fa/iqk_fa_templates.h b/ggml/src/iqk/fa/iqk_fa_templates.h index 3aaff08c..711ae7c0 100644 --- a/ggml/src/iqk/fa/iqk_fa_templates.h +++ b/ggml/src/iqk/fa/iqk_fa_templates.h @@ -32,6 +32,24 @@ namespace { +// Compute effective K boundary by scanning ALL query rows (union-of-masks). +// The original early-termination scanned only the last row, which is correct +// for single-slot (all rows have the same mask) but wrong for multi-slot +// parallel (--parallel N>1) where different slots have different sequence +// lengths and therefore different mask patterns. +// Returns the number of K elements to process (multiple of k_step). +inline int mask_effective_nk1(const char * mask, int n_rows, int stride_m, int nk1, int k_step) { + int ik_max = 0; + for (int j = 0; j < n_rows; ++j) { + auto Mc = (const uint16_t *)(mask + j * stride_m); + int ik = nk1 - k_step; + for (; ik >= 0 && Mc[ik] != 0; ik -= k_step); + ik += k_step; + if (ik > ik_max) ik_max = ik; + } + return ik_max; +} + struct BaseHelper { BaseHelper(const char * data, int stride) : data(data), block(data), stride(stride) {} @@ -1359,11 +1377,8 @@ void compute_helper(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, in KQHelper::convert(q_step, stride_q, q, q_f16); #endif auto mr = mask; - auto Mc = (const uint16_t *)(mr + (q_step - 1)*stride_m); - int ik = nk1 - k_step; - for (; ik >=0 && Mc[ik] != 0; ik -= k_step); - ik += k_step; - for (int k1 = 0; k1 < ik/k_step; ++k1) { + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { #ifdef __aarch64__ KQHelper::multiply_mask_kq(kh, Dk, stride_m, q_f16, mr, fms); #else @@ -1424,11 +1439,8 @@ void compute_helper_q(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, auto q8r = (typename HelperQ80R8::block_q8 *)qptr; HelperQ80::convert(q_step, stride_q, q, q8r); auto mr = mask; - auto Mc = (const uint16_t *)(mr + (q_step - 1)*stride_m); - int ik = nk1 - k_step; - for (; ik >=0 && Mc[ik] != 0; ik -= k_step); - ik += k_step; - for (int k1 = 0; k1 < ik/k_step; ++k1) { + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { HelperQ80R8::repack(k_step, kh.block, kh.stride, q8r8); KQHelper::mul_mask_kq(khr8, stride_m, q8r, mr, fms); fqkv.accumulate_qkv(vh, fms); @@ -1455,11 +1467,8 @@ void compute_helper_q(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, perf.accum_nolock(0, t1); #endif auto mr = mask; - auto Mc = (const uint16_t *)(mr + (q_step - 1)*stride_m); - int ik = nk1 - k_step; - for (; ik >=0 && Mc[ik] != 0; ik -= k_step); - ik += k_step; - for (int k1 = 0; k1 < ik/k_step; ++k1) { + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { #if FA_TIMING t1 = Perf::cur_time(); KQHelper::mul_mask_kq(kh, stride_m, q8, mr, fms); @@ -1977,11 +1986,8 @@ struct FlashAttnBF16 { perf.accum_nolock(0, t1); #endif auto mr = mask; - auto Mc = (const uint16_t *)(mr + (q_step - 1)*stride_m); - int ik = nk1 - k_step; - for (; ik >=0 && Mc[ik] != 0; ik -= k_step); - ik += k_step; - for (int k1 = 0; k1 < ik/k_step; ++k1) { + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { #if FA_TIMING //t1 = Perf::cur_time(); FlashQKbf16::multiply_mask_kq(kh, stride_m, q_bf16, mr, fms, perf);