From c07204df68e6bc264cfc57223d4d49bbd6b630bd Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 05:33:37 -0300 Subject: [PATCH 1/9] qwen4exp: fix sparse-attention block selection Build QSA blocks per sequence in token order and select complete blocks before expanding them to cache cells. Keep only the incomplete tail unconditionally visible and rotate pooled keys with the first token's full M-RoPE position. This prevents unified-cache sequences from sharing pooled indexer keys and avoids replacing padded tail entries with extra history tokens. Synthetic Qwen4 architecture, exact mask, F16 and Q8_0 state, sequence-copy, Metal, and AddressSanitizer checks pass. Assisted-by: Codex --- src/llama-memory-hybrid-idx.cpp | 440 +++++++++++++++++++++++--------- src/llama-memory-hybrid-idx.h | 37 ++- src/models/models.h | 14 +- src/models/qwen4exp.cpp | 342 ++++++++++--------------- tests/test-llama-archs.cpp | 68 ++++- 5 files changed, 551 insertions(+), 350 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d4e59d77e570..68450fc3e1d8 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -9,7 +9,9 @@ #include #include #include +#include #include +#include // // llama_memory_hybrid_idx @@ -137,6 +139,8 @@ void llama_memory_hybrid_idx::clear(bool data) { if (mem_idx) { mem_idx->clear(data); } + + qsa_histories.clear(); } bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { @@ -149,15 +153,96 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po mem_idx->seq_rm(seq_id, p0, p1); } - return get_mem_attn()->seq_rm(seq_id, p0, p1); + const bool res = get_mem_attn()->seq_rm(seq_id, p0, p1); + if (!res) { + return false; + } + + auto remove = [&](qsa_history & history) { + history.erase(std::remove_if(history.begin(), history.end(), [&](const qsa_token & token) { + return (p0 < 0 || token.pos[0] >= p0) && (p1 < 0 || token.pos[0] < p1); + }), history.end()); + }; + + if (seq_id < 0) { + for (auto & item : qsa_histories) { + remove(item.second); + } + } else { + auto it = qsa_histories.find(seq_id); + if (it != qsa_histories.end()) { + remove(it->second); + if (it->second.empty()) { + qsa_histories.erase(it); + } + } + } + + return true; } void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + if (seq_id_src == seq_id_dst) { + return; + } + + qsa_history copied; + const auto & cells_src = get_mem_attn()->get_cells(seq_id_src); + const auto & cells_dst = get_mem_attn()->get_cells(seq_id_dst); + const bool replace = &cells_src != &cells_dst; + const auto src = qsa_histories.find(seq_id_src); + if (src != qsa_histories.end()) { + using pos_key = std::tuple; + std::map> cells_by_pos; + for (uint32_t cell = 0; cell < cells_src.size(); ++cell) { + if (cells_src.is_empty(cell) || !cells_src.seq_has(cell, seq_id_src)) { + continue; + } + + const llama_pos pos = cells_src.pos_get(cell); + if ((p0 >= 0 && pos < p0) || (p1 >= 0 && pos >= p1)) { + continue; + } + + const auto & ext = cells_src.ext_get(cell); + cells_by_pos[{ pos, ext.y, ext.x }].push_back(!replace && cells_src.seq_has(cell, seq_id_dst)); + } + + std::map next_cell; + for (const auto & token : src->second) { + if ((p0 >= 0 && token.pos[0] < p0) || (p1 >= 0 && token.pos[0] >= p1)) { + continue; + } + + const pos_key key = { token.pos[0], token.pos[1], token.pos[2] }; + auto cells = cells_by_pos.find(key); + if (cells == cells_by_pos.end()) { + continue; + } + + size_t & index = next_cell[key]; + if (index < cells->second.size() && !cells->second[index++]) { + copied.push_back(token); + } + } + } + llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1); if (mem_idx) { mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); } + + if (replace) { + if (copied.empty()) { + qsa_histories.erase(seq_id_dst); + } else { + qsa_histories[seq_id_dst] = std::move(copied); + } + } else if (!copied.empty()) { + auto & dst = qsa_histories[seq_id_dst]; + dst.insert(dst.end(), copied.begin(), copied.end()); + } } void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { @@ -166,6 +251,13 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { if (mem_idx) { mem_idx->seq_keep(seq_id); } + + auto it = qsa_histories.find(seq_id); + qsa_history keep = it == qsa_histories.end() ? qsa_history{} : std::move(it->second); + qsa_histories.clear(); + if (!keep.empty()) { + qsa_histories.emplace(seq_id, std::move(keep)); + } } void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { @@ -174,6 +266,15 @@ void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_p if (mem_idx) { mem_idx->seq_add(seq_id, p0, p1, shift); } + + auto it = qsa_histories.find(seq_id); + if (it != qsa_histories.end()) { + for (auto & token : it->second) { + if ((p0 < 0 || token.pos[0] >= p0) && (p1 < 0 || token.pos[0] < p1)) { + token.pos[0] += shift; + } + } + } } void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { @@ -182,6 +283,15 @@ void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_p if (mem_idx) { mem_idx->seq_div(seq_id, p0, p1, d); } + + auto it = qsa_histories.find(seq_id); + if (it != qsa_histories.end()) { + for (auto & token : it->second) { + if ((p0 < 0 || token.pos[0] >= p0) && (p1 < 0 || token.pos[0] < p1)) { + token.pos[0] /= d; + } + } + } } std::map llama_memory_hybrid_idx::memory_breakdown() const { @@ -205,8 +315,28 @@ void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id se if (mem_idx) { mem_idx->state_write(io, seq_id, flags); } - } + uint32_t n_histories = 0; + if (seq_id < 0) { + n_histories = (uint32_t) qsa_histories.size(); + } else if (qsa_histories.count(seq_id) != 0) { + n_histories = 1; + } + io.write(&n_histories, sizeof(n_histories)); + + for (const auto & item : qsa_histories) { + if (seq_id >= 0 && item.first != seq_id) { + continue; + } + + io.write(&item.first, sizeof(item.first)); + const uint64_t n_tokens = item.second.size(); + io.write(&n_tokens, sizeof(n_tokens)); + for (const auto & token : item.second) { + io.write(token.pos.data(), sizeof(token.pos)); + } + } + } } void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { @@ -230,6 +360,34 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_ if (mem_idx) { mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn); } + + uint32_t n_histories; + io.read(&n_histories, sizeof(n_histories)); + if (n_histories > LLAMA_MAX_SEQ) { + throw std::runtime_error("invalid QSA history count"); + } + + if (seq_id < 0) { + qsa_histories.clear(); + } else { + qsa_histories.erase(seq_id); + } + + for (uint32_t ih = 0; ih < n_histories; ++ih) { + llama_seq_id stored_seq; + uint64_t n_tokens; + io.read(&stored_seq, sizeof(stored_seq)); + io.read(&n_tokens, sizeof(n_tokens)); + if (stored_seq < 0 || stored_seq >= LLAMA_MAX_SEQ || n_tokens > get_mem_attn()->get_size()) { + throw std::runtime_error("invalid QSA history"); + } + + auto & history = qsa_histories[seq_id < 0 ? stored_seq : seq_id]; + history.resize(n_tokens); + for (auto & token : history) { + io.read(token.pos.data(), sizeof(token.pos)); + } + } } } catch (...) { @@ -255,12 +413,156 @@ void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) { if (mem_idx) { mem_idx->seq_rm(seq_id, -1, -1); } + + qsa_histories.erase(seq_id); } llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const { return mem_idx.get(); } +void llama_memory_hybrid_idx::commit_qsa_tokens(const llama_ubatch & ubatch) { + if (!mem_idx) { + return; + } + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + qsa_token token = {}; + if (ubatch.token) { + token.pos = { ubatch.pos[i], ubatch.pos[i], ubatch.pos[i], 0 }; + } else { + for (uint32_t ip = 0; ip < token.pos.size(); ++ip) { + token.pos[ip] = ip < ubatch.n_pos ? ubatch.pos[i + ip*ubatch.n_tokens] : ubatch.pos[i]; + } + } + + for (int32_t is = 0; is < ubatch.n_seq_id[i]; ++is) { + qsa_histories[ubatch.seq_id[i][is]].push_back(token); + } + } +} + +void llama_memory_hybrid_idx::set_input_qsa( + ggml_tensor * block_cells, + ggml_tensor * block_pos, + ggml_tensor * block_mask, + ggml_tensor * selected, + const ggml_tensor * kq_mask, + const llama_ubatch * ubatch, + uint32_t ratio, + uint32_t block_topk) const { + GGML_ASSERT(ggml_backend_buffer_is_host(block_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(block_pos->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(block_mask->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(selected->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(kq_mask->buffer)); + + const int64_t n_blocks = block_cells->ne[1]; + const int64_t n_tokens = ubatch->n_tokens; + const int64_t n_pos = block_pos->ne[1]; + const int64_t n_kv = selected->ne[0]; + + GGML_ASSERT(block_cells->type == GGML_TYPE_I32); + GGML_ASSERT(block_pos->type == GGML_TYPE_I32); + GGML_ASSERT(block_mask->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_F32); + GGML_ASSERT(block_cells->ne[0] == ratio && block_cells->ne[2] == n_tokens); + GGML_ASSERT(block_pos->ne[0] == n_blocks && block_pos->ne[2] == n_tokens); + GGML_ASSERT(block_mask->ne[0] == n_blocks && block_mask->ne[1] == n_tokens); + GGML_ASSERT(selected->ne[1] == n_tokens); + GGML_ASSERT(kq_mask->ne[0] == n_kv); + + int32_t * cell_data = (int32_t *) block_cells->data; + int32_t * pos_data = (int32_t *) block_pos->data; + float * mask_data = (float *) block_mask->data; + float * selected_data = (float *) selected->data; + + std::fill(cell_data, cell_data + ggml_nelements(block_cells), 0); + std::fill(pos_data, pos_data + ggml_nelements(block_pos), 0); + std::fill(mask_data, mask_data + ggml_nelements(block_mask), -INFINITY); + std::fill(selected_data, selected_data + ggml_nelements(selected), 0.0f); + + auto mask_visible = [&](int64_t query, uint32_t cell) { + const int64_t index = query*n_kv + cell; + if (kq_mask->type == GGML_TYPE_F16) { + return std::isfinite(ggml_fp16_to_fp32(((const ggml_fp16_t *) kq_mask->data)[index])); + } + return std::isfinite(((const float *) kq_mask->data)[index]); + }; + + for (int64_t iq = 0; iq < n_tokens; ++iq) { + const llama_seq_id seq_id = ubatch->seq_id[iq][0]; + const auto found = qsa_histories.find(seq_id); + if (found == qsa_histories.end()) { + continue; + } + + const auto & cells = get_mem_attn()->get_cells(seq_id); + using pos_key = std::tuple; + std::map> cells_by_pos; + for (uint32_t cell = 0; cell < cells.size() && cell < (uint32_t) n_kv; ++cell) { + if (cells.is_empty(cell) || !cells.seq_has(cell, seq_id)) { + continue; + } + + const auto & ext = cells.ext_get(cell); + cells_by_pos[{ cells.pos_get(cell), ext.y, ext.x }].push_back(cell); + } + + std::map next_cell; + std::vector> visible; + for (const auto & token : found->second) { + const pos_key key = { token.pos[0], token.pos[1], token.pos[2] }; + auto cells = cells_by_pos.find(key); + if (cells == cells_by_pos.end()) { + continue; + } + + size_t & index = next_cell[key]; + if (index >= cells->second.size()) { + continue; + } + + const uint32_t cell = cells->second[index++]; + if (mask_visible(iq, cell)) { + visible.emplace_back(&token, cell); + } + } + + const size_t n_complete = visible.size()/ratio; + const size_t n_write = std::min(n_complete, n_blocks); + std::vector used_cells(n_kv, 0); + for (size_t ib = 0; ib < n_write; ++ib) { + mask_data[iq*n_blocks + ib] = 0.0f; + for (uint32_t ir = 0; ir < ratio; ++ir) { + const uint32_t cell = visible[ib*ratio + ir].second; + cell_data[(iq*n_blocks + ib)*ratio + ir] = cell; + used_cells[cell] = 1; + } + for (int64_t ip = 0; ip < n_pos; ++ip) { + pos_data[(iq*n_pos + ip)*n_blocks + ib] = visible[ib*ratio].first->pos[ip]; + } + } + + uint32_t fallback_cell = 0; + for (size_t ib = n_write; ib < (size_t) n_blocks; ++ib) { + for (uint32_t ir = 0; ir < ratio; ++ir) { + while (fallback_cell < used_cells.size() && used_cells[fallback_cell]) { + ++fallback_cell; + } + GGML_ASSERT(fallback_cell < used_cells.size()); + cell_data[(iq*n_blocks + ib)*ratio + ir] = fallback_cell; + used_cells[fallback_cell++] = 1; + } + } + + const size_t selected_start = n_complete <= block_topk ? 0 : n_complete*ratio; + for (size_t iv = selected_start; iv < visible.size(); ++iv) { + selected_data[iq*n_kv + visible[iv].second] = 1.0f; + } + } +} + // // llama_memory_hybrid_idx_context // @@ -307,7 +609,8 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( mem(mem), ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)), ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : - new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {} + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)), + has_ubatches(true) {} bool llama_memory_hybrid_idx_context::next() { if (ctx_idx) { @@ -326,6 +629,10 @@ bool llama_memory_hybrid_idx_context::apply() { res = res & ctx_idx->apply(); } + if (res && ctx_idx && has_ubatches) { + mem->commit_qsa_tokens(ctx_idx->get_ubatch()); + } + return res; } @@ -340,126 +647,17 @@ uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { } void llama_memory_hybrid_idx_context::set_input_qsa( - ggml_tensor * cell_blk, - ggml_tensor * blk_cells, - ggml_tensor * blk_pos, - ggml_tensor * bias, + ggml_tensor * block_cells, + ggml_tensor * block_pos, + ggml_tensor * block_mask, + ggml_tensor * selected, + const ggml_tensor * kq_mask, const llama_ubatch * ubatch, - uint32_t ratio, - bool blk_bias) const { + uint32_t ratio, + uint32_t block_topk) const { GGML_ASSERT(ratio > 0); GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); - GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); - - const int64_t n_kv = cell_blk->ne[0]; - const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch - const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns); - const int64_t n_tokens = ubatch->n_tokens; - const int64_t r = ratio; - - GGML_ASSERT(n_tokens % n_ns == 0); - const int64_t n_tps = n_tokens/n_ns; // tokens per stream - - int32_t * dst_cell_blk = (int32_t *) cell_blk->data; - int32_t * dst_blk_cells = (int32_t *) blk_cells->data; - int32_t * dst_blk_pos = (int32_t *) blk_pos->data; - float * dst_bias = (float *) bias->data; - - // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio - // all mrope sections carry it: exact for text, approximate for images - for (int64_t sec = 0; sec < 4; ++sec) { - for (int64_t s = 0; s < n_ns; ++s) { - for (int64_t b = 0; b < n_blocks; ++b) { - dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r); - } - } - } - - // one pass per stream: cell j is a different token in each, so no mapping is shared - std::vector blk_of(n_kv); - std::vector filled(n_blocks); - - for (int64_t s = 0; s < n_ns; ++s) { - // ubatch index s*n_tps belongs to this stream; ask which cells array it uses - const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; - const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream); - - int32_t * cur_cell_blk = dst_cell_blk + s*n_kv; - int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks); - - // an incomplete block cannot be pooled; the bias below forces those tail cells in - // -1 means no usable block, and block 0 only keeps the gather in range - std::fill(blk_of.begin(), blk_of.end(), -1); - std::fill(filled.begin(), filled.end(), 0); - std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0); - - // a cell no block covers needs its own -inf, which a per-block bias cannot carry - // every cache path keeps the position below the cell window, so this stays false - bool oor = false; - - for (int64_t j = 0; j < n_kv; ++j) { - if (cells.is_empty(j)) { - continue; - } - - const llama_pos p = cells.pos_get(j); - const int64_t b = p/r; - - if (b >= n_blocks) { - oor = true; - continue; - } - - blk_of[j] = (int32_t) b; - cur_blk_cells[b*r + (p%r)] = (int32_t) j; - filled[b]++; - } - - GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); - - // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it - // per-cell mode carries that -inf itself and only needs the gather in range - for (int64_t j = 0; j < n_kv; ++j) { - if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) { - blk_of[j] = -1; - } - cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; - } - - 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]; - const llama_pos q = ubatch->pos[i]; - - // the tail is an incomplete block and is always visible, as in the reference - const llama_pos tail_start = (q + 1)/r*r; - - if (blk_bias) { - // a block sits wholly inside or outside the tail, so one value covers it - // the caller adds the attention mask, which drops empty, foreign and future cells - float * cur_blk_bias = dst_bias + i*n_blocks; - - for (int64_t b = 0; b < n_blocks; ++b) { - // finite, so it can never meet a -inf and produce a nan - cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f); - } - - continue; - } - - float * cur_bias = dst_bias + i*n_kv; - - for (int64_t j = 0; j < n_kv; ++j) { - float v = -INFINITY; - - if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { - // finite, so it can never meet a -inf and produce a nan - v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); - } - - cur_bias[j] = v; - } - } - } + mem->set_input_qsa(block_cells, block_pos, block_mask, selected, + kq_mask, ubatch, ratio, block_topk); } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e3472646d0f6..5f1ee7794303 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -2,6 +2,8 @@ #include "llama-memory-hybrid.h" +#include +#include #include #include @@ -75,7 +77,20 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer + void set_input_qsa(ggml_tensor * block_cells, ggml_tensor * block_pos, + ggml_tensor * block_mask, ggml_tensor * selected, + const ggml_tensor * kq_mask, const llama_ubatch * ubatch, + uint32_t ratio, uint32_t block_topk) const; + + void commit_qsa_tokens(const llama_ubatch & ubatch); + private: + struct qsa_token { + std::array pos; + }; + + using qsa_history = std::vector; + // forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step // seq_id < 0 drops the whole context, as the caches themselves do on a failed restore void state_drop(llama_seq_id seq_id); @@ -85,6 +100,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + + std::map qsa_histories; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -129,20 +146,14 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified uint32_t get_n_stream() const; - // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache. - // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout: - // cell_blk I32 [n_kv, ns] block each cell belongs to - // blk_cells I32 [ratio*n_blocks, ns] cells making up each block - // blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token - // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible - // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns] - // the caller then adds the attention mask, the only part of the bias that varies within a block - void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, - ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, - bool blk_bias) const; + // QSA blocks follow each sequence's token order, not physical cells or scalar positions. + void set_input_qsa(ggml_tensor * block_cells, ggml_tensor * block_pos, + ggml_tensor * block_mask, ggml_tensor * selected, + const ggml_tensor * kq_mask, const llama_ubatch * ubatch, + uint32_t ratio, uint32_t block_topk) const; private: - const llama_memory_hybrid_idx * mem = nullptr; + llama_memory_hybrid_idx * mem = nullptr; // streams per ubatch, read from the slot infos before ctx_idx takes them // declared first, so it is initialised while sinfos_idx is still intact @@ -151,6 +162,8 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // null unless the model has an indexer and this is a batch or full context const llama_memory_context_ptr ctx_idx; + const bool has_ubatches = false; + // mirrors the base class's ubatch cursor, which is private there size_t i_cur = 0; }; diff --git a/src/models/models.h b/src/models/models.h index af60764c2f7f..0283500a534c 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2310,22 +2310,12 @@ struct llama_model_qwen4exp : public llama_model_base { int * sections, int il); - // dense self-attention restricted to the cells that top_k names - ggml_tensor * build_attn_qsa( - llm_graph_input_attn_kv * inp, - ggml_tensor * q_cur, - ggml_tensor * k_cur, - ggml_tensor * v_cur, - ggml_tensor * top_k, - float kq_scale, - int il); - // the QSA cache layout inputs do not depend on the layer, only on its compress ratio, // so the layers sharing a ratio share one input set std::map qsa_inps; - // QSA: token indices this layer's queries may attend to, or nullptr for dense - ggml_tensor * build_qsa_top_k( + // QSA mask for this layer, or nullptr for dense attention + ggml_tensor * build_qsa_mask( const llama_memory_hybrid_idx_context * mctx_hyb, ggml_tensor * cur, ggml_tensor * inp_pos, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b50038..ac0abe4e2d52 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -99,6 +99,17 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { } } + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + const uint32_t ratio = hparams.dsv4_compress_ratios[il]; + if (hparams.is_recr(il)) { + if (ratio != 0) { + throw std::runtime_error(format("Qwen4-Exp recurrent layer %u has a QSA compression ratio", il)); + } + } else if (ratio == 0 || hparams.indexer_top_k % ratio != 0) { + throw std::runtime_error(format("invalid Qwen4-Exp QSA compression ratio %u at layer %u", ratio, il)); + } + } + switch (hparams.n_layer()) { case 48: type = LLM_TYPE_A3B; break; default: type = LLM_TYPE_UNKNOWN; @@ -416,13 +427,17 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( // one mean-pooled indexer key scores each block; set_input resolves the cache layout class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { public: - llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) : - mctx(mctx), ratio(ratio), blk_bias(blk_bias) {} + llm_graph_input_qsa( + const llama_memory_hybrid_idx_context * mctx, + ggml_tensor * kq_mask, + uint32_t ratio, + uint32_t block_topk) : + mctx(mctx), kq_mask(kq_mask), ratio(ratio), block_topk(block_topk) {} virtual ~llm_graph_input_qsa() = default; void set_input(const llama_ubatch * ubatch) override { - mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); + mctx->set_input_qsa(block_cells, block_pos, block_mask, selected, + kq_mask, ubatch, ratio, block_topk); } bool can_reuse(const llm_graph_params & params) override { @@ -434,39 +449,33 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { } const int64_t n_kv = idx->get_n_kv(); - const int64_t n_stream = mctx->get_n_stream(); - const int64_t n_blocks = (n_kv + ratio - 1)/ratio; + const int64_t n_blocks = n_kv/ratio; bool res = true; - res &= params.ubatch.n_tokens % n_stream == 0; - - res &= k_idxs->ne[0] == params.ubatch.n_tokens; - res &= cell_blk->ne[0] == n_kv; - res &= cell_blk->ne[1] == n_stream; - res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks; - res &= blk_pos->ne[0] == 4*n_blocks*n_stream; - res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv); - res &= bias->ne[1] == params.ubatch.n_tokens/n_stream; + res &= n_kv > (int64_t) block_topk*ratio + ratio - 1; + res &= block_cells->ne[1] == n_blocks; + res &= block_cells->ne[2] == params.ubatch.n_tokens; + res &= block_pos->ne[2] == params.ubatch.n_tokens; + res &= block_mask->ne[1] == params.ubatch.n_tokens; + res &= selected->ne[0] == n_kv; + res &= selected->ne[1] == params.ubatch.n_tokens; return res; } - // per stream: a cell index names a different token in each stream - ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] - ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream] - ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] - ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] - ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream] + ggml_tensor * block_cells = nullptr; + ggml_tensor * block_pos = nullptr; + ggml_tensor * block_mask = nullptr; + ggml_tensor * selected = nullptr; const llama_memory_hybrid_idx_context * mctx; + ggml_tensor * kq_mask; const uint32_t ratio; - - // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded - const bool blk_bias; + const uint32_t block_topk; }; -ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask( const llama_memory_hybrid_idx_context * mctx_hyb, ggml_tensor * cur, ggml_tensor * inp_pos, @@ -482,21 +491,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( GGML_ASSERT(r > 0); - const int64_t n_blocks = (n_kv + r - 1)/r; + const int64_t n_blocks = n_kv/r; + const int64_t block_topk = hparams.indexer_top_k/r; - // build_attn_qsa and the KQ mask need the tokens to divide evenly across the streams const int64_t n_stream = mctx_hyb->get_n_stream(); GGML_ASSERT(n_tokens % n_stream == 0); const int64_t n_tps = n_tokens/n_stream; - // only the "which block is visible" half of the bias varies per block - // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells - // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out - // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there - const bool blk_bias = kq_mask != nullptr && - kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && - cparams.causal_attn && !hparams.use_alibi; - // nothing above depends on the layer, so the layers sharing a ratio share one input set llm_graph_input_qsa * inp = nullptr; @@ -504,59 +505,29 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( if (it != qsa_inps.end()) { inp = it->second; } else { - auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + auto qsa = std::make_unique( + mctx_hyb, kq_mask, (uint32_t) r, (uint32_t) block_topk); - qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); - qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); - qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); - qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); - qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream); + qsa->block_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, r, n_blocks, n_tokens); + qsa->block_pos = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, n_blocks, 4, n_tokens); + qsa->block_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_blocks, n_tokens); + qsa->selected = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); - ggml_set_input(qsa->cell_blk); - ggml_set_input(qsa->blk_cells); - ggml_set_input(qsa->blk_pos); - ggml_set_input(qsa->bias); + ggml_set_input(qsa->block_cells); + ggml_set_input(qsa->block_pos); + ggml_set_input(qsa->block_mask); + ggml_set_input(qsa->selected); inp = qsa.get(); res->add_input(std::move(qsa)); qsa_inps.emplace((uint32_t) r, inp); } - // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither - ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); - k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens); - cb(k_raw, "indexer_k_raw", il); - - ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il)); + kq_mask = inp->kq_mask; // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream]. ggml_tensor * k_all = mctx_idx->get_k(ctx0, il); k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0); - - // gathers per stream: blk_cells row s indexes stream s's own cells - ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); - members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); - - // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows - ggml_tensor * pooled = nullptr; - for (int64_t i = 0; i < r; ++i) { - ggml_tensor * slice = ggml_cont(ctx0, - ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, - members->nb[2], members->nb[3], i*members->nb[1])); - pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; - } - pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); - cb(pooled, "indexer_k_pooled", il); - - // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after. - pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); - pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); - pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, - n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream); - cb(pooled, "indexer_k", il); - ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur); q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens); q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il); @@ -565,128 +536,86 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ext_factor, attn_factor, beta_fast, beta_slow); cb(q, "indexer_q", il); - // rectify each head dot product before the sum, as in the DeepSeek lightning indexer - // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s - ggml_tensor * score = ggml_mul_mat(ctx0, pooled, - ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream)); - score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream); - score = ggml_relu(ctx0, score); - score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); - score = ggml_sum_rows(ctx0, score); - score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); - cb(score, "indexer_score", il); - - // one value per block, so it is cheaper to bias here than after the cells are expanded - if (blk_bias) { - score = ggml_add(ctx0, score, inp->bias); - } - - // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary - ggml_tensor * expanded = ggml_get_rows(ctx0, - ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); - expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); - - if (blk_bias) { - // flash attention keeps the mask in f16; the scores are f32 - ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); - expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream)); - } else { - expanded = ggml_add(ctx0, expanded, inp->bias); - } - cb(expanded, "indexer_score_tokens", il); - - // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail - const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); - - ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); - - // build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask. - top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); - cb(top_k, "indexer_top_k", il); - - return top_k; -} - -// Dense GQA self-attention restricted to the cells that top_k names. -// The mask build below copies the MLA sparse path in llm_graph_context::build_attn. -ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( - llm_graph_input_attn_kv * inp, - ggml_tensor * q_cur, - ggml_tensor * k_cur, - ggml_tensor * v_cur, - ggml_tensor * top_k, - float kq_scale, - int il) { - // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer - // has already scored with its own query in build_qsa_top_k, so top_k is unaffected. - if (inp->self_k_rot) { - q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); - k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); + const ggml_type activation_type = mctx_idx->type_k(); + std::vector selected_streams; + selected_streams.reserve(n_stream); + + for (int64_t is = 0; is < n_stream; ++is) { + ggml_tensor * cache = ggml_view_2d(ctx0, k_all, idx_dim, n_kv, + k_all->nb[1], is*k_all->nb[2]); + ggml_tensor * block_cells = ggml_view_3d(ctx0, inp->block_cells, r, n_blocks, n_tps, + inp->block_cells->nb[1], inp->block_cells->nb[2], is*n_tps*inp->block_cells->nb[2]); + ggml_tensor * block_keys = ggml_get_rows(ctx0, cache, + ggml_reshape_1d(ctx0, block_cells, r*n_blocks*n_tps)); + block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, r, n_blocks, n_tps); + block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys)); + block_keys = ggml_mean(ctx0, block_keys); + block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys)); + block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, 1, n_blocks*n_tps); + + if (block_keys->type != activation_type) { + block_keys = ggml_cast(ctx0, block_keys, activation_type); + block_keys = ggml_cast(ctx0, block_keys, GGML_TYPE_F32); + } + block_keys = build_norm(block_keys, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * block_pos = ggml_view_3d(ctx0, inp->block_pos, n_blocks, 4, n_tps, + inp->block_pos->nb[1], inp->block_pos->nb[2], is*n_tps*inp->block_pos->nb[2]); + block_keys = ggml_rope_multi(ctx0, block_keys, + ggml_reshape_1d(ctx0, block_pos, n_blocks*4*n_tps), nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(block_keys, "indexer_k", il); + + block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, n_blocks, 1, n_tps); + ggml_tensor * query = ggml_view_3d(ctx0, q, idx_dim, n_idx_h, n_tps, + q->nb[1], q->nb[2], is*n_tps*q->nb[2]); + query = ggml_reshape_4d(ctx0, query, idx_dim, n_idx_h, 1, n_tps); + + ggml_tensor * scores = ggml_mul_mat(ctx0, block_keys, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_relu(ctx0, scores); + scores = ggml_sum_rows(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3))); + scores = ggml_scale(ctx0, ggml_reshape_2d(ctx0, scores, n_blocks, n_tps), + 1.0f/sqrtf((float) idx_dim)); + + ggml_tensor * block_mask = ggml_view_2d(ctx0, inp->block_mask, n_blocks, n_tps, + inp->block_mask->nb[1], is*n_tps*inp->block_mask->nb[1]); + scores = ggml_add(ctx0, scores, block_mask); + cb(scores, "indexer_score", il); + + ggml_tensor * top_blocks = ggml_top_k(ctx0, scores, block_topk); + ggml_tensor * top_cells = ggml_get_rows(ctx0, block_cells, top_blocks); + top_cells = ggml_reshape_2d(ctx0, top_cells, r*block_topk, n_tps); + + ggml_tensor * base_selected = ggml_view_2d(ctx0, inp->selected, n_kv, n_tps, + inp->selected->nb[1], is*n_tps*inp->selected->nb[1]); + base_selected = ggml_reshape_3d(ctx0, base_selected, 1, n_kv, n_tps); + ggml_tensor * selected_top = ggml_fill(ctx0, base_selected, 0.0f); + ggml_tensor * ones = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, r*block_topk, n_tps); + ones = ggml_fill(ctx0, ones, 1.0f); + selected_top = ggml_set_rows(ctx0, selected_top, ones, top_cells); + ggml_tensor * selected_stream = ggml_clamp( + ctx0, ggml_add(ctx0, base_selected, selected_top), 0.0f, 1.0f); + selected_streams.push_back(ggml_reshape_2d(ctx0, selected_stream, n_kv, n_tps)); } - if (inp->self_v_rot) { - v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); + ggml_tensor * selected = selected_streams[0]; + for (int64_t is = 1; is < n_stream; ++is) { + selected = ggml_concat(ctx0, selected, selected_streams[is], 1); } + selected = ggml_scale_bias(ctx0, selected, 1e30f, -1e30f); - // these nodes are added to the graph together so that they are not reordered - // by doing so, the number of splits in the graph is reduced - // expand k later to enable rope fusion which directly writes into k-v cache - ggml_build_forward_expand(gf, q_cur); - ggml_build_forward_expand(gf, v_cur); - ggml_build_forward_expand(gf, k_cur); - - const auto * mctx_cur = inp->mctx; - - // store to KV cache - { - const auto & k_idxs = inp->get_k_idxs(); - const auto & v_idxs = inp->get_v_idxs(); - - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + ggml_tensor * base_mask = ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens); + if (base_mask->type != GGML_TYPE_F32) { + base_mask = ggml_cast(ctx0, base_mask, GGML_TYPE_F32); } - - ggml_tensor * kq_mask = inp->get_kq_mask(); - - // prepare new kq mask - starts filled with -INFINITY - ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); - - // reshape KQ mask into tensor with rows of size 1: - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] - kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); - - // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] - ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - - // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] - // this will be our source of zero values for unmasking top k mask elements - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); - zeros = ggml_fill(ctx0, zeros, 0.0f); - - // modify KQ mask by unmasking elements that are in top_k indices - // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) - ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); - - // reshape to restore the original shape of KQ mask: - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] - kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); - - // combine with the original kq mask - kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); - - ggml_tensor * q = q_cur; - ggml_tensor * k = mctx_cur->get_k(ctx0, il); - ggml_tensor * v = mctx_cur->get_v(ctx0, il); - - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il); - cb(cur, "kqv_out", il); - - // the rotation is its own inverse, so undo it on the value side of the output - if (inp->self_v_rot) { - cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); + ggml_tensor * mask = ggml_add(ctx0, base_mask, selected); + if (cparams.flash_attn) { + mask = ggml_cast(ctx0, mask, GGML_TYPE_F16); } - - return cur; + cb(mask, "qsa_mask", il); + return mask; } ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( @@ -699,10 +628,23 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); - // indexer reads the same block input as q/k/v; no cache or no ratio means dense - const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; + const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); + if (mctx_idx) { + ggml_tensor * index_k = build_lora_mm(model.layers[il].index_k_proj, cur); + index_k = ggml_reshape_3d(ctx0, index_k, hparams.indexer_head_size, 1, n_tokens); + cb(index_k, "indexer_k_raw", il); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, index_k, inp->get_k_idxs(), il)); + } - ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + const int64_t ratio = hparams.dsv4_compress_ratios[il]; + const bool qsa = mctx_idx && ratio > 0 && + mctx_idx->get_n_kv() > (int64_t) hparams.indexer_top_k + ratio - 1; + if (qsa) { + inp->self_kq_mask_cnv = build_qsa_mask( + mctx_hyb, cur, inp_pos, inp->self_kq_mask, sections, il); + } else { + inp->self_kq_mask_cnv = inp->self_kq_mask; + } // Qwen3Next uses a single Q projection that outputs query + gate ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] @@ -754,13 +696,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; - if (top_k) { - cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il); - } else { - cur = build_attn(inp, - nullptr, nullptr, nullptr, - Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); - } + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); cb(cur, "attn_pregate", il); ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 35a3286e4a1b..ab05df77e1bd 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -252,8 +252,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_QWEN4EXP) { ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8)); - // without this the QSA layers fall back to dense and go uncovered - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); + std::vector ratios(n_layer, 0); + for (uint32_t il = 1; il < n_layer; il += 2) { + ratios[il] = 4; + } + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, ratios); } // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused @@ -347,7 +350,8 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, + ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -360,6 +364,8 @@ static std::pair get_model_and_ctx( ctx_params.n_ctx = 0; ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; + ctx_params.cb_eval = cb_eval; + ctx_params.cb_eval_user_data = cb_eval_user_data; if (!encode) { ctx_params.n_ubatch = 64; } @@ -412,6 +418,54 @@ static std::vector get_logits( return ret; } +struct qwen4_qsa_mask_check { + int64_t n_seen = 0; + int64_t n_tokens_seen = 0; + bool ok = true; +}; + +static bool check_qwen4_qsa_mask(ggml_tensor * tensor, bool ask, void * user_data) { + if (strncmp(tensor->name, "qsa_mask", 8) != 0) { + return false; + } + if (ask) { + return true; + } + + auto * check = (qwen4_qsa_mask_check *) user_data; + const int64_t n_kv = tensor->ne[0]; + const int64_t n_tokens = tensor->ne[1]; + std::vector mask(ggml_nelements(tensor)); + if (tensor->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(tensor, mask.data(), 0, ggml_nbytes(tensor)); + } else { + GGML_ASSERT(tensor->type == GGML_TYPE_F16); + std::vector mask_f16(ggml_nelements(tensor)); + ggml_backend_tensor_get(tensor, mask_f16.data(), 0, ggml_nbytes(tensor)); + for (size_t i = 0; i < mask.size(); ++i) { + mask[i] = ggml_fp16_to_fp32(mask_f16[i]); + } + } + + for (int64_t it = 0; it < n_tokens; ++it) { + const int64_t n_visible = check->n_tokens_seen + it + 1; + const int64_t n_complete = n_visible/4; + const int64_t expected = n_complete <= 2 ? n_visible : 8 + n_visible%4; + int64_t actual = 0; + for (int64_t ikv = 0; ikv < n_kv; ++ikv) { + actual += mask[it*n_kv + ikv] > -1e20f; + } + if (actual != expected) { + fprintf(stderr, "Qwen4 QSA mask row %lld selects %lld tokens, expected %lld\n", + (long long) (check->n_tokens_seen + it), (long long) actual, (long long) expected); + } + check->ok = check->ok && actual == expected; + } + check->n_tokens_seen += n_tokens; + check->n_seen++; + return true; +} + static bool moe_mandatory(const llm_arch arch) { switch (arch) { case LLM_ARCH_LLAMA4: @@ -686,6 +740,14 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg if (arch == LLM_ARCH_BAILINGMOE3) { GGML_ASSERT(gguf_remove_key(gguf_ctx.get(), "bailingmoe3.kda.safe_gate") >= 0); } + if (arch == LLM_ARCH_QWEN4EXP) { + qwen4_qsa_mask_check check; + auto model_and_ctx = get_model_and_ctx( + gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, + check_qwen4_qsa_mask, &check); + get_logits(model_and_ctx.first.get(), model_and_ctx.second.get(), tokens); + GGML_ASSERT(check.ok && check.n_seen > 0); + } std::pair model_and_ctx_cpu; std::vector logits_cpu; for (device_config & dc : dev_configs) { From fce561dff778812786005701b600f51f6456e4fe Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 05:40:45 -0300 Subject: [PATCH 2/9] qwen4exp: support independent PLE embedding widths Size the PLE key and value projections from the concatenated n-gram embedding instead of assuming it matches the model hidden width. Validate the head count before narrowing it to the stored type. Add a synthetic PLE model with a 64-wide embedding and a 256-wide hidden state, then verify inference and model roundtrip. Assisted-by: Codex --- src/models/qwen4exp.cpp | 26 ++++++++++++++++++-------- tests/test-llama-archs.cpp | 32 +++++++++++++++++++++++++++++++- 2 files changed, 49 insertions(+), 9 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index ac0abe4e2d52..a18cd867a5c3 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -61,14 +61,16 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0); - hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; - hparams.ple_head_dim = hparams.n_embd_per_layer; if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) { throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size)); } - if (hparams.ple_n_heads == 0 || hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS) { - throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads)); + + const uint64_t ple_n_heads = (uint64_t) (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; + hparams.ple_head_dim = hparams.n_embd_per_layer; + if (ple_n_heads == 0 || ple_n_heads > LLAMA_MAX_PLE_HEADS) { + throw std::runtime_error(format("PLE head count %" PRIu64 " is out of range", ple_n_heads)); } + hparams.ple_n_heads = (uint32_t) ple_n_heads; ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); @@ -138,8 +140,15 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back if (hparams.ple_n_heads > 0) { const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); - const auto & ple_w = ml.require_weight(ple_name.c_str()); - const int64_t ple_rows = ple_w.tensor->ne[1]; + const auto * ple_w = ml.get_weight(ple_name.c_str()); + int64_t ple_rows = 0; + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + ple_rows = std::max(ple_rows, + (int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h]); + } + if (ple_w) { + ple_rows = ple_w->tensor->ne[1]; + } // sanity check for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { @@ -201,8 +210,9 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { } if (hparams.is_ple(il)) { - layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0); - layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0); + const int64_t ple_dim = (int64_t) hparams.ple_head_dim * hparams.ple_n_heads; + layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { ple_dim, hc_dim }, 0); + layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { ple_dim, n_embd }, 0); layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, 0); layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0); layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, 0); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index ab05df77e1bd..f84bca8a5955 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -79,7 +79,7 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 return ret; } -static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { +static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const bool qwen_ple = false) { gguf_context_ptr ret(gguf_init_empty()); llama_model_saver ms(arch, ret.get()); const uint32_t n_ctx = 256; @@ -257,6 +257,18 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ratios[il] = 4; } ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, ratios); + + if (qwen_ple) { + ms.add_kv(LLM_KV_PLE_LAYERS, std::vector({0})); + ms.add_kv(LLM_KV_PLE_NGRAM_SIZE, uint32_t(2)); + ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, uint32_t(1)); + ms.add_kv(LLM_KV_PLE_CONV_KERNEL, uint32_t(2)); + ms.add_kv(LLM_KV_PLE_EOS_TOKEN_ID, uint32_t(1)); + ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, uint32_t(64)); + ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector({1, 3})); + ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS, std::vector({0})); + ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, std::vector({16})); + } } // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused @@ -747,6 +759,24 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg check_qwen4_qsa_mask, &check); get_logits(model_and_ctx.first.get(), model_and_ctx.second.get(), tokens); GGML_ASSERT(check.ok && check.n_seen > 0); + + gguf_context_ptr gguf_ctx_ple = get_gguf_ctx(arch, moe, true); + auto model_and_ctx_ple = get_model_and_ctx(gguf_ctx_ple.get(), nullptr, seed, {}); + const std::vector logits_ple = get_logits( + model_and_ctx_ple.first.get(), model_and_ctx_ple.second.get(), tokens); + + FILE * file_ple = tmpfile(); + GGML_ASSERT(file_ple); + llama_model_saver saver_ple(model_and_ctx_ple.first.get()); + saver_ple.add_kv_from_model(); + saver_ple.add_tensors_from_model(); + saver_ple.save(file_ple); + rewind(file_ple); + + auto model_and_ctx_ple_saved = get_model_and_ctx(nullptr, file_ple, seed, {}); + const std::vector logits_ple_saved = get_logits( + model_and_ctx_ple_saved.first.get(), model_and_ctx_ple_saved.second.get(), tokens); + GGML_ASSERT(logits_ple == logits_ple_saved); } std::pair model_and_ctx_cpu; std::vector logits_cpu; From 14b71e58ca78f515dfba2edd59414587625b712e Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 05:44:28 -0300 Subject: [PATCH 3/9] qwen4exp: validate model metadata Reject invalid GDN, hyper-connection, QSA, and PLE dimensions during model loading instead of aborting later while building the graph. Validate PLE array lengths before copying them into fixed storage. The released configuration and synthetic Qwen4 architecture tests pass. Assisted-by: Codex --- src/models/qwen4exp.cpp | 48 ++++++++++++++++++++++++++++++++--------- 1 file changed, 38 insertions(+), 10 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index a18cd867a5c3..cb18ca0366b3 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -18,21 +18,29 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); - GGML_ASSERT(hparams.ssm_d_conv > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 && - hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0); + if (hparams.ssm_d_conv == 0 || hparams.ssm_d_inner == 0 || hparams.ssm_d_state == 0 || + hparams.ssm_dt_rank == 0 || hparams.ssm_n_group == 0 || + hparams.ssm_dt_rank % hparams.ssm_n_group != 0 || + (uint64_t) hparams.ssm_d_state * hparams.ssm_dt_rank != hparams.ssm_d_inner) { + throw std::runtime_error("invalid Qwen4-Exp gated delta net dimensions"); + } // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank) ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); - GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0); + if (hparams.n_embd == 0 || hparams.dsv4_hc_mult <= 1 || hparams.hc_low_rank == 0 || + hparams.dsv4_hc_mult > UINT32_MAX/hparams.n_embd) { + throw std::runtime_error("invalid Qwen4-Exp hyper-connection dimensions"); + } hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); - GGML_ASSERT(hparams.indexer_n_head > 0 - && hparams.indexer_head_size > 0 - && hparams.indexer_top_k > 0); + if (hparams.indexer_n_head == 0 || hparams.indexer_head_size == 0 || hparams.indexer_top_k == 0 || + hparams.n_rot_full > hparams.indexer_head_size) { + throw std::runtime_error("invalid Qwen4-Exp sparse-attention dimensions"); + } ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false); // PLE n-gram hash embeddings; if the key group is absent every field stays zero @@ -44,9 +52,11 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { if (n_ple > 0) { std::vector ple_layers; ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers); - GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer"); + if (n_ple != 1 || ple_layers.size() != n_ple) { + throw std::runtime_error("Qwen4-Exp supports exactly one PLE layer"); + } for (uint32_t il : ple_layers) { - if (il >= hparams.n_layer_all) { + if (il >= hparams.n_layer()) { throw std::runtime_error(format("PLE layer %u is out of range", il)); } hparams.is_ple_impl.set(il); @@ -59,7 +69,9 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { // optional: files written before this key fall back to the EOS token ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false); ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); - GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0); + if (hparams.ple_conv_kernel == 0 || hparams.n_embd_per_layer == 0) { + throw std::runtime_error("invalid Qwen4-Exp PLE dimensions"); + } if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) { throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size)); @@ -72,6 +84,17 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { } hparams.ple_n_heads = (uint32_t) ple_n_heads; + uint32_t n_multipliers = 0; + uint32_t n_offsets = 0; + uint32_t n_vocab_sizes = 0; + ml.get_arr_n(LLM_KV_PLE_LAYER_MULTIPLIERS, n_multipliers); + ml.get_arr_n(LLM_KV_PLE_HEAD_OFFSETS, n_offsets); + ml.get_arr_n(LLM_KV_PLE_HEAD_VOCAB_SIZES, n_vocab_sizes); + if (n_multipliers != hparams.ple_ngram_size || + n_offsets != hparams.ple_n_heads || n_vocab_sizes != hparams.ple_n_heads) { + throw std::runtime_error("invalid Qwen4-Exp PLE metadata lengths"); + } + ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); // the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses @@ -95,7 +118,9 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { uint32_t full_attn_interval = 4; ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); - GGML_ASSERT(full_attn_interval > 0); + if (full_attn_interval == 0) { + throw std::runtime_error("invalid Qwen4-Exp full-attention interval"); + } for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0); } @@ -110,6 +135,9 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { } else if (ratio == 0 || hparams.indexer_top_k % ratio != 0) { throw std::runtime_error(format("invalid Qwen4-Exp QSA compression ratio %u at layer %u", ratio, il)); } + if (hparams.is_ple(il) && !hparams.is_recr(il)) { + throw std::runtime_error(format("Qwen4-Exp PLE layer %u is not recurrent", il)); + } } switch (hparams.n_layer()) { From 6ef243692bf8d0ecd05c309e198c8ca1eebe18cf Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 06:42:59 -0300 Subject: [PATCH 4/9] qwen4exp: update indexer cache after sequence copies Treat cached indexer keys as unrotated data and apply pending cache updates alongside the attention and recurrent state. This copies indexer data during non-unified cross-stream sequence copies without applying RoPE shifts to raw keys. Assisted-by: Codex --- src/llama-memory-hybrid-idx.cpp | 5 ++++- src/llama-memory-hybrid-idx.h | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 68450fc3e1d8..1575f0189349 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -49,6 +49,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( hparams_idx(model.hparams), mem_idx(filter_idx == nullptr ? nullptr : [&] { // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own + hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE; std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; @@ -597,7 +598,9 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( llama_context * lctx, bool optimize) : llama_memory_hybrid_context(mem, lctx, optimize), - mem(mem) {} + mem(mem), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + mem->get_mem_idx()->init_update(lctx, optimize)) {} llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( llama_memory_hybrid_idx * mem, diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 5f1ee7794303..746ce76ab12e 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -140,7 +140,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // llama_memory_hybrid_idx_context specific API // - // nullptr with no indexer, and for the update context, which builds no sparse graph + // nullptr with no indexer const llama_kv_cache_context * get_idx() const; // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified From edb6dec1c4468ad92dc81ebf83fe493feb27eefe Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 09:10:16 -0300 Subject: [PATCH 5/9] qwen4exp: enable recurrent state rollback Assisted-by: Codex --- src/llama-arch.cpp | 1 + tests/CMakeLists.txt | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5e61f61f7f0d..1db1f1835d1e 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1100,6 +1100,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { switch (arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fe3d14ffc552..7067f5f54678 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -220,6 +220,16 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) FIXTURES_REQUIRED generate-models ) + llama_test( + test-recurrent-state-rollback + NAME test-recurrent-state-rollback-qwen4exp + LABEL main + ARGS -m "${MODEL_DIR}/qwen4exp-moe.gguf" + ) + set_tests_properties(test-recurrent-state-rollback-qwen4exp PROPERTIES + FIXTURES_REQUIRED generate-models + ) + llama_test( test-recurrent-state-rollback NAME test-recurrent-state-rollback-nemotron-h From 8427625291a2493ea366230677073e5c0bf80fe3 Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 09:11:15 -0300 Subject: [PATCH 6/9] qwen4exp: disable tensor split Assisted-by: Codex --- src/llama-arch.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 1db1f1835d1e..6f781949d20c 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1142,6 +1142,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_QWEN3TTS: return false; default: From a7fc7e40f4b29835cae66eaa0499b59d7f808cb8 Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Fri, 28 Aug 2026 11:52:04 -0300 Subject: [PATCH 7/9] qwen4exp: fix GDN QK normalization Normalize GDN queries and keys with rsqrt(sum(x^2) + eps), matching the reference instead of ggml_l2_norm's max(sqrt(sum), eps) convention. Assisted-by: Codex --- src/models/qwen4exp.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index cb18ca0366b3..32f510c7a8a4 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -5,6 +5,12 @@ #include #include +#include + +static ggml_tensor * qwen4_l2_norm(ggml_context * ctx, ggml_tensor * input, float eps) { + const float n = input->ne[0]; + return ggml_scale(ctx, ggml_rms_norm(ctx, input, eps/n), 1.0f/std::sqrt(n)); +} void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); @@ -845,8 +851,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( const float eps_norm = hparams.f_norm_rms_eps; - q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); - k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + q_conv = qwen4_l2_norm(ctx0, q_conv, eps_norm); + k_conv = qwen4_l2_norm(ctx0, k_conv, eps_norm); // repeat to match shapes when head keys != value keys; unneeded with the fused GDN if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { From dfc79327fc8e37f2db68a122fa04b8c4c96b61ab Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 29 Aug 2026 06:56:23 -0300 Subject: [PATCH 8/9] qwen4exp: avoid CUDA RMSNorm grid limit Assisted-by: Codex --- src/models/qwen4exp.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 32f510c7a8a4..f7ce7c599fb7 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -595,13 +595,14 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask( block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys)); block_keys = ggml_mean(ctx0, block_keys); block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys)); - block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, 1, n_blocks*n_tps); + block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, n_blocks*n_tps, 1); if (block_keys->type != activation_type) { block_keys = ggml_cast(ctx0, block_keys, activation_type); block_keys = ggml_cast(ctx0, block_keys, GGML_TYPE_F32); } block_keys = build_norm(block_keys, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, 1, n_blocks*n_tps); ggml_tensor * block_pos = ggml_view_3d(ctx0, inp->block_pos, n_blocks, 4, n_tps, inp->block_pos->nb[1], inp->block_pos->nb[2], is*n_tps*inp->block_pos->nb[2]); From 50d6b3db314d4f7161f2152c6aaef3e3423bd804 Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Sat, 29 Aug 2026 13:23:03 -0300 Subject: [PATCH 9/9] qwen4exp: share QSA block tables across queries Assisted-by: Codex --- src/llama-memory-hybrid-idx.cpp | 207 +++++++++++++++++++++++++++----- src/llama-memory-hybrid-idx.h | 12 +- src/models/qwen4exp.cpp | 174 ++++++++++++++++++--------- tests/test-llama-archs.cpp | 19 ++- 4 files changed, 317 insertions(+), 95 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 1575f0189349..e658af09f9bf 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -447,41 +447,54 @@ void llama_memory_hybrid_idx::set_input_qsa( ggml_tensor * block_cells, ggml_tensor * block_pos, ggml_tensor * block_mask, - ggml_tensor * selected, + ggml_tensor * tail_cells, + ggml_tensor * tail_fallback, const ggml_tensor * kq_mask, const llama_ubatch * ubatch, - uint32_t ratio, - uint32_t block_topk) const { + uint32_t n_groups, + uint32_t ratio) const { GGML_ASSERT(ggml_backend_buffer_is_host(block_cells->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(block_pos->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(block_mask->buffer)); - GGML_ASSERT(ggml_backend_buffer_is_host(selected->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(tail_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(tail_fallback->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(kq_mask->buffer)); const int64_t n_blocks = block_cells->ne[1]; const int64_t n_tokens = ubatch->n_tokens; - const int64_t n_pos = block_pos->ne[1]; - const int64_t n_kv = selected->ne[0]; + const int64_t n_pos = n_groups > 0 ? 4 : block_pos->ne[1]; + const int64_t n_kv = kq_mask->ne[0]; + const int64_t n_tail = tail_cells->ne[0]; + const int64_t n_layout = n_groups > 0 ? n_groups : n_tokens; GGML_ASSERT(block_cells->type == GGML_TYPE_I32); GGML_ASSERT(block_pos->type == GGML_TYPE_I32); GGML_ASSERT(block_mask->type == GGML_TYPE_F32); - GGML_ASSERT(selected->type == GGML_TYPE_F32); - GGML_ASSERT(block_cells->ne[0] == ratio && block_cells->ne[2] == n_tokens); - GGML_ASSERT(block_pos->ne[0] == n_blocks && block_pos->ne[2] == n_tokens); + GGML_ASSERT(tail_cells->type == GGML_TYPE_I32); + GGML_ASSERT(tail_fallback->type == GGML_TYPE_F32); + GGML_ASSERT(block_cells->ne[0] == ratio && block_cells->ne[2] == n_layout); + if (n_groups > 0) { + GGML_ASSERT(block_pos->ne[0] == n_blocks*n_pos && block_pos->ne[1] == n_layout); + } else { + GGML_ASSERT(block_pos->ne[0] == n_blocks && block_pos->ne[2] == n_layout); + } GGML_ASSERT(block_mask->ne[0] == n_blocks && block_mask->ne[1] == n_tokens); - GGML_ASSERT(selected->ne[1] == n_tokens); + GGML_ASSERT(n_tail == std::max(ratio - 1, 1) && tail_cells->ne[1] == n_tokens); + GGML_ASSERT(tail_fallback->ne[0] == n_tail && tail_fallback->ne[1] == n_tokens); GGML_ASSERT(kq_mask->ne[0] == n_kv); + GGML_ASSERT(n_groups == 0 || n_tokens % n_groups == 0); - int32_t * cell_data = (int32_t *) block_cells->data; - int32_t * pos_data = (int32_t *) block_pos->data; - float * mask_data = (float *) block_mask->data; - float * selected_data = (float *) selected->data; + int32_t * cell_data = (int32_t *) block_cells->data; + int32_t * pos_data = (int32_t *) block_pos->data; + float * mask_data = (float *) block_mask->data; + int32_t * tail_data = (int32_t *) tail_cells->data; + float * tail_fallback_data = (float *) tail_fallback->data; std::fill(cell_data, cell_data + ggml_nelements(block_cells), 0); std::fill(pos_data, pos_data + ggml_nelements(block_pos), 0); std::fill(mask_data, mask_data + ggml_nelements(block_mask), -INFINITY); - std::fill(selected_data, selected_data + ggml_nelements(selected), 0.0f); + std::fill(tail_data, tail_data + ggml_nelements(tail_cells), 0); + std::fill(tail_fallback_data, tail_fallback_data + ggml_nelements(tail_fallback), 1.0f); auto mask_visible = [&](int64_t query, uint32_t cell) { const int64_t index = query*n_kv + cell; @@ -491,11 +504,11 @@ void llama_memory_hybrid_idx::set_input_qsa( return std::isfinite(((const float *) kq_mask->data)[index]); }; - for (int64_t iq = 0; iq < n_tokens; ++iq) { - const llama_seq_id seq_id = ubatch->seq_id[iq][0]; + using mapped_token = std::pair; + auto map_sequence = [&](llama_seq_id seq_id, std::vector & mapped) { const auto found = qsa_histories.find(seq_id); if (found == qsa_histories.end()) { - continue; + return false; } const auto & cells = get_mem_attn()->get_cells(seq_id); @@ -511,7 +524,7 @@ void llama_memory_hybrid_idx::set_input_qsa( } std::map next_cell; - std::vector> visible; + mapped.reserve(found->second.size()); for (const auto & token : found->second) { const pos_key key = { token.pos[0], token.pos[1], token.pos[2] }; auto cells = cells_by_pos.find(key); @@ -524,9 +537,80 @@ void llama_memory_hybrid_idx::set_input_qsa( continue; } - const uint32_t cell = cells->second[index++]; - if (mask_visible(iq, cell)) { - visible.emplace_back(&token, cell); + mapped.emplace_back(&token, cells->second[index++]); + } + + return true; + }; + + if (n_groups > 0) { + // Share one logical block table across the queries of each sequence. + const int64_t n_qpg = n_tokens/n_groups; + for (uint32_t ig = 0; ig < n_groups; ++ig) { + const int64_t iq0 = ig*n_qpg; + const llama_seq_id seq_id = ubatch->seq_id[iq0][0]; + std::vector mapped; + GGML_ASSERT(map_sequence(seq_id, mapped) && !mapped.empty()); + + const size_t n_complete = mapped.size()/ratio; + const size_t n_write = std::min(n_complete, n_blocks); + std::vector used_cells(n_kv, 0); + for (size_t ib = 0; ib < n_write; ++ib) { + for (uint32_t ir = 0; ir < ratio; ++ir) { + const uint32_t cell = mapped[ib*ratio + ir].second; + cell_data[(ig*n_blocks + ib)*ratio + ir] = cell; + used_cells[cell] = 1; + } + for (int64_t ip = 0; ip < n_pos; ++ip) { + pos_data[(ig*n_pos + ip)*n_blocks + ib] = mapped[ib*ratio].first->pos[ip]; + } + } + + uint32_t fallback_cell = 0; + for (size_t ib = n_write; ib < (size_t) n_blocks; ++ib) { + for (uint32_t ir = 0; ir < ratio; ++ir) { + while (fallback_cell < used_cells.size() && used_cells[fallback_cell]) { + ++fallback_cell; + } + GGML_ASSERT(fallback_cell < used_cells.size()); + cell_data[(ig*n_blocks + ib)*ratio + ir] = fallback_cell; + used_cells[fallback_cell++] = 1; + } + } + + for (int64_t iq = iq0; iq < iq0 + n_qpg; ++iq) { + const size_t n_visible = std::partition_point(mapped.begin(), mapped.end(), [&](const mapped_token & token) { + return mask_visible(iq, token.second); + }) - mapped.begin(); + const size_t n_query_complete = n_visible/ratio; + for (size_t ib = 0; ib < std::min(n_query_complete, n_blocks); ++ib) { + mask_data[iq*n_blocks + ib] = 0.0f; + } + + const size_t tail_start = n_query_complete*ratio; + for (size_t iv = tail_start; iv < n_visible; ++iv) { + const size_t it = iq*n_tail + iv - tail_start; + tail_data[it] = mapped[iv].second; + tail_fallback_data[it] = 0.0f; + } + } + } + + return; + } + + for (int64_t iq = 0; iq < n_tokens; ++iq) { + const llama_seq_id seq_id = ubatch->seq_id[iq][0]; + std::vector mapped; + if (!map_sequence(seq_id, mapped)) { + continue; + } + + std::vector visible; + visible.reserve(mapped.size()); + for (const auto & token : mapped) { + if (mask_visible(iq, token.second)) { + visible.push_back(token); } } @@ -557,9 +641,11 @@ void llama_memory_hybrid_idx::set_input_qsa( } } - const size_t selected_start = n_complete <= block_topk ? 0 : n_complete*ratio; - for (size_t iv = selected_start; iv < visible.size(); ++iv) { - selected_data[iq*n_kv + visible[iv].second] = 1.0f; + const size_t tail_start = n_complete*ratio; + for (size_t iv = tail_start; iv < visible.size(); ++iv) { + const size_t it = iq*n_tail + iv - tail_start; + tail_data[it] = visible[iv].second; + tail_fallback_data[it] = 0.0f; } } } @@ -649,18 +735,79 @@ uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { return ns_ubatch[i_cur]; } +uint32_t llama_memory_hybrid_idx_context::get_qsa_compact_groups(const llama_ubatch & ubatch) const { + // A shared table is exact when every query sees a prefix of one monotonic sequence history. + if (!ubatch.equal_seqs() || ubatch.n_seqs == 0 || ubatch.n_seqs != ubatch.n_seqs_unq || + ubatch.n_tokens != ubatch.n_seq_tokens*ubatch.n_seqs) { + return 0; + } + + const uint32_t n_stream = get_n_stream(); + if (n_stream > 1 && ubatch.n_seqs != n_stream) { + return 0; + } + + std::vector group_seqs; + group_seqs.reserve(ubatch.n_seqs); + std::map seen; + for (uint32_t ig = 0; ig < ubatch.n_seqs; ++ig) { + const uint32_t first = ig*ubatch.n_seq_tokens; + if (ubatch.n_seq_id[first] != 1) { + return 0; + } + + const llama_seq_id seq_id = ubatch.seq_id[first][0]; + if (!seen.emplace(seq_id, true).second || (n_stream > 1 && seq_id != ubatch.seq_id_unq[ig])) { + return 0; + } + + for (uint32_t i = first + 1; i < first + ubatch.n_seq_tokens; ++i) { + if (ubatch.n_seq_id[i] != 1 || ubatch.seq_id[i][0] != seq_id) { + return 0; + } + } + group_seqs.push_back(seq_id); + } + + if (!has_ubatches) { + return ubatch.n_seqs; + } + + auto pos_less = [](const llama_memory_hybrid_idx::qsa_token & a, + const llama_memory_hybrid_idx::qsa_token & b) { + return std::tie(a.pos[0], a.pos[1], a.pos[2]) < std::tie(b.pos[0], b.pos[1], b.pos[2]); + }; + + for (llama_seq_id seq_id : group_seqs) { + const auto found = mem->qsa_histories.find(seq_id); + if (found == mem->qsa_histories.end() || found->second.empty()) { + return 0; + } + + const auto & history = found->second; + for (size_t i = 1; i < history.size(); ++i) { + if (pos_less(history[i], history[i - 1])) { + return 0; + } + } + } + + return ubatch.n_seqs; +} + void llama_memory_hybrid_idx_context::set_input_qsa( ggml_tensor * block_cells, ggml_tensor * block_pos, ggml_tensor * block_mask, - ggml_tensor * selected, + ggml_tensor * tail_cells, + ggml_tensor * tail_fallback, const ggml_tensor * kq_mask, const llama_ubatch * ubatch, - uint32_t ratio, - uint32_t block_topk) const { + uint32_t n_groups, + uint32_t ratio) const { GGML_ASSERT(ratio > 0); GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); - mem->set_input_qsa(block_cells, block_pos, block_mask, selected, - kq_mask, ubatch, ratio, block_topk); + mem->set_input_qsa(block_cells, block_pos, block_mask, tail_cells, tail_fallback, + kq_mask, ubatch, n_groups, ratio); } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 746ce76ab12e..97f641f0fefc 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -78,13 +78,15 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer void set_input_qsa(ggml_tensor * block_cells, ggml_tensor * block_pos, - ggml_tensor * block_mask, ggml_tensor * selected, + ggml_tensor * block_mask, ggml_tensor * tail_cells, ggml_tensor * tail_fallback, const ggml_tensor * kq_mask, const llama_ubatch * ubatch, - uint32_t ratio, uint32_t block_topk) const; + uint32_t n_groups, uint32_t ratio) const; void commit_qsa_tokens(const llama_ubatch & ubatch); private: + friend class llama_memory_hybrid_idx_context; + struct qsa_token { std::array pos; }; @@ -146,11 +148,13 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified uint32_t get_n_stream() const; + uint32_t get_qsa_compact_groups(const llama_ubatch & ubatch) const; + // QSA blocks follow each sequence's token order, not physical cells or scalar positions. void set_input_qsa(ggml_tensor * block_cells, ggml_tensor * block_pos, - ggml_tensor * block_mask, ggml_tensor * selected, + ggml_tensor * block_mask, ggml_tensor * tail_cells, ggml_tensor * tail_fallback, const ggml_tensor * kq_mask, const llama_ubatch * ubatch, - uint32_t ratio, uint32_t block_topk) const; + uint32_t n_groups, uint32_t ratio) const; private: llama_memory_hybrid_idx * mem = nullptr; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 09e4bdfff7f5..98e08edbb88c 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -482,14 +482,16 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { llm_graph_input_qsa( const llama_memory_hybrid_idx_context * mctx, ggml_tensor * kq_mask, + uint32_t n_groups, + bool allow_compact, uint32_t ratio, uint32_t block_topk) : - mctx(mctx), kq_mask(kq_mask), ratio(ratio), block_topk(block_topk) {} + mctx(mctx), kq_mask(kq_mask), n_groups(n_groups), allow_compact(allow_compact), ratio(ratio), block_topk(block_topk) {} virtual ~llm_graph_input_qsa() = default; void set_input(const llama_ubatch * ubatch) override { - mctx->set_input_qsa(block_cells, block_pos, block_mask, selected, - kq_mask, ubatch, ratio, block_topk); + mctx->set_input_qsa(block_cells, block_pos, block_mask, tail_cells, tail_fallback, + kq_mask, ubatch, n_groups, ratio); } bool can_reuse(const llm_graph_params & params) override { @@ -502,16 +504,25 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { const int64_t n_kv = idx->get_n_kv(); const int64_t n_blocks = n_kv/ratio; + const uint32_t new_groups = allow_compact ? mctx->get_qsa_compact_groups(params.ubatch) : 0; bool res = true; res &= n_kv > (int64_t) block_topk*ratio + ratio - 1; res &= block_cells->ne[1] == n_blocks; - res &= block_cells->ne[2] == params.ubatch.n_tokens; - res &= block_pos->ne[2] == params.ubatch.n_tokens; + res &= block_cells->ne[2] == (new_groups > 0 ? new_groups : params.ubatch.n_tokens); + if (new_groups > 0) { + res &= block_pos->ne[0] == 4*n_blocks; + res &= block_pos->ne[1] == new_groups; + } else { + res &= block_pos->ne[0] == n_blocks; + res &= block_pos->ne[2] == params.ubatch.n_tokens; + } res &= block_mask->ne[1] == params.ubatch.n_tokens; - res &= selected->ne[0] == n_kv; - res &= selected->ne[1] == params.ubatch.n_tokens; + res &= tail_cells->ne[0] == std::max(ratio - 1, 1); + res &= tail_cells->ne[1] == params.ubatch.n_tokens; + res &= tail_fallback->ne[1] == params.ubatch.n_tokens; + res &= new_groups == n_groups; return res; } @@ -519,10 +530,13 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { ggml_tensor * block_cells = nullptr; ggml_tensor * block_pos = nullptr; ggml_tensor * block_mask = nullptr; - ggml_tensor * selected = nullptr; + ggml_tensor * tail_cells = nullptr; + ggml_tensor * tail_fallback = nullptr; const llama_memory_hybrid_idx_context * mctx; ggml_tensor * kq_mask; + const uint32_t n_groups; + const bool allow_compact; const uint32_t ratio; const uint32_t block_topk; }; @@ -548,7 +562,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask( const int64_t n_stream = mctx_hyb->get_n_stream(); GGML_ASSERT(n_tokens % n_stream == 0); - const int64_t n_tps = n_tokens/n_stream; + const bool allow_compact = cparams.causal_attn && !hparams.use_alibi; + const uint32_t n_groups = allow_compact ? mctx_hyb->get_qsa_compact_groups(ubatch) : 0; + const int64_t n_layout = n_groups > 0 ? n_groups : n_tokens; // nothing above depends on the layer, so the layers sharing a ratio share one input set llm_graph_input_qsa * inp = nullptr; @@ -558,17 +574,21 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask( inp = it->second; } else { auto qsa = std::make_unique( - mctx_hyb, kq_mask, (uint32_t) r, (uint32_t) block_topk); + mctx_hyb, kq_mask, n_groups, allow_compact, (uint32_t) r, (uint32_t) block_topk); - qsa->block_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, r, n_blocks, n_tokens); - qsa->block_pos = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, n_blocks, 4, n_tokens); + qsa->block_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, r, n_blocks, n_layout); + qsa->block_pos = n_groups > 0 ? + ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 4*n_blocks, n_layout) : + ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, n_blocks, 4, n_layout); qsa->block_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_blocks, n_tokens); - qsa->selected = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); + qsa->tail_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, std::max(r - 1, 1), n_tokens); + qsa->tail_fallback = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, std::max(r - 1, 1), n_tokens); ggml_set_input(qsa->block_cells); ggml_set_input(qsa->block_pos); ggml_set_input(qsa->block_mask); - ggml_set_input(qsa->selected); + ggml_set_input(qsa->tail_cells); + ggml_set_input(qsa->tail_fallback); inp = qsa.get(); res->add_input(std::move(qsa)); @@ -589,84 +609,120 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask( cb(q, "indexer_q", il); const ggml_type activation_type = mctx_idx->type_k(); - std::vector selected_streams; - selected_streams.reserve(n_stream); - - for (int64_t is = 0; is < n_stream; ++is) { + const bool compact = inp->n_groups > 0; + const int64_t n_selection_groups = compact ? inp->n_groups : n_stream; + GGML_ASSERT(n_tokens % n_selection_groups == 0); + GGML_ASSERT(!compact || n_stream == 1 || n_selection_groups == n_stream); + const int64_t n_qpg = n_tokens/n_selection_groups; + const int64_t n_block_layout = compact ? 1 : n_qpg; + std::vector selected_groups; + selected_groups.reserve(n_selection_groups); + + for (int64_t ig = 0; ig < n_selection_groups; ++ig) { + const int64_t is = compact && n_stream == 1 ? 0 : ig; ggml_tensor * cache = ggml_view_2d(ctx0, k_all, idx_dim, n_kv, k_all->nb[1], is*k_all->nb[2]); - ggml_tensor * block_cells = ggml_view_3d(ctx0, inp->block_cells, r, n_blocks, n_tps, - inp->block_cells->nb[1], inp->block_cells->nb[2], is*n_tps*inp->block_cells->nb[2]); + ggml_tensor * block_cells; + ggml_tensor * block_pos; + if (compact) { + if (n_selection_groups == 1) { + block_cells = inp->block_cells; + block_pos = inp->block_pos; + } else { + block_cells = ggml_view_2d(ctx0, inp->block_cells, r, n_blocks, + inp->block_cells->nb[1], ig*inp->block_cells->nb[2]); + block_pos = ggml_view_1d(ctx0, inp->block_pos, 4*n_blocks, ig*inp->block_pos->nb[1]); + } + } else { + block_cells = ggml_view_3d(ctx0, inp->block_cells, r, n_blocks, n_qpg, + inp->block_cells->nb[1], inp->block_cells->nb[2], ig*n_qpg*inp->block_cells->nb[2]); + block_pos = ggml_view_3d(ctx0, inp->block_pos, n_blocks, 4, n_qpg, + inp->block_pos->nb[1], inp->block_pos->nb[2], ig*n_qpg*inp->block_pos->nb[2]); + } ggml_tensor * block_keys = ggml_get_rows(ctx0, cache, - ggml_reshape_1d(ctx0, block_cells, r*n_blocks*n_tps)); - block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, r, n_blocks, n_tps); + ggml_reshape_1d(ctx0, block_cells, r*n_blocks*n_block_layout)); + block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, r, n_blocks, n_block_layout); block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys)); block_keys = ggml_mean(ctx0, block_keys); block_keys = ggml_cont(ctx0, ggml_transpose(ctx0, block_keys)); - block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, n_blocks*n_tps, 1); + block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, n_blocks*n_block_layout, 1); if (block_keys->type != activation_type) { block_keys = ggml_cast(ctx0, block_keys, activation_type); block_keys = ggml_cast(ctx0, block_keys, GGML_TYPE_F32); } block_keys = build_norm(block_keys, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); - block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, 1, n_blocks*n_tps); + block_keys = ggml_reshape_3d(ctx0, block_keys, idx_dim, 1, n_blocks*n_block_layout); - ggml_tensor * block_pos = ggml_view_3d(ctx0, inp->block_pos, n_blocks, 4, n_tps, - inp->block_pos->nb[1], inp->block_pos->nb[2], is*n_tps*inp->block_pos->nb[2]); - block_keys = ggml_rope_multi(ctx0, block_keys, - ggml_reshape_1d(ctx0, block_pos, n_blocks*4*n_tps), nullptr, + ggml_tensor * block_pos_flat = compact ? block_pos : + ggml_reshape_1d(ctx0, block_pos, n_blocks*4*n_block_layout); + block_keys = ggml_rope_multi(ctx0, block_keys, block_pos_flat, nullptr, n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); cb(block_keys, "indexer_k", il); - block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, n_blocks, 1, n_tps); - ggml_tensor * query = ggml_view_3d(ctx0, q, idx_dim, n_idx_h, n_tps, - q->nb[1], q->nb[2], is*n_tps*q->nb[2]); - query = ggml_reshape_4d(ctx0, query, idx_dim, n_idx_h, 1, n_tps); + block_keys = ggml_reshape_4d(ctx0, block_keys, idx_dim, n_blocks, 1, n_block_layout); + ggml_tensor * query = ggml_view_3d(ctx0, q, idx_dim, n_idx_h, n_qpg, + q->nb[1], q->nb[2], ig*n_qpg*q->nb[2]); + query = ggml_reshape_4d(ctx0, query, idx_dim, n_idx_h, compact ? n_qpg : 1, compact ? 1 : n_qpg); ggml_tensor * scores = ggml_mul_mat(ctx0, block_keys, query); ggml_mul_mat_set_prec(scores, GGML_PREC_F32); scores = ggml_relu(ctx0, scores); scores = ggml_sum_rows(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3))); - scores = ggml_scale(ctx0, ggml_reshape_2d(ctx0, scores, n_blocks, n_tps), + scores = ggml_scale(ctx0, ggml_reshape_2d(ctx0, scores, n_blocks, n_qpg), 1.0f/sqrtf((float) idx_dim)); - ggml_tensor * block_mask = ggml_view_2d(ctx0, inp->block_mask, n_blocks, n_tps, - inp->block_mask->nb[1], is*n_tps*inp->block_mask->nb[1]); + ggml_tensor * block_mask = n_selection_groups == 1 ? inp->block_mask : + ggml_view_2d(ctx0, inp->block_mask, n_blocks, n_qpg, + inp->block_mask->nb[1], ig*n_qpg*inp->block_mask->nb[1]); scores = ggml_add(ctx0, scores, block_mask); cb(scores, "indexer_score", il); ggml_tensor * top_blocks = ggml_top_k(ctx0, scores, block_topk); - ggml_tensor * top_cells = ggml_get_rows(ctx0, block_cells, top_blocks); - top_cells = ggml_reshape_2d(ctx0, top_cells, r*block_topk, n_tps); - - ggml_tensor * base_selected = ggml_view_2d(ctx0, inp->selected, n_kv, n_tps, - inp->selected->nb[1], is*n_tps*inp->selected->nb[1]); - base_selected = ggml_reshape_3d(ctx0, base_selected, 1, n_kv, n_tps); - ggml_tensor * selected_top = ggml_fill(ctx0, base_selected, 0.0f); - ggml_tensor * ones = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, r*block_topk, n_tps); - ones = ggml_fill(ctx0, ones, 1.0f); - selected_top = ggml_set_rows(ctx0, selected_top, ones, top_cells); - ggml_tensor * selected_stream = ggml_clamp( - ctx0, ggml_add(ctx0, base_selected, selected_top), 0.0f, 1.0f); - selected_streams.push_back(ggml_reshape_2d(ctx0, selected_stream, n_kv, n_tps)); - } + ggml_tensor * top_cells; + if (compact) { + ggml_tensor * top_rows = ggml_reshape_1d( + ctx0, ggml_cont(ctx0, top_blocks), block_topk*n_qpg); + top_cells = ggml_get_rows(ctx0, block_cells, top_rows); + } else { + top_cells = ggml_get_rows(ctx0, block_cells, top_blocks); + } + top_cells = ggml_reshape_2d(ctx0, top_cells, r*block_topk, n_qpg); + + ggml_tensor * selected_cells = top_cells; + if (r > 1) { + ggml_tensor * tail_cells = n_selection_groups == 1 ? inp->tail_cells : + ggml_view_2d(ctx0, inp->tail_cells, r - 1, n_qpg, + inp->tail_cells->nb[1], ig*n_qpg*inp->tail_cells->nb[1]); + ggml_tensor * tail_fallback = n_selection_groups == 1 ? inp->tail_fallback : + ggml_view_2d(ctx0, inp->tail_fallback, r - 1, n_qpg, + inp->tail_fallback->nb[1], ig*n_qpg*inp->tail_fallback->nb[1]); + ggml_tensor * top_cell = ggml_view_2d(ctx0, top_cells, 1, n_qpg, top_cells->nb[1], 0); + top_cell = ggml_cast(ctx0, top_cell, GGML_TYPE_F32); + tail_fallback = ggml_mul(ctx0, ggml_repeat(ctx0, top_cell, tail_fallback), tail_fallback); + tail_cells = ggml_cast(ctx0, tail_cells, GGML_TYPE_F32); + tail_cells = ggml_cast(ctx0, ggml_add(ctx0, tail_cells, tail_fallback), GGML_TYPE_I32); + selected_cells = ggml_concat(ctx0, selected_cells, tail_cells, 0); + } - ggml_tensor * selected = selected_streams[0]; - for (int64_t is = 1; is < n_stream; ++is) { - selected = ggml_concat(ctx0, selected, selected_streams[is], 1); + ggml_tensor * selected_all = n_selection_groups == 1 ? kq_mask : + ggml_view_2d(ctx0, kq_mask, n_kv, n_qpg, + kq_mask->nb[1], ig*n_qpg*kq_mask->nb[1]); + selected_all = ggml_reshape_3d(ctx0, ggml_fill(ctx0, selected_all, -1e30f), 1, n_kv, n_qpg); + ggml_tensor * zeros = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, selected_cells->ne[0], n_qpg); + zeros = ggml_fill(ctx0, zeros, 0.0f); + selected_all = ggml_set_rows(ctx0, selected_all, zeros, selected_cells); + selected_groups.push_back(ggml_reshape_2d(ctx0, selected_all, n_kv, n_qpg)); } - selected = ggml_scale_bias(ctx0, selected, 1e30f, -1e30f); - ggml_tensor * base_mask = ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens); - if (base_mask->type != GGML_TYPE_F32) { - base_mask = ggml_cast(ctx0, base_mask, GGML_TYPE_F32); + ggml_tensor * selected = selected_groups[0]; + for (int64_t ig = 1; ig < n_selection_groups; ++ig) { + selected = ggml_concat(ctx0, selected, selected_groups[ig], 1); } + ggml_tensor * base_mask = n_selection_groups == 1 ? kq_mask : + ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens); ggml_tensor * mask = ggml_add(ctx0, base_mask, selected); - if (cparams.flash_attn) { - mask = ggml_cast(ctx0, mask, GGML_TYPE_F16); - } cb(mask, "qsa_mask", il); return mask; } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index f84bca8a5955..941fbd9ba6f4 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -433,10 +433,21 @@ static std::vector get_logits( struct qwen4_qsa_mask_check { int64_t n_seen = 0; int64_t n_tokens_seen = 0; + int64_t n_indexer_keys_seen = 0; + int64_t n_indexer_key_rows_max = 0; bool ok = true; }; static bool check_qwen4_qsa_mask(ggml_tensor * tensor, bool ask, void * user_data) { + auto * check = (qwen4_qsa_mask_check *) user_data; + if (strncmp(tensor->name, "indexer_k-", 10) == 0) { + if (!ask) { + check->n_indexer_keys_seen++; + check->n_indexer_key_rows_max = std::max(check->n_indexer_key_rows_max, tensor->ne[2]); + } + return true; + } + if (strncmp(tensor->name, "qsa_mask", 8) != 0) { return false; } @@ -444,9 +455,13 @@ static bool check_qwen4_qsa_mask(ggml_tensor * tensor, bool ask, void * user_dat return true; } - auto * check = (qwen4_qsa_mask_check *) user_data; const int64_t n_kv = tensor->ne[0]; const int64_t n_tokens = tensor->ne[1]; + if (check->n_indexer_key_rows_max > n_kv/4) { + fprintf(stderr, "Qwen4 QSA materialized %lld block-key rows, expected at most %lld\n", + (long long) check->n_indexer_key_rows_max, (long long) (n_kv/4)); + check->ok = false; + } std::vector mask(ggml_nelements(tensor)); if (tensor->type == GGML_TYPE_F32) { ggml_backend_tensor_get(tensor, mask.data(), 0, ggml_nbytes(tensor)); @@ -758,7 +773,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, check_qwen4_qsa_mask, &check); get_logits(model_and_ctx.first.get(), model_and_ctx.second.get(), tokens); - GGML_ASSERT(check.ok && check.n_seen > 0); + GGML_ASSERT(check.ok && check.n_seen > 0 && check.n_indexer_keys_seen > 0); gguf_context_ptr gguf_ctx_ple = get_gguf_ctx(arch, moe, true); auto model_and_ctx_ple = get_model_and_ctx(gguf_ctx_ple.get(), nullptr, seed, {});