diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 72db486cace..64258424c00 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); + // qwen4exp's QSA gather graphs never reference the mask, so it has no buffer; + // the same guard the other attention inputs carry + 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/models/models.h b/src/models/models.h index 9b87a40d5af..49677d29cba 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 * qsa_bias, 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 100a6de4269..832aef18b8a 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -528,7 +528,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; @@ -549,7 +550,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( // 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 && + // the gather path needs the per-cell bias: it carries the full visibility of every + // cell (foreign, future, unpooled, tail), so the gathered cells' bias values can + // serve directly as the attention mask over the gathered set + const bool blk_bias = !gather && 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; @@ -662,7 +666,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); + // the gather path pads the width to a multiple of 256 so the gathered K/V satisfy flash + // attention's padding without extra ops; surplus cells arrive with a -inf bias and are + // masked out of the gathered attention + 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)); @@ -681,8 +690,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k_cur, ggml_tensor * v_cur, ggml_tensor * top_k, + ggml_tensor * qsa_bias, 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) { @@ -712,6 +723,70 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } + // 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(qsa_bias != nullptr && qsa_bias->ne[0] == n_kv && "QSA gather requires the per-cell bias"); + + // 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); + + // the gathered cells' bias values already encode visibility: 0 for a selected block + // member, 1e9 for the always-visible tail, -inf for anything the padded top-k width + // pulled in that the query must not see. clamping to (-inf, 0] turns that into the + // attention mask over the gathered set. + ggml_tensor * b1 = ggml_view_4d(ctx0, qsa_bias, 1, n_kv, 1, ns, + qsa_bias->nb[0], qsa_bias->nb[1], qsa_bias->nb[3], 0); + ggml_tensor * m_g = ggml_get_rows(ctx0, b1, idx); // F32 [1, n_topk, 1, ns] + m_g = ggml_reshape_4d(ctx0, m_g, n_topk, 1, 1, ns); + m_g = ggml_clamp(ctx0, m_g, -INFINITY, 0.0f); + 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; + } + ggml_tensor * kq_mask = inp->get_kq_mask(); // prepare new kq mask - starts filled with -INFINITY @@ -768,7 +843,26 @@ 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 ] @@ -821,7 +915,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); + ggml_tensor * qsa_bias = gather ? qsa_inps.at((uint32_t) hparams.dsv4_compress_ratios[il])->bias : nullptr; + + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, qsa_bias, kq_scale, il, gather); } else { cur = build_attn(inp, nullptr, nullptr, nullptr,