diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 8fafcd15304e..33a4fc24c6db 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1868,16 +1868,33 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st for (uint32_t s = 0; s < n_stream; ++s) { // p_max inclusive: an embd token looks up cells at its own (shared) position - v_cells[s].for_each_token_in(seqs, 0, p_max + 1, + v_cells[s].for_each_token_in(seqs, w0, p_max + 1, [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) { - if (pos >= w0) { - hist[key(seq_id, pos)] = tok; - } else if (pos > below[seq_id].first) { - below[seq_id] = { pos, tok }; - } + hist[key(seq_id, pos)] = tok; }); } + // below[] only answers when an M-RoPE gap leaves the window empty, so it costs a pass of + // its own that contiguous positions never pay + bool below_ready = false; + + const auto ensure_below = [&]() { + if (below_ready) { + return; + } + + below_ready = true; + + for (uint32_t s = 0; s < n_stream; ++s) { + v_cells[s].for_each_token_in(seqs, 0, w0, + [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) { + if (pos > below[seq_id].first) { + below[seq_id] = { pos, tok }; + } + }); + } + }; + // the token at pos p, or the nearest earlier one when p falls in an M-RoPE gap const auto lookup = [&](llama_seq_id seq_id, llama_pos p) -> llama_token { for (llama_pos q = p; q >= w0; --q) { @@ -1886,6 +1903,9 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st return it->second; } } + + ensure_below(); + return below[seq_id].second; }; diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 5167c037db3e..2914d35f5100 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -10,6 +10,10 @@ #include #include +#if defined(_MSC_VER) +#include +#endif + struct llama_kv_cell_ext { // 2D spatial positions, typically used for M-RoPE llama_pos x = 0; @@ -31,6 +35,105 @@ struct llama_kv_cell_ext { } }; +// index of the lowest set bit, and of the highest; the word is never zero here +#if defined(_MSC_VER) +static inline uint32_t llama_kv_ctz64(uint64_t x) { + unsigned long r; + _BitScanForward64(&r, x); + return (uint32_t) r; +} + +static inline uint32_t llama_kv_clz64(uint64_t x) { + unsigned long r; + _BitScanReverse64(&r, x); + return (uint32_t) r; +} +#else +static inline uint32_t llama_kv_ctz64(uint64_t x) { + return (uint32_t) __builtin_ctzll(x); +} + +static inline uint32_t llama_kv_clz64(uint64_t x) { + return (uint32_t) (63 - __builtin_clzll(x)); +} +#endif + +// the cells tracked are a large fraction of the cache, so a bitmap is both smaller and faster +// to walk than a tree node per index +class llama_kv_idx_set { +public: + void resize(uint32_t n) { + bits.assign((n + 63)/64, 0); + n_set = 0; + } + + void clear() { + std::fill(bits.begin(), bits.end(), 0); + n_set = 0; + } + + void insert(uint32_t i) { + uint64_t & w = bits[i/64]; + const uint64_t b = 1ull << (i%64); + + n_set += (w & b) == 0; + w |= b; + } + + void erase(uint32_t i) { + uint64_t & w = bits[i/64]; + const uint64_t b = 1ull << (i%64); + + n_set -= (w & b) != 0; + w &= ~b; + } + + bool contains(uint32_t i) const { + return (bits[i/64] >> (i%64)) & 1; + } + + uint32_t size() const { return n_set; } + bool empty() const { return n_set == 0; } + + uint32_t first() const { + for (size_t w = 0; w < bits.size(); ++w) { + if (bits[w]) { + return 64*w + llama_kv_ctz64(bits[w]); + } + } + + return 0; + } + + uint32_t last() const { + for (size_t w = bits.size(); w-- > 0; ) { + if (bits[w]) { + return 64*w + llama_kv_clz64(bits[w]); + } + } + + return 0; + } + + // ascending order, as the set it replaces guaranteed + template + void for_each(F && f) const { + for (size_t w = 0; w < bits.size(); ++w) { + uint64_t m = bits[w]; + + while (m) { + f((uint32_t) (64*w + llama_kv_ctz64(m))); + m &= m - 1; + } + } + } + +private: + std::vector bits; + + uint32_t n_set = 0; +}; + // meta information about KV cells that can be part of multiple sequences at the same time // TODO: add unit tests class llama_kv_cells { @@ -69,6 +172,7 @@ class llama_kv_cells { ext.resize(n); shift.resize(n); seq.resize(n); + used.resize(n); reset(); } @@ -87,13 +191,13 @@ class llama_kv_cells { // the index of the first cell that is used // return 0 if no cells are used uint32_t used_min() const { - return used.empty() ? 0 : *used.begin(); + return used.empty() ? 0 : used.first(); } // the index of the last cell that is used + 1 // return 0 if no cells are used uint32_t used_max_p1() const { - return used.empty() ? 0 : *used.rbegin() + 1; + return used.empty() ? 0 : used.last() + 1; } bool get_has_shift() const { @@ -314,22 +418,24 @@ class llama_kv_cells { // note: used by n-gram input embeddings to recover the tokens preceding a ubatch template void for_each_token_in(const std::bitset & seqs, llama_pos p0, llama_pos p1, F && f) const { - for (const auto & i : used) { + used.for_each([&](uint32_t i) { if (pos[i] < p0 || pos[i] >= p1) { - continue; + return; } const auto m = seq[i] & seqs; - if (m.none()) { - continue; - } - for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + // a cell carries a handful of sequences at most, so stop once they are all seen + // instead of walking the whole LLAMA_MAX_SEQ width + size_t left = m.count(); + + for (llama_seq_id s = 0; left > 0 && s < (llama_seq_id) LLAMA_MAX_SEQ; ++s) { if (m.test(s)) { f(s, pos[i], ext[i].tok); + --left; } } - } + }); } // note: call only if the cell is not empty and the seq_id is not in the cell @@ -486,7 +592,7 @@ class llama_kv_cells { bool has_shift = false; // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id) - std::set used; + llama_kv_idx_set used; std::vector pos; diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af9..7180051bf600 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2320,6 +2320,22 @@ struct llama_model_qwen4exp : public llama_model_base { float kq_scale, int il); + // attention over the whole window, masked down to the selection + ggml_tensor * build_qsa_scan( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + + // attention over the selected cells alone, gathered into a window per query + ggml_tensor * build_qsa_gather( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_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; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index abf6a0502fbf..f6dcdcce0703 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -579,9 +579,17 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( 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); + + // the heads sit side by side on ne[1] and there are few of them, so summing slices beats a + // transpose that would carry the whole block by token surface twice over + ggml_tensor * summed = nullptr; + for (int64_t h = 0; h < n_idx_h; ++h) { + ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream, + score->nb[2], score->nb[3], h*score->nb[1]); + summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice); + } + + score = summed; cb(score, "indexer_score", il); // one value per block, so it is cheaper to bias here than after the cells are expanded @@ -654,6 +662,35 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } + const int64_t width = top_k->ne[0]; + const int64_t n_tps = top_k->ne[1]; + const int64_t n_kv = mctx_cur->get_n_kv(); + + // the two costs meet at 2*n_tps*width == n_kv; the margin keeps the windows small enough + // for the compute buffer of a decode graph. flash attention is what keeps the values as rows + ggml_tensor * cur = cparams.flash_attn && 4*n_tps*width < n_kv + ? build_qsa_gather(inp, q_cur, top_k, kq_scale, il) + : build_qsa_scan (inp, q_cur, top_k, 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); + } + + return cur; +} + +// The window stays whole and the mask hides every cell the selection leaves out. +// The mask build below copies the MLA sparse path in llm_graph_context::build_attn. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_scan( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + const auto * mctx_cur = inp->mctx; + ggml_tensor * kq_mask = inp->get_kq_mask(); // prepare new kq mask - starts filled with -INFINITY @@ -686,15 +723,64 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( 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); + return build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, 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); - } +// Key and value traffic follows the budget instead of the whole cache. The queries ride the +// stream axis of the attention, each carrying the window its own selection named. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_gather( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + const auto * mctx_cur = inp->mctx; - return cur; + const int64_t width = top_k->ne[0]; + const int64_t n_tps = top_k->ne[1]; + const int64_t n_stream = top_k->ne[3]; + const int64_t n_q = n_tps*n_stream; + + ggml_tensor * k_all = mctx_cur->get_k(ctx0, il); + ggml_tensor * v_all = mctx_cur->get_v(ctx0, il); + + const int64_t n_kv = k_all->ne[2]; + + // a cell holds its heads back to back, so one row of the gather is one whole cell + ggml_tensor * k_cells = ggml_view_3d(ctx0, k_all, k_all->ne[0]*k_all->ne[1], n_kv, n_stream, + k_all->nb[2], k_all->nb[3], 0); + ggml_tensor * v_cells = ggml_view_3d(ctx0, v_all, v_all->ne[0]*v_all->ne[1], n_kv, n_stream, + v_all->nb[2], v_all->nb[3], 0); + + // a cell index names a cell of its own stream, and the ubatch lays the queries of a stream + // out contiguously, so the flat index of a window is the index of its query + ggml_tensor * idx_stream = ggml_reshape_2d(ctx0, top_k, width*n_tps, n_stream); + + ggml_tensor * k_sel = ggml_get_rows(ctx0, k_cells, idx_stream); + ggml_tensor * v_sel = ggml_get_rows(ctx0, v_cells, idx_stream); + + k_sel = ggml_reshape_4d(ctx0, k_sel, k_all->ne[0], k_all->ne[1], width, n_q); + v_sel = ggml_reshape_4d(ctx0, v_sel, v_all->ne[0], v_all->ne[1], width, n_q); + cb(k_sel, "qsa_k_sel", il); + cb(v_sel, "qsa_v_sel", il); + + // gathering the attention mask at the selected cells leaves the same values the scan path + // would put there, so the window carries the reach of its query + ggml_tensor * kq_mask = inp->get_kq_mask(); + + GGML_ASSERT(kq_mask->nb[3] == kq_mask->nb[1]*n_tps); + + ggml_tensor * mask_cells = ggml_view_3d(ctx0, kq_mask, 1, n_kv, n_q, + kq_mask->nb[0], kq_mask->nb[1], 0); + + ggml_tensor * idx_query = ggml_reshape_3d(ctx0, top_k, width, n_q, 1); + + ggml_tensor * mask = ggml_get_rows(ctx0, mask_cells, idx_query); + + mask = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mask, width, 1, 1, n_q), GGML_TYPE_F16); + cb(mask, "qsa_mask_sel", il); + + return build_attn_mha(q_cur, k_sel, v_sel, nullptr, mask, nullptr, nullptr, kq_scale, il); } ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(