From 3a20e537b0794c3b2f3adfdd8d044213424a30fd Mon Sep 17 00:00:00 2001 From: Akio Nishimura Date: Sun, 20 Sep 2026 09:19:09 +0900 Subject: [PATCH] qwen4exp: fix per-block bias indexing when a unified cache holds several sequences The per-block (blk_bias) path of set_input_qsa indexes the per-sequence bid arrays by block number (bid_cell[b], bid_idx[b]), but those arrays are indexed by bid, and the block table keys a block on (sequence set, bucket): with more than one sequence in a unified cache the two no longer coincide, so a block gets the visibility of a different sequence's group - one of our blocks can be masked out, one of another sequence's can pass through. Build an inverse block -> bid map once per ubatch, keeping the entry that owns the stream's sequence, and use it for both the seq_has test and the tail comparison. Single-sequence caches are unaffected (one_seq keeps the old correspondence). Reproduced with two live sequences in one stream (llama-server serves each slot as its own sequence from position 0): the model stops seeing its last messages and reports the input as empty. With the change, the same concurrent shape is answered correctly in 4/4 attempts (2/4 before). --- src/llama-memory-hybrid-idx.cpp | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 3972ce9ce293..f891c3d8ddec 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -511,6 +511,24 @@ void llama_memory_hybrid_idx::set_input_qsa( cur_cell_blk[j] = blk_of[j] < 0 ? dead_bid : blk_of[j]; } + // [TAG_QSA_BIASFIX] the block bias below is indexed by block, but the bid arrays are + // indexed by bid, and a unified cache with a second sequence interleaves one entry per + // sequence at the same position block. Build the inverse map once per ubatch, keeping the + // entry that owns this stream's cells: an entry at index b then really describes block b. + std::vector blk_entry(n_blocks, -1); + + for (int32_t t = 0; t < n_bid; ++t) { + const int64_t pb = bid_idx[t]/r; + + if (pb >= n_blocks || blk_entry[pb] >= 0) { + continue; + } + + if (one_seq || cells.seq_has((uint32_t) bid_cell[t], seq_of_stream)) { + blk_entry[pb] = t; + } + } + for (int64_t ii = 0; ii < n_tps; ++ii) { const int64_t i = s*n_tps + ii; const llama_seq_id seq_id = ubatch->seq_id[i][0]; @@ -549,13 +567,17 @@ void llama_memory_hybrid_idx::set_input_qsa( float * cur_blk_bias = dst_bias + i*n_blocks; for (int64_t b = 0; b < n_blocks; ++b) { - if (b >= n_bid || !cells.seq_has((uint32_t) bid_cell[b], seq_id)) { + // [TAG_QSA_BIASFIX] the entry that describes block b (not entry b) + const int32_t t = blk_entry[b]; + + if (t < 0) { + // no complete group of ours at this block: invisible, as before cur_blk_bias[b] = -INFINITY; continue; } // finite, so it can never meet a -inf and produce a nan - cur_blk_bias[b] = bid_idx[b] >= tail_start ? 1e9f : 0.0f; + cur_blk_bias[b] = bid_idx[t] >= tail_start ? 1e9f : 0.0f; } // the spare block holds the unpooled cells, which are the incomplete tail, so