From 2b61da06023c8cf861bfd340b4ad33a724733993 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Wed, 8 Jul 2026 09:42:40 +0100 Subject: [PATCH] fix(encoder): keep padded query rows finite in limited-context attention masks In a heterogeneous-length batch on a chunked_limited / local-attention model (e.g. nemotron-3.5-asr-streaming-0.6b), any clip whose padding tail exceeded the attention left context (~4.5 s of length spread at 8x subsampling) decoded to an EMPTY transcript. Mechanism: a deeply padded query row has its ENTIRE additive mask row at -inf, softmax over all--inf yields NaN, the post-softmax query-row mask keeps it (NaN*0 == NaN), and the NaN key columns then poison every softmax row of the next layer -- 24 layers later the whole item is NaN. Full attention was unaffected (padded query rows can still see valid keys), which is why only limited-context models in batches with enough length spread were hit; the batch JSON entry points the LocalAI backend calls were silently affected. Leave fully-masked padded query rows unmasked instead: their finite garbage output is zeroed by the existing post-softmax query mask (output reduces to linear_out.bias, exactly as before), and valid rows' masks are untouched. Applied to every band/window mask builder (single + batched, full/chunked-limited + banded local + chunked kernel). Verified on nemotron-3.5-asr-streaming-0.6b q8_0 (att_context [56,3]): in a {16.4 s, 7.4 s} batch the short item's encoder output was 100% NaN before the fix and now byte-matches its solo decode; transcripts match solo at every tested T_max spread (8.4-16.4 s). Also verified on parakeet_realtime_eou_120m-v1 f16 (att_context [70,1]): pre-fix the 7.4 s fixture decoded empty in any batch with T_max >= 14.4 s, post-fix it byte-matches solo at every spread. Cache-aware streaming and single-clip offline transcripts are byte-identical between pre-fix and post-fix builds (those paths never form padded batches), so the fix changes nothing outside the broken case. Assisted-by: Claude:claude-fable-5 [Claude Code] --- src/relpos_attention.cpp | 60 ++++++++++++++++++++++++++++++++++------ 1 file changed, 51 insertions(+), 9 deletions(-) diff --git a/src/relpos_attention.cpp b/src/relpos_attention.cpp index ef6ffa1..c82038e 100644 --- a/src/relpos_attention.cpp +++ b/src/relpos_attention.cpp @@ -111,6 +111,11 @@ ggml_tensor* RelPosAttention::build_graph(ggml_context* ctx, ggml_tensor* xt, // Additive mask [T_k, T_q]: 0 where query qi may attend to key kj, -inf // otherwise. (1) pad mask: key kj valid iff kj < valid_len. (2) chunked- // limited window for streaming models. See header / NeMo _create_masks. + // PADDED query rows (qi >= valid_len) are left fully unmasked: with a + // limited window a deeply padded row can have its ENTIRE row masked, and + // softmax over all -inf is NaN — which the post-softmax query mask does NOT + // clear (NaN*0 == NaN) and which poisons every softmax row of the next + // layer through the key path. Their finite garbage is zeroed below. const int chunk_size = chunked_limited_ ? (att_right_ + 1) : 0; const int left_chunks = (chunked_limited_ && chunk_size > 0) ? (att_left_ / chunk_size) : 0; @@ -119,6 +124,10 @@ ggml_tensor* RelPosAttention::build_graph(ggml_context* ctx, ggml_tensor* xt, float* md = mask_host.data(); const float ninf = -INFINITY; for (int qi = 0; qi < T; ++qi) { + if (qi >= valid_len) { // dead (padded) query row: see above + for (int kj = 0; kj < T; ++kj) md[(size_t)qi * T + kj] = 0.0f; + continue; + } const int cq = chunked_limited_ ? (qi / chunk_size) : 0; for (int kj = 0; kj < T; ++kj) { bool ok = (kj < valid_len); @@ -256,6 +265,13 @@ ggml_tensor* RelPosAttention::build_graph_batched( // i12 = i02 % ne12 (head, here ne12=1 -> always 0) and i13 = i03 % ne13 // (batch, here ne13=B -> exact per-item), and ggml_soft_max_impl asserts // a->ne[2] % mask->ne[2] == 0 and a->ne[3] % mask->ne[3] == 0. + // PADDED query rows (qi >= valid_len[b]) are left fully unmasked — with the + // chunked-limited window a deeply padded row (pad tail > att_left) has its + // ENTIRE row masked, softmax over all -inf is NaN, the post-softmax query + // mask keeps NaN (NaN*0), and the NaN key columns then poison EVERY softmax + // row of the next layer. This corrupted whole items of a heterogeneous- + // length batch on streaming models (clips differing by more than ~att_left + // encoder frames decoded empty). Their finite garbage is zeroed below. const int chunk_size = chunked_limited_ ? (att_right_ + 1) : 0; const int left_chunks = (chunked_limited_ && chunk_size > 0) ? (att_left_ / chunk_size) : 0; @@ -266,6 +282,11 @@ ggml_tensor* RelPosAttention::build_graph_batched( for (int b = 0; b < B; ++b) { const int vl = valid_len[b]; for (int qi = 0; qi < T; ++qi) { + if (qi >= vl) { // dead (padded) query row: see above + for (int kj = 0; kj < T; ++kj) + md[(size_t)b * T * T + (size_t)qi * T + kj] = 0.0f; + continue; + } const int cq = chunked_limited_ ? (qi / chunk_size) : 0; for (int kj = 0; kj < T; ++kj) { bool ok = (kj < vl); @@ -365,15 +386,20 @@ ggml_tensor* RelPosAttention::build_graph_batched_local( ggml_tensor* bd = ggml_mul_mat(ctx, php, qv); // [P, T, H, B] (php broadcasts over B) ggml_tensor* scores = ggml_add(ctx, ac, bd); // [P, T, H, B] - // Per-item band mask [P, T, 1, B]. + // Per-item band mask [P, T, 1, B]. Padded query rows whose entire band lies + // beyond vl are left unmasked (all -inf -> softmax NaN -> NaN*0 stays NaN + // and poisons later layers); the query mask below zeroes their output. std::vector& mh = pool.alloc_f32((size_t)B * T * P); for (int b = 0; b < B; ++b) { const int vl = valid_len[b]; - for (int t = 0; t < T; ++t) + for (int t = 0; t < T; ++t) { + const bool dead_row = (t - att_left >= vl); for (int c = 0; c < P; ++c) { const int key = t - att_left + c; - mh[(size_t)b * T * P + (size_t)t * P + c] = (key >= 0 && key < vl) ? 0.0f : -INFINITY; + mh[(size_t)b * T * P + (size_t)t * P + c] = + (dead_row || (key >= 0 && key < vl)) ? 0.0f : -INFINITY; } + } } int64_t mne[4] = {P, T, 1, B}; ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 4, mne, @@ -492,13 +518,22 @@ ggml_tensor* RelPosAttention::build_graph_local(ggml_context* ctx, ggml_tensor* ggml_tensor* scores = ggml_add(ctx, ac, bd); // [P, T, H] // Band mask [P, T]: 0 if key in [0, valid_len), else -inf (covers the - // out-of-sequence window corners and pad frames). + // out-of-sequence window corners and pad frames). A padded query row + // whose ENTIRE band lies beyond valid_len (pad tail > att_left, i.e. a + // short item in a heavily padded batch) must NOT be all -inf: softmax + // would yield NaN, the post-softmax query mask keeps NaN (NaN*0), and + // the NaN keys poison later layers' valid frames. Leave such rows fully + // unmasked instead — their finite garbage output is zeroed by the query + // mask below. std::vector& mh = pool.alloc_f32((size_t)P * T); - for (int t = 0; t < T; ++t) + for (int t = 0; t < T; ++t) { + const bool dead_row = (t - att_left >= valid_len); for (int c = 0; c < P; ++c) { const int key = t - att_left + c; - mh[(size_t)t * P + c] = (key >= 0 && key < valid_len) ? 0.0f : -INFINITY; + mh[(size_t)t * P + c] = + (dead_row || (key >= 0 && key < valid_len)) ? 0.0f : -INFINITY; } + } int64_t mne[2] = {P, T}; ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, mne, mh.data(), mh.size() * sizeof(float)); @@ -628,13 +663,20 @@ ggml_tensor* RelPosAttention::build_graph_local_chunked( ggml_tensor* bd = ggml_mul_mat(ctx, php, qv); // [P, T, H] ggml_tensor* scores = ggml_add(ctx, ac, bd); // [P, T, H] - // Band mask [P, T]: 0 if key in [0, valid_len), else -inf. + // Band mask [P, T]: 0 if key in [0, valid_len), else -inf. As in + // build_graph_local: a padded query row whose entire band lies beyond + // valid_len must not be all -inf (softmax -> NaN, NaN*0 stays NaN and + // poisons later layers through the key path); leave such rows fully + // unmasked — the query mask below zeroes their output. std::vector& mh = pool.alloc_f32((size_t)P * T); - for (int t = 0; t < T; ++t) + for (int t = 0; t < T; ++t) { + const bool dead_row = (t - att_left >= valid_len); for (int c = 0; c < P; ++c) { const int key = t - att_left + c; - mh[(size_t)t * P + c] = (key >= 0 && key < valid_len) ? 0.0f : -INFINITY; + mh[(size_t)t * P + c] = + (dead_row || (key >= 0 && key < valid_len)) ? 0.0f : -INFINITY; } + } int64_t mne[2] = {P, T}; ggml_tensor* mask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, mne, mh.data(), mh.size() * sizeof(float));