diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 5167c037db3..648facd4b3d 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -50,6 +50,8 @@ class llama_kv_cells { for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) { seq_pos[s].clear(); } + + generation++; } void reset_shift() { @@ -100,6 +102,12 @@ class llama_kv_cells { return has_shift; } + // incremented by every mutator above; used by llama_memory_hybrid_idx to detect + // that the cell contents changed and cached host-side data must be recomputed + uint32_t get_generation() const { + return generation; + } + // move cell isrc to idst (used during defrag) //void mv(uint32_t isrc, uint32_t idst) { // assert(isrc < pos.size()); @@ -164,6 +172,8 @@ class llama_kv_cells { void set(uint32_t i, const llama_kv_cells & other) { assert(i + other.pos.size() <= pos.size()); + generation++; + for (uint32_t j = 0; j < other.pos.size(); ++j) { const auto idx = i + j; @@ -195,6 +205,8 @@ class llama_kv_cells { void set(const std::vector & idxs, const llama_kv_cells & other) { assert(idxs.size() == other.pos.size()); + generation++; + for (uint32_t j = 0; j < other.pos.size(); ++j) { const auto idx = idxs[j]; @@ -227,6 +239,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + generation++; + seq_pos_rm(i); seq[i].reset(); @@ -245,6 +259,8 @@ class llama_kv_cells { assert(pos[i] != -1); assert(seq_id >= 0); + generation++; + seq[i].reset(seq_id); seq_pos_dec(seq_id, pos[i]); @@ -265,6 +281,8 @@ class llama_kv_cells { bool seq_keep(uint32_t i, llama_seq_id seq_id) { assert(i < pos.size()); + generation++; + if (seq[i].test(seq_id)) { seq_pos_rm(i); seq[i].reset(); @@ -338,6 +356,8 @@ class llama_kv_cells { assert(pos[i] != -1); assert(!seq[i].test(seq_id)); + generation++; + seq[i].set(seq_id); seq_pos_inc(seq_id, pos[i]); } @@ -424,6 +444,8 @@ class llama_kv_cells { assert(pos[i] == -1); assert(seq[i].none()); + generation++; + pos[i] = p; used.insert(i); @@ -431,6 +453,9 @@ class llama_kv_cells { void ext_set(uint32_t i, llama_kv_cell_ext p) { assert(i < ext.size()); + + generation++; + ext[i] = p; } @@ -441,6 +466,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + generation++; + seq_pos_rm(i); pos[i] += d; @@ -470,6 +497,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + generation++; + const llama_pos p_old = pos[i]; seq_pos_rm(i); @@ -485,6 +514,9 @@ class llama_kv_cells { private: bool has_shift = false; + // bumped on every mutation of the cell contents (see get_generation()) + uint32_t generation = 0; + // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id) std::set used; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d4e59d77e57..a1742dd8c13 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -366,8 +367,57 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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 + // the graph calls this once per attention layer with identical inputs and the + // result depends only on the indexer cells and the ubatch, so memoize it on the + // memory object (see qsa_memo). only memoize small batches (decode): during + // prefill the bias matrix is n_kv*n_tokens floats and the cost is amortized. + const size_t bytes_cell_blk = sizeof(int32_t)*n_kv*n_ns; + const size_t bytes_blk_cells = sizeof(int32_t)*r*n_blocks*n_ns; + const size_t bytes_blk_pos = sizeof(int32_t)*4*n_blocks*n_ns; + const size_t bytes_bias = sizeof(float)*n_kv*n_tokens; + + auto & memo = mem->memo_qsa; + + const bool use_memo = bytes_bias < (size_t) 256*1024*1024; + + uint64_t key = 0; + if (use_memo) { + key = 1469598103934665603ull; + const auto fnv = [&key](uint64_t v) { + key ^= v; + key *= 1099511628211ull; + }; + + fnv((uint64_t) n_kv); + fnv((uint64_t) n_ns); + fnv((uint64_t) n_blocks); + fnv((uint64_t) n_tokens); + fnv((uint64_t) r); + + for (int64_t s = 0; s < n_ns; ++s) { + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + + fnv((uint64_t) (int64_t) seq_of_stream); + fnv(mem->get_mem_idx()->get_cells(seq_of_stream).get_generation()); + } + + for (int64_t i = 0; i < n_tokens; ++i) { + fnv((uint64_t) (int64_t) ubatch->seq_id[i][0]); + fnv((uint64_t) (int64_t) ubatch->pos[i]); + } + + if (memo.valid && memo.key == key) { + memcpy(dst_cell_blk, memo.cell_blk.data(), bytes_cell_blk); + memcpy(dst_blk_cells, memo.blk_cells.data(), bytes_blk_cells); + memcpy(dst_blk_pos, memo.blk_pos.data(), bytes_blk_pos); + memcpy(dst_bias, memo.bias.data(), bytes_bias); + + return; + } + } + + // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio. All three + // mrope sections carry it: exact for text, approximate for images. Positions repeat per stream. 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) { @@ -462,4 +512,19 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } } + + if (use_memo) { + memo.cell_blk.resize(n_kv*n_ns); + memo.blk_cells.resize(r*n_blocks*n_ns); + memo.blk_pos.resize(4*n_blocks*n_ns); + memo.bias.resize(n_kv*n_tokens); + + memcpy(memo.cell_blk.data(), dst_cell_blk, bytes_cell_blk); + memcpy(memo.blk_cells.data(), dst_blk_cells, bytes_blk_cells); + memcpy(memo.blk_pos.data(), dst_blk_pos, bytes_blk_pos); + memcpy(memo.bias.data(), dst_bias, bytes_bias); + + memo.key = key; + memo.valid = true; + } } diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e3472646d0f..51cee21f0e8 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -85,6 +85,23 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + +public: + // memo for set_input_qsa: its result depends only on the indexer cells content and + // the ubatch, but the graph asks for it once per attention layer with identical + // arguments. lives here (not in the short-lived batch context) so the memo survives + // across decode steps. keyed by a fingerprint of cells generations + ubatch tokens. + struct qsa_memo { + bool valid = false; + uint64_t key = 0; + + std::vector cell_blk; + std::vector blk_cells; + std::vector blk_pos; + std::vector bias; + }; + + mutable qsa_memo memo_qsa; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { diff --git a/src/models/models.h b/src/models/models.h index af60764c2f7..4980d4e254f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2351,7 +2351,19 @@ struct llama_model_qwen4exp : public llama_model_base { // build_rs writes the state tensor in place, so one gather per cache tensor is reused std::map rs_rows; - // one conv history per cache tensor: delta-net and PLE each have their own + // QSA cache-layout inputs depend only on the cells and the ubatch, so all + // layers sharing a compress ratio reuse one tensor set instead of building + // (and filling, and copying to the device) 48 identical copies + struct qsa_inps { + ggml_tensor * k_idxs; + ggml_tensor * cell_blk; + ggml_tensor * blk_cells; + ggml_tensor * blk_pos; + ggml_tensor * bias; + }; + std::map qsa_shared; + + // conv history at an explicit offset: delta-net and PLE share the row ggml_tensor * build_conv_state_at( llm_graph_input_rs * inp, ggml_tensor * conv_states_all, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b5003..851fe49366b 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -489,52 +489,49 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( 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; - - const auto it = qsa_inps.find((uint32_t) r); - if (it != qsa_inps.end()) { - inp = it->second; - } else { - auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + // the cache-layout inputs are identical for every layer with this ratio: + // build and register them once, reuse the tensors afterwards + auto sh = qsa_shared.find((uint32_t) r); + if (sh == qsa_shared.end()) { + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r); 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->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 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); - inp = qsa.get(); + sh = qsa_shared.emplace((uint32_t) r, qsa_inps{ + qsa->k_idxs, qsa->cell_blk, qsa->blk_cells, qsa->blk_pos, qsa->bias, + }).first; + res->add_input(std::move(qsa)); - qsa_inps.emplace((uint32_t) r, inp); } + ggml_tensor * in_k_idxs = sh->second.k_idxs; + ggml_tensor * in_cell_blk = sh->second.cell_blk; + ggml_tensor * in_blk_cells = sh->second.blk_cells; + ggml_tensor * in_blk_pos = sh->second.blk_pos; + ggml_tensor * in_bias = sh->second.bias; + // 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)); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, in_k_idxs, il)); // 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); + ggml_tensor * members = ggml_get_rows(ctx0, k_all, in_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 @@ -551,7 +548,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // 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, + pooled = ggml_rope_multi(ctx0, pooled, in_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); @@ -583,16 +580,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // 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); + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), in_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); - } + expanded = ggml_add(ctx0, expanded, in_bias); cb(expanded, "indexer_score_tokens", il); // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail