diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 72db486cace..1f8adc04c58 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1087,7 +1087,11 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); - mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + // the mask can be left unallocated when the graph never attends over the full + // cache (e.g. the qwen4exp gather path reads its own compact per-cell mask) + if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) { + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + } if (inp_attn->self_k_rot) { mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot); diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index d4e59d77e57..98753e9ad9b 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -343,6 +343,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + ggml_tensor * mask_row, ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, @@ -350,7 +351,10 @@ void llama_memory_hybrid_idx_context::set_input_qsa( GGML_ASSERT(ratio > 0); GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); - GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); + // the block-gather decode path never references cell_blk in the graph, so the + // scheduler leaves it unallocated; skip its (O(n_kv) per ubatch) fill in that case + const bool has_cell_blk = cell_blk->buffer != nullptr; + GGML_ASSERT(!has_cell_blk || 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 @@ -361,7 +365,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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_cell_blk = has_cell_blk ? (int32_t *) cell_blk->data : nullptr; 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; @@ -385,7 +389,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( 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_cell_blk = has_cell_blk ? dst_cell_blk + s*n_kv : nullptr; 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 @@ -420,7 +424,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // 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) { + for (int64_t j = 0; has_cell_blk && j < n_kv; ++j) { if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) { blk_of[j] = -1; } @@ -432,6 +436,33 @@ void llama_memory_hybrid_idx_context::set_input_qsa( const llama_seq_id seq_id = ubatch->seq_id[i][0]; const llama_pos q = ubatch->pos[i]; + // per-cell visibility for the gather path: what row i of the attention + // kq_mask would say, minus SWA/alibi (the gather gate excludes those). + // unallocated (and skipped) when the graph takes the masked path. + if (mask_row != nullptr && mask_row->buffer != nullptr) { + float * dst_row = (float *) mask_row->data + i*n_kv; + + const bool is_2d = ubatch->is_pos_2d(); + + const llama_pos qx = is_2d ? ubatch->pos[i + ubatch->n_tokens*2] : 0; + const llama_pos qy = is_2d ? ubatch->pos[i + ubatch->n_tokens] : 0; + + 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) { + v = 0.0f; + + // M-RoPE: image tokens can share the temporal position of the query + if (is_2d && cells.pos_get(j) == q && cells.ext_get(j).is_2d_gt(qx, qy)) { + v = -INFINITY; + } + } + + dst_row[j] = v; + } + } + // the tail is an incomplete block and is always visible, as in the reference const llama_pos tail_start = (q + 1)/r*r; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e3472646d0f..d5c88be823d 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -137,7 +137,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // 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, + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, ggml_tensor * mask_row, ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af..f69681d5d1d 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2317,8 +2317,10 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * mask_row, float kq_scale, - int il); + int il, + bool gather = false); // 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 @@ -2331,7 +2333,8 @@ struct llama_model_qwen4exp : public llama_model_base { ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il); + int il, + bool gather = false); ggml_tensor * build_layer_attn_linear( llm_graph_input_rs * inp, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index abf6a0502fb..e6a702c8a4c 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -430,7 +430,7 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { 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(cell_blk, blk_cells, blk_pos, mask_row, bias, ubatch, ratio, blk_bias); } bool can_reuse(const llm_graph_params & params) override { @@ -457,6 +457,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv); res &= bias->ne[1] == params.ubatch.n_tokens/n_stream; + res &= mask_row->ne[0] == n_kv; + res &= mask_row->ne[1] == params.ubatch.n_tokens/n_stream; + res &= mask_row->ne[2] == n_stream; + return res; } @@ -467,6 +471,10 @@ class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { 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] + // per-cell visibility for the gather path (0 or -INF); left unallocated by the + // masked path, which reads the attention kq_mask instead. F32 [n_kv, n_tps, n_stream] + ggml_tensor * mask_row = nullptr; + const llama_memory_hybrid_idx_context * mctx; const uint32_t ratio; @@ -480,7 +488,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( ggml_tensor * inp_pos, ggml_tensor * kq_mask, int * sections, - int il) { + int il, + bool gather) { const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); const int64_t idx_dim = hparams.indexer_head_size; @@ -519,7 +528,9 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( 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->mask_row = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream); + ggml_set_input(qsa->mask_row); ggml_set_input(qsa->cell_blk); ggml_set_input(qsa->blk_cells); ggml_set_input(qsa->blk_pos); @@ -589,6 +600,31 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( score = ggml_add(ctx0, score, inp->bias); } + // block-level top-k for the gather path: the scores are one-per-block anyway, so pick + // the winning blocks first and expand only those to cell indices via blk_cells. this + // sorts n_blocks instead of n_kv entries and skips every O(n_kv) op below. the cell + // count r*K_blk stays a multiple of 256, which keeps flash attention's padding happy. + // needs blk_bias: the per-block bias is what carries visibility into the selection. + if (gather && blk_bias && 256 % r == 0) { + const int64_t K_blk = std::min(n_blocks, + GGML_PAD(((int64_t) hparams.indexer_top_k + r - 1)/r, 256/r)); + + ggml_tensor * blk_idx = ggml_cont(ctx0, ggml_top_k(ctx0, score, K_blk)); // I32 [K_blk, n_tps, n_stream] + blk_idx = ggml_reshape_4d(ctx0, blk_idx, K_blk, 1, n_stream, 1); + + // blk_cells maps (block, slot-in-block) -> cell: view one row of r cells per block + ggml_tensor * bc = ggml_view_4d(ctx0, inp->blk_cells, r, n_blocks, 1, n_stream, + r*ggml_element_size(inp->blk_cells), + r*n_blocks*ggml_element_size(inp->blk_cells), + inp->blk_cells->nb[1], 0); + + ggml_tensor * cells = ggml_get_rows(ctx0, bc, blk_idx); // I32 [r, K_blk, 1, n_stream] + cells = ggml_reshape_4d(ctx0, cells, r*K_blk, n_tps, 1, n_stream); + cb(cells, "indexer_top_k_cells", il); + + return cells; + } + // 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); @@ -604,7 +640,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( 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); + // for the gather path the width is padded to a multiple of GGML_KQ_MASK_PAD so the + // gathered K/V/mask satisfy flash attention's padding requirements without extra ops; + // the gathered mask carries -INF for any surplus cells the padded top-k pulls in + const int64_t width = gather + ? std::min(n_kv, GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256)) + : 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)); @@ -623,8 +664,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * mask_row, float kq_scale, - int il) { + int il, + bool gather) { // 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) { @@ -656,6 +699,72 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * kq_mask = inp->get_kq_mask(); + // decode-time sparse gather: instead of masking the full cache (which makes flash + // attention scan all n_kv cells to use only indexer_top_k of them), copy the selected + // cells out of the cache and attend densely over just those. this bounds the per-token + // attention cost by the top-k width instead of the context depth. + // restricted to single-token-per-stream ubatches (decode), where every stream carries + // exactly one top-k list, so the gathered K/V stay one tensor per stream. + if (gather) { + ggml_tensor * kf = mctx_cur->get_k(ctx0, il); // [hd_k, n_head_kv, n_kv, ns] + ggml_tensor * vf = mctx_cur->get_v(ctx0, il); // [hd_v, n_head_kv, n_kv, ns] + + const int64_t hd_k = kf->ne[0]; + const int64_t hd_v = vf->ne[0]; + const int64_t n_h_kv = kf->ne[1]; + const int64_t n_kv = kf->ne[2]; + const int64_t ns = kf->ne[3]; + const int64_t n_topk = top_k->ne[0]; + + GGML_ASSERT(top_k->ne[1] == 1 && "QSA gather requires single-token-per-stream ubatches"); + GGML_ASSERT(kq_mask->ne[0] == n_kv); + + // the heads of a cell are contiguous in the cache, so a cell can be gathered as one row + GGML_ASSERT(kf->nb[2] == ggml_row_size(kf->type, hd_k*n_h_kv)); + GGML_ASSERT(vf->nb[2] == ggml_row_size(vf->type, hd_v*n_h_kv)); + + ggml_tensor * k_cells = ggml_view_4d(ctx0, kf, hd_k*n_h_kv, n_kv, 1, ns, + kf->nb[2], kf->nb[2]*n_kv, kf->nb[3], 0); + ggml_tensor * v_cells = ggml_view_4d(ctx0, vf, hd_v*n_h_kv, n_kv, 1, ns, + vf->nb[2], vf->nb[2]*n_kv, vf->nb[3], 0); + + // top_k [n_topk, 1, 1, ns] -> the index layout ggml_get_rows expects: [n_topk, 1, ns, 1] + ggml_tensor * idx = ggml_reshape_4d(ctx0, top_k, n_topk, 1, ns, 1); + + // get_rows dequantizes the cells to F32; build_attn_mha casts to F16 for flash attention + ggml_tensor * k_g = ggml_get_rows(ctx0, k_cells, idx); // F32 [hd_k*n_h_kv, n_topk, 1, ns] + ggml_tensor * v_g = ggml_get_rows(ctx0, v_cells, idx); // F32 [hd_v*n_h_kv, n_topk, 1, ns] + + k_g = ggml_reshape_4d(ctx0, k_g, hd_k, n_h_kv, n_topk, ns); + v_g = ggml_reshape_4d(ctx0, v_g, hd_v, n_h_kv, n_topk, ns); + cb(k_g, "qsa_k_gathered", il); + cb(v_g, "qsa_v_gathered", il); + + // gather the same cells' visibility: keeps -INF for any invalid cell the padded + // top-k width pulled in (e.g. when fewer than n_topk cells are visible). + // reading the compact per-cell mask_row instead of the attention kq_mask leaves + // the latter unreferenced, so it is neither filled (O(n_kv * pad) host work) nor + // uploaded (n_kv * GGML_KQ_MASK_PAD * 2 bytes, ~18 MB/token at 141K ctx) + GGML_ASSERT(mask_row != nullptr); + GGML_ASSERT(mask_row->ne[0] == n_kv); + ggml_tensor * m1 = ggml_view_4d(ctx0, mask_row, 1, n_kv, 1, ns, + mask_row->nb[0], mask_row->nb[1], mask_row->nb[3], 0); + ggml_tensor * m_g = ggml_get_rows(ctx0, m1, idx); // F32 [1, n_topk, 1, ns] + m_g = ggml_reshape_4d(ctx0, m_g, n_topk, 1, 1, ns); + m_g = ggml_cast(ctx0, m_g, GGML_TYPE_F16); // FA wants contiguous F16 + cb(m_g, "qsa_mask_gathered", il); + + ggml_tensor * cur = build_attn_mha(q_cur, k_g, v_g, nullptr, m_g, 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); + } + + return cur; + } + // prepare new kq mask - starts filled with -INFINITY ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); @@ -710,7 +819,27 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( // 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; - ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + // gather-based QSA decode: worth it once the cache is meaningfully deeper than the + // top-k width; below that the masked path costs about the same. QWEN4EXP_QSA_GATHER=0 + // disables it (A/B lever, and an escape hatch). + static const bool gather_enabled = [] { + const char * e = getenv("QWEN4EXP_QSA_GATHER"); + return e == nullptr || atoi(e) != 0; + }(); + + bool gather = false; + if (qsa && gather_enabled) { + const int64_t r = hparams.dsv4_compress_ratios[il]; + const int64_t n_kv = mctx_hyb->get_idx()->get_n_kv(); + const int64_t width = GGML_PAD((int64_t) hparams.indexer_top_k + r - 1, 256); + + const int64_t n_stream = mctx_hyb->get_n_stream(); + + gather = n_tokens == n_stream && n_kv >= 2*width; + + } + + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il, gather) : nullptr; // 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 ] @@ -763,7 +892,13 @@ 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); + ggml_tensor * qsa_mask_row = nullptr; + if (gather) { + const auto it = qsa_inps.find((uint32_t) hparams.dsv4_compress_ratios[il]); + GGML_ASSERT(it != qsa_inps.end()); + qsa_mask_row = it->second->mask_row; + } + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, qsa_mask_row, kq_scale, il, gather); } else { cur = build_attn(inp, nullptr, nullptr, nullptr,