diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5e61f61f7f0d..6f781949d20c 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: @@ -1141,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: diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d4e59d77e570..e658af09f9bf 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 @@ -47,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; @@ -137,6 +140,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 +154,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 +252,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 +267,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 +284,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 +316,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 +361,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 +414,242 @@ 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 * tail_cells, + ggml_tensor * tail_fallback, + const ggml_tensor * kq_mask, + const llama_ubatch * ubatch, + 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(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 = 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(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(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; + 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(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; + 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]); + }; + + 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()) { + return false; + } + + 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; + 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); + if (cells == cells_by_pos.end()) { + continue; + } + + size_t & index = next_cell[key]; + if (index >= cells->second.size()) { + continue; + } + + 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); + } + } + + 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 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; + } + } +} + // // llama_memory_hybrid_idx_context // @@ -295,7 +684,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, @@ -307,7 +698,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 +718,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; } @@ -339,127 +735,79 @@ uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { return ns_ubatch[i_cur]; } -void llama_memory_hybrid_idx_context::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 { - 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); - } - } +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; } - // 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; - } + const uint32_t n_stream = get_n_stream(); + if (n_stream > 1 && ubatch.n_seqs != n_stream) { + return 0; + } - blk_of[j] = (int32_t) b; - cur_blk_cells[b*r + (p%r)] = (int32_t) j; - filled[b]++; + 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; } - GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); + 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; + } - // 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; + 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; } - cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; } + group_seqs.push_back(seq_id); + } - 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 (!has_ubatches) { + return ubatch.n_seqs; + } - 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; + 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 (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); - } + 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; + } - continue; + const auto & history = found->second; + for (size_t i = 1; i < history.size(); ++i) { + if (pos_less(history[i], history[i - 1])) { + return 0; } + } + } - float * cur_bias = dst_bias + i*n_kv; - - for (int64_t j = 0; j < n_kv; ++j) { - float v = -INFINITY; + return ubatch.n_seqs; +} - 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); - } +void llama_memory_hybrid_idx_context::set_input_qsa( + ggml_tensor * block_cells, + ggml_tensor * block_pos, + ggml_tensor * block_mask, + ggml_tensor * tail_cells, + ggml_tensor * tail_fallback, + const ggml_tensor * kq_mask, + const llama_ubatch * ubatch, + uint32_t n_groups, + uint32_t ratio) const { + GGML_ASSERT(ratio > 0); + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); - cur_bias[j] = v; - } - } - } + 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 e3472646d0f6..97f641f0fefc 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,22 @@ 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 * tail_cells, ggml_tensor * tail_fallback, + const ggml_tensor * kq_mask, const llama_ubatch * ubatch, + 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; + }; + + 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 +102,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 { @@ -123,26 +142,22 @@ 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 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; + 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 * tail_cells, ggml_tensor * tail_fallback, + const ggml_tensor * kq_mask, const llama_ubatch * ubatch, + uint32_t n_groups, uint32_t ratio) 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 +166,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 9b87a40d5af9..7b7c5ef48f4f 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 abf6a0502fbf..98e08edbb88c 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); @@ -18,21 +24,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 +58,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,15 +75,30 @@ 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"); + } - 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; + + 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); @@ -93,12 +124,28 @@ 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); } } + 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)); + } + 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()) { case 48: type = LLM_TYPE_A3B; break; default: type = LLM_TYPE_UNKNOWN; @@ -127,8 +174,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) { @@ -190,8 +244,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); @@ -424,13 +479,19 @@ 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 n_groups, + bool allow_compact, + uint32_t ratio, + uint32_t 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->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, tail_cells, tail_fallback, + kq_mask, ubatch, n_groups, ratio); } bool can_reuse(const llm_graph_params & params) override { @@ -442,39 +503,45 @@ 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; + const uint32_t new_groups = allow_compact ? mctx->get_qsa_compact_groups(params.ubatch) : 0; 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] == (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 &= 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; } - // 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 * 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; - - // 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, @@ -490,20 +557,14 @@ 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; + 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; @@ -512,59 +573,33 @@ 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); - - 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); - - ggml_set_input(qsa->cell_blk); - ggml_set_input(qsa->blk_cells); - ggml_set_input(qsa->blk_pos); - ggml_set_input(qsa->bias); + auto qsa = std::make_unique( + 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_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->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->tail_cells); + ggml_set_input(qsa->tail_fallback); 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); @@ -573,128 +608,123 @@ 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); - } - - if (inp->self_v_rot) { - v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); - } - - // 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(); + const ggml_type activation_type = mctx_idx->type_k(); + 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_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_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_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_block_layout); + + 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_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_qpg), + 1.0f/sqrtf((float) idx_dim)); + + 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; + 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_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 * 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)); } - 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 * selected = selected_groups[0]; + for (int64_t ig = 1; ig < n_selection_groups; ++ig) { + selected = ggml_concat(ctx0, selected, selected_groups[ig], 1); } - - return cur; + 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); + cb(mask, "qsa_mask", il); + return mask; } ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( @@ -707,10 +737,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 ] @@ -762,13 +805,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); @@ -877,8 +916,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)) { 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 diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 35a3286e4a1b..941fbd9ba6f4 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; @@ -252,8 +252,23 @@ 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); + + 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 @@ -347,7 +362,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 +376,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 +430,69 @@ static std::vector get_logits( return ret; } +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; + } + if (ask) { + return true; + } + + 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)); + } 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 +767,32 @@ 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 && 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, {}); + 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; for (device_config & dc : dev_configs) {