diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6fd69a82..18de9c16 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -94,6 +94,7 @@ add_library(llama graphs/build_qwen3.cpp graphs/build_mellum.cpp graphs/build_qwen3next.cpp + graphs/build_qwen4exp.cpp graphs/build_qwen35.cpp graphs/build_phi2.cpp graphs/build_phi3.cpp diff --git a/src/graphs/build_qwen4exp.cpp b/src/graphs/build_qwen4exp.cpp new file mode 100644 index 00000000..134723bf --- /dev/null +++ b/src/graphs/build_qwen4exp.cpp @@ -0,0 +1,539 @@ +#include "../llama-build-context.h" +#include "../llama-model.h" +#include "../llama-context.h" +#include "../llama-delta-net.h" + +// the [hc_dim] gamma is wider than the per-stream reduction, so ggml_fused_rms_norm cannot +// express this and the two ops stay separate +static ggml_tensor * qwen4exp_grouped_rms( + ggml_context * ctx0, + const llama_hparams & hparams, + ggml_tensor * x, + ggml_tensor * w, + int32_t hc_dim, + int32_t nt) { + ggml_tensor * t = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps); + t = ggml_reshape_2d(ctx0, t, hc_dim, nt); + return ggml_mul(ctx0, t, w); +} + +// the low-rank down/up pair is a gate over the normalised streams, not a mixing matrix +static ggml_tensor * qwen4exp_hc_mix( + llm_build_context & bctx, + ggml_context * ctx0, + llama_context & lctx, + const llama_hparams & hparams, + ggml_tensor * x, + ggml_tensor * w_norm, + ggml_tensor * w_down, + ggml_tensor * w_up, + ggml_tensor * w_inject, + ggml_tensor ** inject, + int32_t n_embd, + int il, + const llm_build_cb & cb) { + const int32_t hc = hparams.dsv4_hc_mult; + const int32_t hc_dim = hc * n_embd; + const int64_t nt = x->ne[2]; + + ggml_tensor * xn = qwen4exp_grouped_rms(ctx0, hparams, x, w_norm, hc_dim, nt); + cb(xn, "hc_norm", il); + + ggml_tensor * lo = llm_build_context::llm_build_lora_mm(lctx, ctx0, w_down, xn); + lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc)); + ggml_tensor * gate = ggml_sigmoid(ctx0, llm_build_context::llm_build_lora_mm(lctx, ctx0, w_up, lo)); + cb(gate, "hc_gate", il); + + ggml_tensor * gated = ggml_mul(ctx0, xn, gate); + gated = ggml_reshape_3d(ctx0, gated, n_embd, hc, nt); + + ggml_tensor * mixed = ggml_multi_add(ctx0, + ggml_view_2d(ctx0, gated, n_embd, nt, gated->nb[2], 0), hc); + mixed = ggml_scale(ctx0, mixed, 1.0f / (float) hc); + cb(mixed, "hc_mixed", il); + + if (inject) { + *inject = llm_build_context::llm_build_lora_mm(lctx, ctx0, w_inject, xn); + cb(*inject, "hc_inject", il); + } + + GGML_UNUSED(bctx); + return mixed; +} + +// the factor of two centres the weights on one, so an untrained injection is a plain +// residual add +static ggml_tensor * qwen4exp_hc_combine( + ggml_context * ctx0, + const llama_hparams & hparams, + ggml_tensor * residual, + ggml_tensor * block_out, + ggml_tensor * inject, + int32_t n_embd, + int il, + const llm_build_cb & cb) { + const int32_t hc = hparams.dsv4_hc_mult; + const int64_t nt = residual->ne[2]; + + ggml_tensor * w = ggml_sigmoid(ctx0, ggml_scale(ctx0, inject, 1.0f / (float) hc)); + w = ggml_scale(ctx0, w, 2.0f); + w = ggml_reshape_3d(ctx0, w, 1, hc, nt); + + ggml_tensor * b = ggml_reshape_3d(ctx0, block_out, n_embd, 1, nt); + b = ggml_repeat_4d(ctx0, b, n_embd, hc, nt, 1); + + ggml_tensor * cur = ggml_add(ctx0, residual, ggml_mul(ctx0, b, w)); + cb(cur, "hc_combine", il); + + return cur; +} + +// out[c, t] = sum_k w[k, c] * x[c, t - (K-1-k)*dilation] +// +// the (K-1)*dilation positions reached back live in the tail of this layer's state row. +// Prepending them puts every tap inside the tensor, so none needs a pad or a range test +static ggml_tensor * qwen4exp_ple_conv( + ggml_context * ctx0, + ggml_cgraph * gf, + const llama_hparams & hparams, + const llama_model & model, + ggml_tensor * state_all, + ggml_tensor * xt, // [n_tokens, hc_dim] + int32_t hc_dim, + int32_t n_tokens, + uint32_t slot, + bool reset, + int il, + const llm_build_cb & cb) { + const int32_t kern = hparams.ple_conv_kernel; + const int32_t dil = hparams.ple_ngram_size; + const int32_t hist = hparams.ple_conv_state(); + + // the delta-net state occupies the front of the row; this history follows it + const size_t esz = ggml_element_size(state_all); + const size_t row_off = esz * (state_all->ne[0] - hist*hc_dim); + + ggml_tensor * state = ggml_cont(ctx0, + ggml_view_2d(ctx0, state_all, hist, hc_dim, hist*esz, slot*state_all->nb[1] + row_off)); + if (reset) { + state = ggml_scale(ctx0, state, 0.0f); + } + cb(state, "ple_conv_state", il); + + ggml_tensor * conv_in = ggml_concat(ctx0, state, xt, 0); // [hist + n_tokens, hc_dim] + + ggml_tensor * conv_out = nullptr; + for (int32_t k = 0; k < kern; ++k) { + const int32_t start = hist - (kern - 1 - k)*dil; + + ggml_tensor * shifted = ggml_cont(ctx0, + ggml_view_2d(ctx0, conv_in, n_tokens, hc_dim, conv_in->nb[1], start*conv_in->nb[0])); + + ggml_tensor * wk = ggml_cont(ctx0, + ggml_view_2d(ctx0, model.layers[il].ple_conv1d, 1, hc_dim, + model.layers[il].ple_conv1d->nb[1], + k * model.layers[il].ple_conv1d->nb[0])); + wk = ggml_reshape_1d(ctx0, wk, hc_dim); + if (wk->type != GGML_TYPE_F32) { + wk = ggml_cast(ctx0, wk, GGML_TYPE_F32); + } + + ggml_tensor * term = ggml_mul(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, shifted)), wk); + term = ggml_cont(ctx0, ggml_transpose(ctx0, term)); // [n_tokens, hc_dim] + + conv_out = conv_out ? ggml_add(ctx0, conv_out, term) : term; + } + + // the last `hist` columns are what the next ubatch reaches back into. When the ubatch is + // shorter than the window they still carry part of the incoming state, which is correct. + ggml_tensor * tail = ggml_cont(ctx0, + ggml_view_2d(ctx0, conv_in, hist, hc_dim, conv_in->nb[1], + (conv_in->ne[0] - hist) * conv_in->nb[0])); + ggml_tensor * dst = ggml_view_2d(ctx0, state_all, hist, hc_dim, hist*esz, + slot*state_all->nb[1] + row_off); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, tail, dst)); + + return conv_out; +} + +static ggml_tensor * qwen4exp_ple( + llm_build_context & bctx, + ggml_context * ctx0, + ggml_cgraph * gf, + llama_context & lctx, + const llama_model & model, + const llama_hparams & hparams, + const delta_net & delta, + ggml_tensor * hidden, + ggml_tensor * rows, + ggml_tensor * state_all, + bool reset_state, + const std::vector & reset_pos, + int32_t n_embd, + int32_t n_tokens, + int il, + const llm_build_cb & cb) { + const int32_t hc = hparams.dsv4_hc_mult; + const int32_t hc_dim = hc * n_embd; + const int32_t n_heads = hparams.ple_n_heads; + + // get_rows lays the head dimension out slowest, which is the flatten order the + // projections expect + ggml_tensor * emb = ggml_get_rows(ctx0, model.tok_embd_per_layer, rows); + emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens); + cb(emb, "ple_embd", il); + + ggml_tensor * key = llm_build_context::llm_build_lora_mm(lctx, ctx0, model.layers[il].ple_key, emb); + ggml_tensor * value = llm_build_context::llm_build_lora_mm(lctx, ctx0, model.layers[il].ple_value, emb); + + auto grouped_norm = [&](ggml_tensor * x, ggml_tensor * w) { + ggml_tensor * t = qwen4exp_grouped_rms(ctx0, hparams, + ggml_reshape_3d(ctx0, x, n_embd, hc, n_tokens), w, hc_dim, n_tokens); + return ggml_reshape_3d(ctx0, t, n_embd, hc, n_tokens); + }; + + key = grouped_norm(key, model.layers[il].ple_norm_key); + ggml_tensor * query = grouped_norm(hidden, model.layers[il].ple_norm_query); + + ggml_tensor * s = ggml_sum_rows(ctx0, ggml_mul(ctx0, key, query)); + s = ggml_scale(ctx0, s, 1.0f / sqrtf((float) n_embd)); + + ggml_tensor * mag = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, s), 1e-6f, INFINITY)); + ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, ggml_sgn(ctx0, s), mag)); + cb(gate, "ple_gate", il); + + ggml_tensor * v3 = ggml_reshape_3d(ctx0, value, n_embd, 1, n_tokens); + v3 = ggml_repeat_4d(ctx0, v3, n_embd, hc, n_tokens, 1); + + ggml_tensor * gated = ggml_mul(ctx0, v3, gate); + cb(gated, "ple_gated_value", il); + + ggml_tensor * normalized = grouped_norm( + ggml_reshape_2d(ctx0, gated, hc_dim, n_tokens), + model.layers[il].ple_norm_conv); + normalized = ggml_reshape_2d(ctx0, normalized, hc_dim, n_tokens); + + ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, normalized)); // [n_tokens, hc_dim] + + ggml_tensor * conv_out = nullptr; + if (delta.batch_shares_one_seq()) { + conv_out = qwen4exp_ple_conv(ctx0, gf, hparams, model, state_all, xt, hc_dim, n_tokens, + delta.state_slot(0), reset_state, il, cb); + } else { + // A mixed-sequence ubatch reads a different history per token, exactly as the + // delta-net path splits it. + for (int32_t i = 0; i < n_tokens; ++i) { + ggml_tensor * x_i = ggml_cont(ctx0, + ggml_view_2d(ctx0, xt, 1, hc_dim, xt->nb[1], i*xt->nb[0])); + ggml_tensor * out_i = qwen4exp_ple_conv(ctx0, gf, hparams, model, state_all, x_i, hc_dim, 1, + delta.state_slot(i), reset_pos[i], il, cb); + conv_out = conv_out ? ggml_concat(ctx0, conv_out, out_i, 0) : out_i; + } + } + + conv_out = ggml_silu(ctx0, conv_out); + conv_out = ggml_cont(ctx0, ggml_transpose(ctx0, conv_out)); // [hc_dim, n_tokens] + conv_out = ggml_reshape_3d(ctx0, conv_out, n_embd, hc, n_tokens); + cb(conv_out, "ple_conv_out", il); + + GGML_UNUSED(bctx); + return ggml_add(ctx0, hidden, ggml_add(ctx0, gated, conv_out)); +} + +// a query keeps a budget of whole blocks plus the incomplete tail it sits in, where a block is +// compress_ratio consecutive cells scored through the mean of its members' indexer keys +static ggml_tensor * qwen4exp_qsa_mask( + llm_build_context & bctx, + ggml_context * ctx0, + llama_context & lctx, + ggml_cgraph * gf, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * KQ_mask, + int il, + const llm_build_cb & cb) { + const llama_hparams & hparams = bctx.hparams; + const llama_model & model = bctx.model; + const llama_kv_cache & kv_self = bctx.kv_self; + + ggml_tensor * kr_cache = il < (int) kv_self.kr_l.size() ? kv_self.kr_l[il] : nullptr; + ggml_tensor * kp_cache = il < (int) kv_self.kp_l.size() ? kv_self.kp_l[il] : nullptr; + if (!kr_cache || !kp_cache || !model.layers[il].indexer_k_proj) { + return KQ_mask; + } + + const int32_t idx_dim = hparams.indexer_head_size; + const int32_t n_idx_h = hparams.indexer_n_head; + const int32_t r = hparams.dsv4_compress_ratios[il]; + const int32_t n_kv = bctx.n_kv; + const int32_t n_tokens = bctx.n_tokens; + + // the cached indexer keys are raw: pooling precedes both the norm and the rotation + ggml_tensor * k_raw = llm_build_context::llm_build_lora_mm(lctx, ctx0, model.layers[il].indexer_k_proj, cur); + k_raw = ggml_reshape_2d(ctx0, k_raw, idx_dim, n_tokens); + cb(k_raw, "qsa_k_raw", il); + { + ggml_tensor * kr_view = ggml_view_2d(ctx0, kr_cache, idx_dim, n_tokens, + ggml_row_size(kr_cache->type, idx_dim), + ggml_row_size(kr_cache->type, idx_dim) * bctx.kv_head); + ggml_tensor * kr_cpy = ggml_cpy(ctx0, k_raw, kr_view); + // the view above bakes kv_head, so register the copy for the offset fixup that + // update_cache_copies() already applies to the K and V writes of a reused graph + if (il < (int) lctx.dsa_cache_copies.size()) { + lctx.dsa_cache_copies[il].cpy = kr_cpy; + lctx.dsa_cache_copies[il].step = kr_cache->nb[1]; + } + ggml_build_forward_expand(gf, kr_cpy); + } + + const int32_t n_blocks = (n_kv + r - 1)/r; + + // a raw key never changes once written, so only the blocks this ubatch wrote need pooling + // again: n_tokens consecutive cells span n_tokens/r, plus one when the run straddles + const int32_t n_win = lctx.qsa_pooled_stale ? n_blocks : std::min(n_blocks, (n_tokens + r - 1)/r + 1); + + llama_context::qsa_input * inp = nullptr; + for (auto & q : lctx.inp_qsa) { + if (q.ratio == (int32_t) r) { + inp = &q; + break; + } + } + if (!inp) { + lctx.inp_qsa.emplace_back(); + inp = &lctx.inp_qsa.back(); + inp->ratio = (int32_t) r; + inp->cell_blk = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_kv); + inp->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); + inp->win_blocks = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_win); + inp->win_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r*n_win); + inp->win_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, GGML_MROPE_SECTIONS*n_win); + inp->head_w = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_idx_h, n_tokens); + cb(inp->cell_blk, "qsa_cell_blk", -1); + cb(inp->bias, "qsa_bias", -1); + cb(inp->win_blocks, "qsa_win_blocks", -1); + cb(inp->win_cells, "qsa_win_cells", -1); + cb(inp->win_pos, "qsa_win_pos", -1); + cb(inp->head_w, "qsa_head_w", -1); + for (ggml_tensor * t : {inp->cell_blk, inp->bias, inp->win_blocks, inp->win_cells, inp->win_pos, inp->head_w}) { + ggml_set_input(t); + ggml_build_forward_expand(gf, t); + } + } + + ggml_tensor * k_all = ggml_view_2d(ctx0, kr_cache, idx_dim, n_kv, + ggml_row_size(kr_cache->type, idx_dim), 0); + + ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->win_cells); + members = ggml_reshape_3d(ctx0, members, idx_dim, r, n_win); + + // the gather put each block's members contiguously in one row, which is what multi_add sums + ggml_tensor * pooled = ggml_multi_add(ctx0, + ggml_view_2d(ctx0, members, idx_dim, n_win, members->nb[2], 0), r); + pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); + cb(pooled, "qsa_k_pooled", il); + + int sections[GGML_MROPE_SECTIONS]; + std::copy(hparams.rope_sections.begin(), hparams.rope_sections.begin() + GGML_MROPE_SECTIONS, sections); + + // rope reads [n_dims, n_head, n_tokens], and a block stands in for one token here + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_win); + pooled = llm_build_context::llm_build_norm(ctx0, pooled, hparams, + model.layers[il].indexer_k_norm, nullptr, LLM_NORM_RMS, cb, il); + pooled = ggml_rope_multi(ctx0, pooled, inp->win_pos, nullptr, + bctx.n_rot, sections, bctx.rope_type, bctx.n_ctx_orig, bctx.freq_base, bctx.freq_scale, + bctx.ext_factor, bctx.attn_factor, bctx.beta_fast, bctx.beta_slow); + pooled = ggml_reshape_2d(ctx0, pooled, idx_dim, n_win); + cb(pooled, "qsa_k_win", il); + + // score against the scatter's result, not the cache, so the read depends on the write + // rather than merely following it into the graph + ggml_tensor * kp_all = ggml_set_rows(ctx0, kp_cache, pooled, inp->win_blocks); + cb(kp_all, "qsa_k_scatter", il); + + pooled = ggml_view_2d(ctx0, kp_all, idx_dim, n_blocks, + ggml_row_size(kp_cache->type, idx_dim), 0); + cb(pooled, "qsa_k", il); + + // at n_kv the cut keeps every cell and scoring cannot change the mask. The pooling above + // still had to run: it is the only chance these blocks get to enter the cache + const int32_t top_k_cells = lctx.cparams.dsa_top_k > 0 ? lctx.cparams.dsa_top_k : (int32_t) hparams.indexer_top_k; + const int32_t width = top_k_cells + r - 1; + if (width >= n_kv) { + ggml_build_forward_expand(gf, kp_all); + return KQ_mask; + } + + ggml_tensor * q = llm_build_context::llm_build_lora_mm(lctx, ctx0, model.layers[il].indexer_q_proj, cur); + q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens); + q = llm_build_context::llm_build_norm(ctx0, q, hparams, + model.layers[il].indexer_q_norm, nullptr, LLM_NORM_RMS, cb, il); + q = ggml_rope_multi(ctx0, q, inp_pos, nullptr, + bctx.n_rot, sections, bctx.rope_type, bctx.n_ctx_orig, bctx.freq_base, bctx.freq_scale, + bctx.ext_factor, bctx.attn_factor, bctx.beta_fast, bctx.beta_slow); + cb(q, "qsa_q", il); + + // causality and sequence membership come from KQ_mask rather than being rebuilt on the + // host; inp->bias carries only what is specific to the block cut + ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0); + if (causal->type != GGML_TYPE_F32) { + causal = ggml_cast(ctx0, causal, GGML_TYPE_F32); + } + ggml_tensor * cut_mask = ggml_add(ctx0, inp->bias, causal); + cb(cut_mask, "qsa_cut_mask", il); + + // top-k selects cells, not blocks: every cell carries its own block's pooled key, so the + // causal mask can still drop cells inside a selected block. The op scores and sorts + // internally, so the n_kv by n_tokens score matrix is never materialized. + if (lctx.cparams.fused_idx_topk) { + ggml_tensor * k_cells = ggml_get_rows_ext(ctx0, pooled, inp->cell_blk, true, false); + k_cells = ggml_reshape_3d(ctx0, k_cells, idx_dim, n_kv, 1); + ggml_tensor * fused = ggml_indexer_topk(ctx0, k_cells, q, inp->head_w, cut_mask, + GGML_UNARY_OP_RELU, width); + cb(fused, "qsa_top_k", il); + ggml_build_forward_expand(gf, fused); + ggml_tensor * mask = ggml_indexer_mask(ctx0, KQ_mask, fused); + cb(mask, "qsa_mask", il); + return mask; + } + + ggml_tensor * score = ggml_mul_mat(ctx0, pooled, + ggml_reshape_2d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tokens)); + score = ggml_reshape_3d(ctx0, score, n_blocks, n_idx_h, n_tokens); + 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_2d(ctx0, score, n_blocks, n_tokens); + cb(score, "qsa_score", il); + + // expanding the block indices instead would need an integer multiply-add ggml has no op + // for. get_rows gathers rows, so the scores are transposed first + ggml_tensor * expanded = ggml_get_rows(ctx0, + ggml_cont(ctx0, ggml_transpose(ctx0, score)), inp->cell_blk); + expanded = ggml_cont(ctx0, ggml_transpose(ctx0, expanded)); + + expanded = ggml_add(ctx0, expanded, cut_mask); + cb(expanded, "qsa_score_cells", il); + + ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); + cb(top_k, "qsa_top_k", il); + + ggml_tensor * mask = ggml_indexer_mask(ctx0, KQ_mask, top_k); + cb(mask, "qsa_mask", il); + + return mask; +} + +ggml_cgraph * llm_build_context::build_qwen4exp() { + + ggml_cgraph * gf = new_graph_custom(); + + delta_net delta(lctx, batch); + + const int32_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); + + const int32_t hc = hparams.dsv4_hc_mult; + + ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + ggml_tensor * KQ_mask = build_inp_KQ_mask(); + + lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); + cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); + ggml_set_input(lctx.inp_s_seq_qnext); + + float KQ_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + // the wide residual starts as hc identical copies of the embedding + ggml_tensor * res_hc = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), + n_embd, hc, n_tokens, 1); + cb(res_hc, "hc_residual", -1); + + if (hparams.ple_n_heads > 0) { + lctx.inp_ple_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, hparams.ple_n_heads * n_tokens); + cb(lctx.inp_ple_rows, "inp_ple_rows", -1); + ggml_set_input(lctx.inp_ple_rows); + } else { + lctx.inp_ple_rows = nullptr; + } + + // the same test the delta-net path uses for its own state + const bool ple_reset_state = batch.pos != nullptr && batch.pos[0] == 0; + std::vector ple_reset_pos(n_tokens, false); + for (int32_t i = 0; i < n_tokens && batch.pos != nullptr; ++i) { + ple_reset_pos[i] = batch.pos[i] == 0; + } + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inject = nullptr; + + if (hparams.is_ple(il)) { + res_hc = qwen4exp_ple(*this, ctx0, gf, lctx, model, hparams, delta, res_hc, lctx.inp_ple_rows, + kv_self.s_l[il], ple_reset_state, ple_reset_pos, n_embd, n_tokens, il, cb); + } + + ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, res_hc, + model.layers[il].hc_attn_norm, model.layers[il].hc_attn_down, + model.layers[il].hc_attn_up, model.layers[il].hc_attn_inject, + &inject, n_embd, il, cb); + + if (hparams.is_recurrent(il)) { + cur = delta.build_layer_attn_linear(ctx0, gf, cur, nullptr, il, cb, /* external_residual */ true, + GGML_UNARY_OP_SIGMOID); + } else { + // the indexer reads the same block input as q/k/v, and returns the causal mask + // itself when the layer carries no compression ratio + ggml_tensor * mask = hparams.is_qsa(il) + ? qwen4exp_qsa_mask(*this, ctx0, lctx, gf, cur, inp_pos, KQ_mask, il, cb) + : KQ_mask; + + cur = build_std_attention(gf, nullptr, cur, inp_pos, nullptr, nullptr, + mask, nullptr, nullptr, KQ_scale, 0.0f, 0, il, true, false, + /* add_input */ false, /* is_norm */ false, /* is_multi */ true); + } + + res_hc = qwen4exp_hc_combine(ctx0, hparams, res_hc, cur, inject, n_embd, il, cb); + + cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, res_hc, + model.layers[il].hc_ffn_norm, model.layers[il].hc_ffn_down, + model.layers[il].hc_ffn_up, model.layers[il].hc_ffn_inject, + &inject, n_embd, il, cb); + + cur = llm_build_std_moe_ffn(ctx0, lctx, nullptr, cur, + model.layers[il].ffn_gate_inp, nullptr, + model.layers[il].ffn_up_exps, nullptr, + model.layers[il].ffn_gate_exps, nullptr, + model.layers[il].ffn_down_exps, nullptr, + nullptr, + model.layers[il].ffn_up_shexp, nullptr, + model.layers[il].ffn_gate_shexp, nullptr, + model.layers[il].ffn_down_shexp, nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, false, 0.0f, + LLM_EXPERT_GATING_FUNC_SOFTMAX, + LLM_FFN_SILU, cb, il, gf, /* add_input */ false, model.layers[il].ffn_up_gate_exps, nullptr, + model.layers[il].ffn_gate_inp_shexp); + + res_hc = qwen4exp_hc_combine(ctx0, hparams, res_hc, cur, inject, n_embd, il, cb); + res_hc = lctx.cvec.apply_to(ctx0, res_hc, il); + cb(res_hc, "l_out", il); + } + + ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, res_hc, + model.hc_head_norm, model.hc_head_down, model.hc_head_up, + nullptr, nullptr, n_embd, -1, cb); + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + cb(cur, "result_output", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 4a33ccd5..48678c68 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -33,6 +33,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_QWEN3VLMOE, "qwen3vlmoe" }, { LLM_ARCH_QWEN35MOE, "qwen35moe" }, { LLM_ARCH_QWEN35, "qwen35" }, + { LLM_ARCH_QWEN4EXP, "qwen4exp" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_PHI2, "phi2" }, { LLM_ARCH_PHI3, "phi3" }, @@ -216,6 +217,17 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_HYPER_CONNECTION_COUNT, "%s.hyper_connection.count" }, { LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, "%s.hyper_connection.sinkhorn_iterations" }, { LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" }, + { LLM_KV_HYPER_CONNECTION_LOW_RANK, "%s.hyper_connection.low_rank" }, + + { LLM_KV_PLE_LAYERS, "%s.ple.layers" }, + { LLM_KV_PLE_NGRAM_SIZE, "%s.ple.ngram_size" }, + { LLM_KV_PLE_HEADS_PER_NGRAM, "%s.ple.heads_per_ngram" }, + { LLM_KV_PLE_CONV_KERNEL, "%s.ple.conv_kernel" }, + { LLM_KV_PLE_LAYER_MULTIPLIERS, "%s.ple.layer_multipliers" }, + { LLM_KV_PLE_HEAD_OFFSETS, "%s.ple.head_offsets" }, + { LLM_KV_PLE_HEAD_VOCAB_SIZES, "%s.ple.head_vocab_sizes" }, + { LLM_KV_PLE_EOS_TOKEN_ID, "%s.ple.eos_token_id" }, + { LLM_KV_PLE_IMAGE_TOKEN_ID, "%s.ple.image_token_id" }, { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, @@ -325,6 +337,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN35: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_BAILINGMOE3: return true; default: diff --git a/src/llama-arch.h b/src/llama-arch.h index fe3cc0f7..b210bd60 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -31,6 +31,7 @@ enum llm_arch { LLM_ARCH_QWEN3VLMOE, LLM_ARCH_QWEN35MOE, LLM_ARCH_QWEN35, + LLM_ARCH_QWEN4EXP, LLM_ARCH_MELLUM, LLM_ARCH_PHI2, LLM_ARCH_PHI3, @@ -199,6 +200,17 @@ enum llm_kv { LLM_KV_HYPER_CONNECTION_COUNT, LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, LLM_KV_HYPER_CONNECTION_EPSILON, + LLM_KV_HYPER_CONNECTION_LOW_RANK, + + LLM_KV_PLE_LAYERS, + LLM_KV_PLE_NGRAM_SIZE, + LLM_KV_PLE_HEADS_PER_NGRAM, + LLM_KV_PLE_CONV_KERNEL, + LLM_KV_PLE_LAYER_MULTIPLIERS, + LLM_KV_PLE_HEAD_OFFSETS, + LLM_KV_PLE_HEAD_VOCAB_SIZES, + LLM_KV_PLE_EOS_TOKEN_ID, + LLM_KV_PLE_IMAGE_TOKEN_ID, LLM_KV_HASH_LAYER_COUNT, @@ -423,6 +435,31 @@ enum llm_tensor { LLM_TENSOR_HC_FFN_FN, LLM_TENSOR_HC_FFN_SCALE, + // qwen4exp low-rank hyper-connections: a gated mix over the wide residual, + // unrelated to the HC_*_BASE/FN/SCALE set above + LLM_TENSOR_HC_HEAD_NORM, + LLM_TENSOR_HC_HEAD_DOWN, + LLM_TENSOR_HC_HEAD_UP, + LLM_TENSOR_HC_ATTN_NORM, + LLM_TENSOR_HC_ATTN_DOWN, + LLM_TENSOR_HC_ATTN_UP, + LLM_TENSOR_HC_ATTN_INJECT, + LLM_TENSOR_HC_FFN_NORM, + LLM_TENSOR_HC_FFN_DOWN, + LLM_TENSOR_HC_FFN_UP, + LLM_TENSOR_HC_FFN_INJECT, + + LLM_TENSOR_PLE_KEY, + LLM_TENSOR_PLE_VALUE, + LLM_TENSOR_PLE_NORM_KEY, + LLM_TENSOR_PLE_NORM_QUERY, + LLM_TENSOR_PLE_NORM_CONV, + LLM_TENSOR_PLE_CONV1D, + + LLM_TENSOR_INDEXER_Q_PROJ, + LLM_TENSOR_INDEXER_K_PROJ, + LLM_TENSOR_INDEXER_Q_NORM, + LLM_TENSOR_PER_LAYER_TOKEN_EMBD, LLM_TENSOR_PER_LAYER_MODEL_PROJ, LLM_TENSOR_PER_LAYER_INP_GATE, // 100 diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 669c92b4..61373eca 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -122,6 +122,7 @@ void llm_build_context::init() { lctx.inp_KQ_mask_cross = nullptr; lctx.inp_dsa_sink = nullptr; lctx.inp_mtp_carry = nullptr; + lctx.inp_qsa.clear(); lctx.dflash.inputs.target_features = nullptr; lctx.dflash.inputs.pos_ctx = nullptr; lctx.dflash.inputs.kq_mask = nullptr; @@ -2810,6 +2811,10 @@ ggml_cgraph * llm_build_context::llama_build_graph( { result = llm.build_mellum(); } break; + case LLM_ARCH_QWEN4EXP: + { + result = llm.build_qwen4exp(); + } break; case LLM_ARCH_QWEN3NEXT: { result = llm.build_qwen3next(); @@ -3160,7 +3165,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens auto the_k_norm = model.layers[il].attn_k_norm ? model.layers[il].attn_k_norm->extra ? ((ggml_split_tensor_t *)model.layers[il].attn_k_norm->extra)->splits[id] : model.layers[il].attn_k_norm : nullptr; ggml_tensor *Qcur, *Kcur, *Vcur, *gate = nullptr; - if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) { + if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN4EXP) { auto [Q, K, V, G] = llm_build_mul_mat_qkv_gated(gf, cur, split_wq, split_wk, split_wv, the_q_norm, the_k_norm, il); Qcur = Q; Kcur = K; Vcur = V; gate = G; @@ -3408,7 +3413,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens auto input_normed = cur; ggml_tensor *Qcur, *Kcur, *Vcur, *gate = nullptr; - if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) { + if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN4EXP) { auto [Q, K, V, G] = llm_build_mul_mat_qkv_gated(gf, cur, model.layers[il].wq, model.layers[il].wk, model.layers[il].wv, model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, il); Qcur = Q; Kcur = K; Vcur = V; gate = G; diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 81f1792b..9c3ce942 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -247,6 +247,7 @@ struct llm_build_context { ggml_cgraph * build_qwen3vlmoe(); ggml_cgraph * build_qwen3next(); + ggml_cgraph * build_qwen4exp(); ggml_cgraph * build_qwen35moe(); diff --git a/src/llama-context.h b/src/llama-context.h index dc7dd136..c1c4a548 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -131,6 +131,10 @@ struct llama_kv_cache { // Empty unless the model has the DSA indexer. std::vector kr_l; + // Pooled block keys for Qwen sparse attention: [indexer_head_size, kv_size/compress_ratio], + // already normalised and rotated. Empty unless the model scores blocks. + std::vector kp_l; + // When true, the delta_net graph builder will enable per-step SSM state saves bool save_per_step_ssm = false; @@ -625,6 +629,7 @@ struct llama_context { struct ggml_tensor * inp_s_mask; // F32 [1, n_kv] struct ggml_tensor * inp_s_seq; // I32 [n_kv, n_batch] struct ggml_tensor * inp_s_seq_qnext; // I32 [1, n_batch] + struct ggml_tensor * inp_ple_rows = nullptr; // I32 [ple_n_heads * n_batch], qwen4exp n-gram rows struct ggml_tensor * inp_pos_bucket; // I32 [n_batch|n_kv, n_batch] struct ggml_tensor * inp_embd_enc; // F32 [n_embd, n_outputs_enc] struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch] @@ -633,6 +638,32 @@ struct llama_context { struct ggml_tensor * inp_mtp_carry = nullptr; // F32 [n_embd, nextn-1] per-head hidden at the last committed position struct ggml_tensor * inp_dsa_sink = nullptr; // F32 [n_kv, n_tokens] per-sequence attention-sink boost for DSA indexer top-k + // Qwen sparse attention: everything that depends on cache layout is computed on the host, + // so the graph only gathers, pools and scores. One entry per distinct compress ratio. + struct qsa_input { + int32_t ratio = 0; + struct ggml_tensor * cell_blk = nullptr; // I32 [n_kv] block each cell belongs to + struct ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens] causal mask plus the tail boost + // only the blocks this ubatch writes into are pooled again; the rest are read from kp_l + struct ggml_tensor * win_blocks = nullptr; // I32 [n_win] which block each entry rebuilds + struct ggml_tensor * win_cells = nullptr; // I32 [ratio*n_win] member cells of those blocks + struct ggml_tensor * win_pos = nullptr; // I32 [4*n_win] mrope position of those blocks + struct ggml_tensor * head_w = nullptr; // F32 [n_idx_h, n_tokens] all ones; the head sum is unweighted + }; + std::vector inp_qsa; + + // the pooled block keys no longer match the raw indexer keys and the next built graph must + // pool every block; set on state restore and defrag, cleared by that graph's host fill + bool qsa_pooled_stale = false; + + // each sequence's recent tokens, read by the n-gram hash when a ubatch does not carry its + // first tokens' predecessors; trusted only while contiguous with the incoming position + struct ple_history { + llama_pos next_pos = -1; + std::vector toks; + }; + std::map ple_hist; + struct swa_window_view_state { bool active = false; bool compacted = false; diff --git a/src/llama-delta-net.cpp b/src/llama-delta-net.cpp index aad31c5f..416d52cc 100644 --- a/src/llama-delta-net.cpp +++ b/src/llama-delta-net.cpp @@ -406,14 +406,18 @@ ggml_tensor * delta_net::build_qkv(ggml_context * ctx0, ggml_tensor * state_stor } ggml_tensor * delta_net::build_gated_output(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_norm, ggml_tensor * ssm_out, ggml_tensor * output, ggml_tensor * z, - int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb) { + int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb, ggml_unary_op gate_op) { ggml_tensor * attn_out_2d = ggml_reshape_2d(ctx0, output, head_v_dim, num_v_heads * n_tok); ggml_tensor * z_2d = ggml_reshape_2d(ctx0, z, head_v_dim, num_v_heads * n_tok); ggml_tensor * attn_out_norm = llm_build_context::llm_build_norm(ctx0, attn_out_2d, lctx.model.hparams, ssm_norm, nullptr, LLM_NORM_RMS, cb, il); cb(attn_out_norm, "attn_rms_norm", il); - attn_out_norm = ggml_fused_mul_unary(ctx0, z_2d, attn_out_norm, GGML_UNARY_OP_SILU); + // the fused form only takes SILU when both operands share a shape, so the sigmoid + // gate is spelled out + attn_out_norm = gate_op == GGML_UNARY_OP_SIGMOID + ? ggml_mul(ctx0, attn_out_norm, ggml_sigmoid(ctx0, z_2d)) + : ggml_fused_mul_unary(ctx0, z_2d, attn_out_norm, gate_op); cb(attn_out_norm, "attn_out_norm", il); ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim*num_v_heads, n_tok); @@ -446,7 +450,8 @@ static ggml_tensor * get_input_tensor_sm_graph(ggml_context * ctx, ggml_tensor * ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_cgraph * gf, ggml_tensor * delta_input, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids, - uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb) const { + uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb, + bool external_residual, ggml_unary_op gate_op) const { const int64_t n_tok = delta_input->ne[1]; const int64_t n_seqs = 1; @@ -582,8 +587,11 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_ input->view_src = input->src[idx]; } } - auto norm = model.layers[il].attn_norm->extra ? ((ggml_split_tensor_t *)model.layers[il].attn_norm->extra)->splits[idx] : model.layers[il].attn_norm; - auto cur = llm_build_context::llm_build_norm(ctx0, input, hparams, norm, nullptr, LLM_NORM_RMS, cb, il); + ggml_tensor * cur = input; + if (model.layers[il].attn_norm) { + auto norm = model.layers[il].attn_norm->extra ? ((ggml_split_tensor_t *)model.layers[il].attn_norm->extra)->splits[idx] : model.layers[il].attn_norm; + cur = llm_build_context::llm_build_norm(ctx0, input, hparams, norm, nullptr, LLM_NORM_RMS, cb, il); + } auto [qkv_mixed, z] = build_qkvz(lctx, ctx0, model.layers[il].wqkv, model.layers[il].wqkv_gate, model.layers[il].ssm_in, head_k_dim, num_k_heads, head_v_dim, num_v_heads, cur, il, cb, gf); @@ -606,19 +614,20 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_ state_seq_id_local, qnext_state_slots, reset_state_local, hparams.f_norm_rms_eps, model.layers[il].ssm_beta_alpha ? 0 : 1, il, cb, gf, per_step_ckpt, per_step_conv); - auto gated_output = build_gated_output(lctx, ctx0, model.layers[il].ssm_norm, model.layers[il].ssm_out, output, z, head_v_dim, num_v_heads, n_tok, il, cb); + auto gated_output = build_gated_output(lctx, ctx0, model.layers[il].ssm_norm, model.layers[il].ssm_out, output, z, head_v_dim, num_v_heads, n_tok, il, cb, gate_op); if (inp_out_ids) { gated_output = ggml_get_rows(ctx0, gated_output, inp_out_ids); input = ggml_get_rows(ctx0, input, inp_out_ids); } - output = ggml_add(ctx0, gated_output, input); + output = external_residual ? gated_output : ggml_add(ctx0, gated_output, input); cb(output, "ssm_output", il); return output; } ggml_tensor * delta_net::build_layer_attn_linear(ggml_context * ctx0, ggml_cgraph * gf, - ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb) const { + ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb, + bool external_residual, ggml_unary_op gate_op) const { GGML_ASSERT(lctx.inp_s_seq_qnext != nullptr); auto & model = lctx.model; @@ -636,7 +645,7 @@ ggml_tensor * delta_net::build_layer_attn_linear(ggml_context * ctx0, ggml_cgrap if (all_same_seq) { bool reset_state = batch.pos != nullptr && batch.pos[0] == 0; - return build_layer_attn_linear_core(ctx0, gf, cur, lctx.inp_s_seq_qnext, inp_out_ids, token_seq_ids.front(), reset_state, il, cb); + return build_layer_attn_linear_core(ctx0, gf, cur, lctx.inp_s_seq_qnext, inp_out_ids, token_seq_ids.front(), reset_state, il, cb, external_residual, gate_op); } GGML_ASSERT(has_unique_seq_ids && "qwen3next mixed-sequence batches require unique sequence IDs per token"); @@ -648,7 +657,7 @@ ggml_tensor * delta_net::build_layer_attn_linear(ggml_context * ctx0, ggml_cgrap const bool reset_state_i = batch.pos != nullptr && batch.pos[i] == 0; const uint32_t state_seq_id_i = (uint32_t) token_seq_ids[i]; - ggml_tensor * out_i = build_layer_attn_linear_core(ctx0, gf, cur_i, inp_s_seq_qnext_i, inp_out_ids, state_seq_id_i, reset_state_i, il, cb); + ggml_tensor * out_i = build_layer_attn_linear_core(ctx0, gf, cur_i, inp_s_seq_qnext_i, inp_out_ids, state_seq_id_i, reset_state_i, il, cb, external_residual, gate_op); out = out == nullptr ? out_i : ggml_concat(ctx0, out, out_i, 1); } diff --git a/src/llama-delta-net.h b/src/llama-delta-net.h index 36722d8f..6be33bc8 100644 --- a/src/llama-delta-net.h +++ b/src/llama-delta-net.h @@ -19,14 +19,23 @@ struct delta_net { ggml_tensor * build_layer_attn_linear_core(ggml_context * ctx0, ggml_cgraph * gf, ggml_tensor * cur, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids, - uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb) const; + uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb, + bool external_residual = false, ggml_unary_op gate_op = GGML_UNARY_OP_SILU) const; + // external_residual: the caller has already normalised the input and owns the residual + // add, as a hyper-connection stack does. The layer then returns just its own output. ggml_tensor * build_layer_attn_linear(ggml_context * ctx0, ggml_cgraph * gf, - ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb) const; + ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb, + bool external_residual = false, ggml_unary_op gate_op = GGML_UNARY_OP_SILU) const; ggml_tensor * build_layer_attn_kda(ggml_context * ctx0, ggml_cgraph * gf, ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb) const; + // which recurrent state slot a token owns. Other per-sequence state in the same row shares + // the slot, so callers read it here rather than resolving the batch again + bool batch_shares_one_seq() const { return all_same_seq; } + uint32_t state_slot(int32_t i) const { return (uint32_t) token_seq_ids[i]; } + private: llama_context & lctx; @@ -60,8 +69,10 @@ private: float eps_norm, int repeat_type, int il, const llm_build_cb & cb, ggml_cgraph * gf, ggml_tensor * per_step_ssm = nullptr, ggml_tensor * per_step_conv = nullptr); + // gate_op selects the output gate: SiLU for the Qwen3-Next family, SIGMOID for qwen4exp static ggml_tensor * build_gated_output(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_norm, ggml_tensor * ssm_out, - ggml_tensor * output, ggml_tensor * z, int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb); + ggml_tensor * output, ggml_tensor * z, int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb, + ggml_unary_op gate_op = GGML_UNARY_OP_SILU); ggml_tensor * build_layer_attn_kda_core(ggml_context * ctx0, ggml_cgraph * gf, ggml_tensor * cur, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 4af1743b..7ecaf109 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -608,6 +608,95 @@ void llm_load_hparams( default: model.type = e_model::MODEL_UNKNOWN; } } break; + case LLM_ARCH_QWEN4EXP: + { + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); + 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); + + 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); + + 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); + + { + uint32_t full_attn_interval = 4; + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + for (uint32_t i = 0; i < hparams.n_layer; ++i) { + hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); + } + } + + { + uint32_t n_ratios = 0; + if (ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_ratios, false) && n_ratios >= hparams.n_layer) { + std::vector ratios; + ml.get_arr(ml.llm_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS), ratios); + std::copy_n(ratios.begin(), hparams.n_layer, hparams.dsv4_compress_ratios.begin()); + } + } + + // ple.layers is absent when the model carries no n-gram table, which leaves + // the module inert + { + uint32_t n_ple_layers = 0; + if (ml.get_arr_n(LLM_KV_PLE_LAYERS, n_ple_layers, false) && n_ple_layers > 0) { + std::vector ple_layers; + ml.get_arr(ml.llm_kv(LLM_KV_PLE_LAYERS), ple_layers); + for (uint32_t il : ple_layers) { + if (il < hparams.n_layer) { + hparams.ple_layer_arr[il] = true; + } + } + ml.get_key(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size); + ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram); + ml.get_key(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel); + ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim); + ml.get_key(LLM_KV_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id, false); + ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false); + + hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; + if (hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) { + throw std::runtime_error("qwen4exp: PLE geometry exceeds the supported bounds"); + } + + std::vector mults, offs, vocabs; + ml.get_arr(ml.llm_kv(LLM_KV_PLE_LAYER_MULTIPLIERS), mults); + ml.get_arr(ml.llm_kv(LLM_KV_PLE_HEAD_OFFSETS), offs); + ml.get_arr(ml.llm_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES), vocabs); + // the derived geometry indexes these arrays, and a short one would leave + // zeros that reach a modulo in the host fill + if (mults.size() < hparams.ple_ngram_size || + offs.size() < hparams.ple_n_heads || + vocabs.size() < hparams.ple_n_heads) { + throw std::runtime_error("qwen4exp: the PLE arrays are shorter than the declared geometry"); + } + for (size_t i = 0; i < mults.size() && i < LLAMA_MAX_PLE_NGRAM; ++i) hparams.ple_layer_multipliers[i] = mults[i]; + for (size_t i = 0; i < offs.size() && i < LLAMA_MAX_PLE_HEADS; ++i) hparams.ple_head_offsets[i] = offs[i]; + for (size_t i = 0; i < vocabs.size() && i < LLAMA_MAX_PLE_HEADS; ++i) hparams.ple_head_vocab_sizes[i] = vocabs[i]; + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + if (hparams.ple_head_vocab_sizes[h] == 0) { + throw std::runtime_error("qwen4exp: a PLE head declares a zero vocabulary size"); + } + } + } + } + + switch (hparams.n_layer) { + case 48: model.type = e_model::MODEL_125B_A6B; break; + default: model.type = e_model::MODEL_UNKNOWN; + } + } break; case LLM_ARCH_QWEN35MOE: { ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index ac70d591..5a14f7c7 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -7,6 +7,8 @@ #include #define LLAMA_MAX_LAYERS 512 +#define LLAMA_MAX_PLE_NGRAM 8 +#define LLAMA_MAX_PLE_HEADS 64 enum llm_expert_gating_func_type { LLM_EXPERT_GATING_FUNC_TYPE_NONE = 0, @@ -152,6 +154,21 @@ struct llama_hparams { float dsv4_hc_eps = 0.0f; std::array dsv4_compress_ratios = {}; + // qwen4exp. hc_low_rank 0 means the full-rank hyper-connection form; the + // ple_* group is inert unless the model carries an n-gram embedding layer. + uint32_t hc_low_rank = 0; + uint32_t ple_ngram_size = 0; + uint32_t ple_heads_per_ngram = 0; + uint32_t ple_conv_kernel = 0; + uint32_t ple_n_heads = 0; // (ngram_size - 1) * heads_per_ngram + uint32_t ple_head_dim = 0; + uint32_t ple_eos_token_id = 0; + uint32_t ple_image_token_id = 0; + std::array ple_layer_arr = {}; + std::array ple_layer_multipliers = {}; + std::array ple_head_offsets = {}; + std::array ple_head_vocab_sizes = {}; + // qwen3vl deepstack uint32_t n_deepstack_layers = 0; @@ -378,6 +395,26 @@ struct llama_hparams { return il < n_layer ? recurrent_layer_arr[il] : false; } + bool is_ple(uint32_t il) const { + return il < n_layer ? ple_layer_arr[il] : false; + } + + // the layer runs Qwen sparse attention over pooled blocks; deepseek4 fills the same + // ratio array for its CSA/HCA layers and reads the ratio value directly instead + bool is_qsa(uint32_t il) const { + return il < n_layer ? dsv4_compress_ratios[il] > 0 : false; + } + + // rows the PLE convolution history adds to a layer's recurrent state row: the taps reach + // (kernel - 1) * ngram_size positions back, over every channel of the wide residual + uint32_t n_embd_ple_conv(uint32_t il) const { + return is_ple(il) ? ple_conv_state() * dsv4_hc_mult * n_embd : 0; + } + + uint32_t ple_conv_state() const { + return ple_conv_kernel > 0 ? (ple_conv_kernel - 1) * ple_ngram_size : 0; + } + static bool is_float_close(float a, float b, float abs_tol) { // Check for non-negative tolerance if (abs_tol < 0.0) { diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 2ca2a212..45b346be 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -81,6 +81,8 @@ struct create_tensors_helper : public create_tensors_helper_interface { bool create_qwen3next_tensors(const LLM_TN & tn); + bool create_qwen4exp_tensors(const LLM_TN & tn); + bool create_qwen35moe_tensors(const LLM_TN & tn); bool create_qwen35_tensors(const LLM_TN & tn); @@ -1672,6 +1674,129 @@ bool create_tensors_helper::create_qwen3next_tensors(const LLM_TN & tn) { return use_mmap_buffer; } +bool create_tensors_helper::create_qwen4exp_tensors(const LLM_TN & tn) { + LOADING_PRELUDE + + const int32_t hc = hparams.dsv4_hc_mult; + const int32_t hc_dim = hc * n_embd; + const int32_t hc_rank = hparams.hc_low_rank; + + model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); + + // The wide residual is normalised and collapsed by a hyper-connection mix rather + // than by an output_norm, so this architecture carries none. + model.hc_head_norm = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), {hc_dim}); + model.hc_head_down = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), {hc_dim, hc_rank}); + model.hc_head_up = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_UP, "weight"), {hc_rank, hc_dim}); + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED); + if (model.output == NULL) { + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); + } + + if (hparams.ple_n_heads > 0) { + // The row count comes from the table itself. The converter shards the n-gram + // table, so the sum of ple_head_vocab_sizes is not the stored row count. + const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"); + const auto * ple_w = ml.get_weight(ple_name.c_str()); + if (ple_w == nullptr) { + throw std::runtime_error("qwen4exp: the PLE n-gram table is missing"); + } + const int64_t ple_rows = ple_w->tensor->ne[1]; + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + // every row this head can reach must exist in the table and fit the I32 row index + const uint64_t last = (uint64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h]; + if (last > (uint64_t) ple_rows || last > INT32_MAX) { + throw std::runtime_error("qwen4exp: a PLE head reaches past the n-gram table"); + } + } + model.tok_embd_per_layer = create_tensor(ctx_input, tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), + {hparams.ple_head_dim, ple_rows}); + } + + const bool has_moe_hparams = n_expert > 0 && n_expert_used > 0; + const int32_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : (has_moe_hparams ? n_ff / n_expert_used : n_ff); + const int32_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff_exp; + + const int32_t head_k_dim = hparams.ssm_d_state; + const int32_t num_k_heads = hparams.ssm_n_group; + const int32_t num_v_heads = hparams.ssm_dt_rank; + const int32_t head_v_dim = hparams.ssm_d_inner / num_v_heads; + const int32_t key_dim = head_k_dim * num_k_heads; + const int32_t value_dim = head_v_dim * num_v_heads; + const int32_t conv_dim = key_dim * 2 + value_dim; + + const int32_t idx_head = hparams.indexer_head_size; + const int32_t idx_n_head = hparams.indexer_n_head; + + if (n_expert_used == 0) { + throw std::runtime_error("n_expert_used must be > 0 when QWEN4EXP MoE tensors are present"); + } + + for (int i = 0; i < n_layer; ++i) { + ggml_context * ctx_layer = ctx_for_layer(i); + ggml_context * ctx_split = ctx_for_layer_split(i); + + auto & layer = model.layers[i]; + + layer.hc_attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_NORM, "weight", i), {hc_dim}); + layer.hc_attn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", i), {hc_dim, hc_rank}); + layer.hc_attn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_UP, "weight", i), {hc_rank, hc_dim}); + layer.hc_attn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", i), {hc_dim, hc}); + layer.hc_ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_NORM, "weight", i), {hc_dim}); + layer.hc_ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_DOWN, "weight", i), {hc_dim, hc_rank}); + layer.hc_ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_UP, "weight", i), {hc_rank, hc_dim}); + layer.hc_ffn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_INJECT, "weight", i), {hc_dim, hc}); + + if (hparams.is_ple(i)) { + layer.ple_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_KEY, "weight", i), {n_embd, hc_dim}); + layer.ple_value = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_VALUE, "weight", i), {n_embd, n_embd}); + layer.ple_norm_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_KEY, "weight", i), {hc_dim}); + layer.ple_norm_query = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", i), {hc_dim}); + layer.ple_norm_conv = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_CONV, "weight", i), {hc_dim}); + layer.ple_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_CONV1D, "weight", i), {hparams.ple_conv_kernel, hc_dim}); + } + + if (!hparams.is_recurrent(i)) { + // wq carries the query and an equal-width gate + layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head * 2}); + layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}); + layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}); + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}); + + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}); + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}); + + // one indexer key is shared across the indexer heads + layer.indexer_q_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, idx_head * idx_n_head}); + layer.indexer_k_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, idx_head}); + layer.indexer_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {idx_head}); + layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_head}); + } else { + layer.wqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, key_dim * 2 + value_dim}); + layer.wqkv_gate = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, value_dim}); + layer.ssm_conv1d = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {hparams.ssm_d_conv, conv_dim}); + layer.ssm_dt = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_DT, "bias", i), {hparams.ssm_dt_rank}); + layer.ssm_a = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_A_NOSCAN, i), {hparams.ssm_dt_rank}); + layer.ssm_beta = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, num_v_heads}); + layer.ssm_alpha = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), {n_embd, num_v_heads}); + layer.ssm_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_v_dim}); + layer.ssm_out = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_OUT, "weight", i), {value_dim, n_embd}); + } + + auto ffn_ctx = ctx_split; + + layer.ffn_gate_inp = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}); + use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, 0, n_ff_exp); + + layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), {n_embd}); + layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}); + layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}); + layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}); + } + + return use_mmap_buffer; +} + bool create_tensors_helper::create_qwen35moe_tensors(const LLM_TN & tn) { LOADING_PRELUDE model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); @@ -5271,6 +5396,8 @@ bool create_tensors_helper::create_tensors() { use_mmap_buffer = create_mellum_tensors(tn); break; case LLM_ARCH_QWEN3NEXT: use_mmap_buffer = create_qwen3next_tensors(tn); break; + case LLM_ARCH_QWEN4EXP: + use_mmap_buffer = create_qwen4exp_tensors(tn); break; case LLM_ARCH_QWEN35MOE: use_mmap_buffer = create_qwen35moe_tensors(tn); break; case LLM_ARCH_QWEN35: diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 86d59cc4..37d27ae9 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -721,8 +721,10 @@ bool llama_model_loader::get_arr(const std::string & key, std::vector & resul case GGUF_TYPE_UINT32: case GGUF_TYPE_BOOL: case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same_v) || (std::is_same_v)); break; + case GGUF_TYPE_UINT64: + case GGUF_TYPE_INT64: GGML_ASSERT((std::is_same_v) || (std::is_same_v)); break; default: - throw std::runtime_error(format("%s is not a float32, int32, uint32 or bool array", key.c_str())); + throw std::runtime_error(format("%s is not a float32, int32, uint32, int64, uint64 or bool array", key.c_str())); } result.resize(arr_info.length); @@ -1416,5 +1418,6 @@ template bool llama_model_loader::get_key_or_arr>(enum ll template std::enable_if::value, bool>::type llama_model_loader::get_arr_n(const std::string &, unsigned int &, bool); template std::enable_if::value, bool>::type llama_model_loader::get_arr_n(enum llm_kv, unsigned int&, bool); template bool llama_model_loader::get_arr(const std::string &, std::vector &, bool); +template bool llama_model_loader::get_arr(const std::string &, std::vector &, bool); template bool llama_model_loader::get_arr(const std::string &, std::array &, bool); template bool llama_model_loader::get_arr(const std::string &, std::array &, bool); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4e38a3b5..d1834619 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -481,6 +481,59 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" }, }, }, + { + LLM_ARCH_QWEN4EXP, + { + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "per_layer_token_embd" }, + { LLM_TENSOR_HC_HEAD_NORM, "output_hc_norm" }, + { LLM_TENSOR_HC_HEAD_DOWN, "output_hc_down" }, + { LLM_TENSOR_HC_HEAD_UP, "output_hc_up" }, + { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, + { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, + { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, + { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, + { LLM_TENSOR_ATTN_QKV, "blk.%d.attn_qkv" }, + { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, + { LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" }, + { LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" }, + { LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" }, + { LLM_TENSOR_INDEXER_K_NORM, "blk.%d.indexer.k_norm" }, + { LLM_TENSOR_SSM_CONV1D, "blk.%d.ssm_conv1d" }, + { LLM_TENSOR_SSM_DT, "blk.%d.ssm_dt" }, + { LLM_TENSOR_SSM_A_NOSCAN, "blk.%d.ssm_a" }, + { LLM_TENSOR_SSM_ALPHA, "blk.%d.ssm_alpha" }, + { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, + { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, + { LLM_TENSOR_SSM_OUT, "blk.%d.ssm_out" }, + { LLM_TENSOR_HC_ATTN_NORM, "blk.%d.hc_attn_norm" }, + { LLM_TENSOR_HC_ATTN_DOWN, "blk.%d.hc_attn_down" }, + { LLM_TENSOR_HC_ATTN_UP, "blk.%d.hc_attn_up" }, + { LLM_TENSOR_HC_ATTN_INJECT, "blk.%d.hc_attn_inject" }, + { LLM_TENSOR_HC_FFN_NORM, "blk.%d.hc_ffn_norm" }, + { LLM_TENSOR_HC_FFN_DOWN, "blk.%d.hc_ffn_down" }, + { LLM_TENSOR_HC_FFN_UP, "blk.%d.hc_ffn_up" }, + { LLM_TENSOR_HC_FFN_INJECT, "blk.%d.hc_ffn_inject" }, + { LLM_TENSOR_PLE_KEY, "blk.%d.ple_key" }, + { LLM_TENSOR_PLE_VALUE, "blk.%d.ple_value" }, + { LLM_TENSOR_PLE_NORM_KEY, "blk.%d.ple_norm_key" }, + { LLM_TENSOR_PLE_NORM_QUERY, "blk.%d.ple_norm_query" }, + { LLM_TENSOR_PLE_NORM_CONV, "blk.%d.ple_norm_conv" }, + { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, + { LLM_TENSOR_FFN_GATE_INP, "blk.%d.ffn_gate_inp" }, + { LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" }, + { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, + { LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" }, + { LLM_TENSOR_FFN_GATE_UP_EXPS, "blk.%d.ffn_gate_up_exps" }, + { LLM_TENSOR_FFN_GATE_INP_SHEXP, "blk.%d.ffn_gate_inp_shexp" }, + { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, + { LLM_TENSOR_FFN_DOWN_SHEXP, "blk.%d.ffn_down_shexp" }, + { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, + }, + }, { LLM_ARCH_QWEN3NEXT, { @@ -2301,6 +2354,7 @@ const char * llama_model_type_name(e_model type) { case MODEL_33B_A3B: return "33B.A3B"; case MODEL_35B_A3B: return "35B.A3B"; case MODEL_80B_A3B: return "80B.A3B"; + case MODEL_125B_A6B: return "125B.A6B"; case MODEL_80B_A13B: return "80B.A13B"; case MODEL_100B_A6B: return "100B.A6B"; case MODEL_106B_A12B: return "106B.A12B"; @@ -2502,7 +2556,7 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_ if (il < 0 || il >= hparams.n_layer) return 0; if (hparams.recurrent_layer_arr[il]) { auto state_sots = std::min(std::max(1, n_seq_max), kv_size); - return hparams.n_embd_v_s() * state_sots * sizeof(float); + return (hparams.n_embd_v_s() + hparams.n_embd_ple_conv(il)) * (size_t) state_sots * sizeof(float); } if (arch == LLM_ARCH_OPENPANGU) { // MLA-latent cache: K row [ckv | roped k_pe]. The value-side latent is @@ -2565,5 +2619,16 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_ llama_kv_cache::get_padding(flash_attn)); auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k(il)) * n_head_kv*rows; auto v_size = ggml_row_size(type_v, hparams.n_embd_v_gqa(il)) * rows; + // a qwen4exp sparse-attention layer also caches one raw indexer key per cell and one + // pooled key per block of `ratio` cells + if (arch == LLM_ARCH_QWEN4EXP && hparams.indexer_head_size > 0 && hparams.is_qsa(il)) { + const uint32_t ratio = hparams.dsv4_compress_ratios[il]; + k_size += ggml_row_size(idx_type_k, hparams.indexer_head_size) * (rows + (rows + ratio - 1)/ratio); + } + // a PLE layer that is not recurrent still gets a state row for its convolution history + if (hparams.n_embd_ple_conv(il) > 0) { + auto state_sots = std::min(std::max(1, n_seq_max), kv_size); + k_size += hparams.n_embd_ple_conv(il) * (size_t) state_sots * sizeof(float); + } return k_size + v_size; } diff --git a/src/llama-model.h b/src/llama-model.h index 88446bc2..af2174ad 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -114,6 +114,7 @@ enum e_model { MODEL_33B_A3B, MODEL_35B_A3B, MODEL_80B_A3B, // Qwen3-Next + MODEL_125B_A6B, // Qwen3.8-Flash-Next MODEL_80B_A13B, MODEL_100B_A6B, MODEL_106B_A12B, @@ -397,6 +398,29 @@ struct llama_layer { struct ggml_tensor * hc_ffn_base = nullptr; struct ggml_tensor * hc_ffn_fn = nullptr; struct ggml_tensor * hc_ffn_scale = nullptr; + + // qwen4exp low-rank hyper-connections + struct ggml_tensor * hc_attn_norm = nullptr; + struct ggml_tensor * hc_attn_down = nullptr; + struct ggml_tensor * hc_attn_up = nullptr; + struct ggml_tensor * hc_attn_inject = nullptr; + struct ggml_tensor * hc_ffn_norm = nullptr; + struct ggml_tensor * hc_ffn_down = nullptr; + struct ggml_tensor * hc_ffn_up = nullptr; + struct ggml_tensor * hc_ffn_inject = nullptr; + + // qwen4exp QSA indexer + struct ggml_tensor * indexer_q_proj = nullptr; + struct ggml_tensor * indexer_k_proj = nullptr; + struct ggml_tensor * indexer_q_norm = nullptr; + + // qwen4exp per-layer n-gram embedding (PLE); present on ple layers only + struct ggml_tensor * ple_key = nullptr; + struct ggml_tensor * ple_value = nullptr; + struct ggml_tensor * ple_norm_key = nullptr; + struct ggml_tensor * ple_norm_query = nullptr; + struct ggml_tensor * ple_norm_conv = nullptr; + struct ggml_tensor * ple_conv1d = nullptr; struct ggml_tensor * attn_comp_wkv = nullptr; struct ggml_tensor * attn_comp_wgate = nullptr; struct ggml_tensor * attn_comp_ape = nullptr; @@ -508,6 +532,11 @@ struct llama_model { struct ggml_tensor * hc_head_fn = nullptr; struct ggml_tensor * hc_head_scale = nullptr; + // qwen4exp: final low-rank hyper-connection mix, plus the n-gram embedding table + struct ggml_tensor * hc_head_norm = nullptr; + struct ggml_tensor * hc_head_down = nullptr; + struct ggml_tensor * hc_head_up = nullptr; + // openPangu-2.0: global mHC stream-merge module (non-block) struct ggml_tensor * mhc_merge_phi = nullptr; struct ggml_tensor * mhc_merge_alpha = nullptr; @@ -588,7 +617,7 @@ struct llama_model { size_t max_nodes(int n_tokens) const { auto n_tensors = tensors_by_name.size(); if (split_mode == LLAMA_SPLIT_MODE_GRAPH && !devices.empty()) n_tensors *= devices.size(); - if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN35) { + if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN4EXP) { return std::max(n_tokens * 40, 32u * n_tensors); } //return std::max(1024, 8*n_tensors); diff --git a/src/llama.cpp b/src/llama.cpp index b54327dd..47683357 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -660,6 +660,7 @@ static void why_not_reuse_previous(const llama_batch & u_batch, const llama_cont bool llama_context::can_reuse_graph(const llama_batch & u_batch, uint64_t seq_fingerprint, uint64_t model_state_hash) { if (!cparams.graph_reuse) return false; + if (qsa_pooled_stale) return false; // the rebuild needs a graph with the full pooling window auto the_prev = cparams.mtp_op_type == MTP_OP_NONE ? prev.get() : prev_mtp.get(); if (!the_prev || !the_prev->graph) return false; if (u_batch.embd) return false; @@ -1344,9 +1345,16 @@ static bool llama_kv_cache_init( const bool has_glm_dsa_indexer = model.arch == LLM_ARCH_GLM_DSA && hparams.indexer_head_size > 0; const bool has_openpangu_dsa_indexer = model.arch == LLM_ARCH_OPENPANGU && hparams.indexer_head_size > 0 && hparams.n_swa > 0; - if (has_glm_dsa_indexer || has_openpangu_dsa_indexer) { + // qwen4exp scores whole blocks of compress_ratio tokens, so its cached indexer keys are raw: + // pooling precedes the norm and the rotation that the other two arches fold in before caching. + const bool has_qwen4exp_indexer = + model.arch == LLM_ARCH_QWEN4EXP && hparams.indexer_head_size > 0; + if (has_glm_dsa_indexer || has_openpangu_dsa_indexer || has_qwen4exp_indexer) { cache.kr_l.resize(n_layer, nullptr); } + if (has_qwen4exp_indexer) { + cache.kp_l.resize(n_layer, nullptr); + } std::vector mem_split(model.splits.size(), 0); @@ -1382,7 +1390,10 @@ static bool llama_kv_cache_init( ggml_tensor * v = nullptr; ggml_tensor * s = nullptr; if (qnext_recurrent) { - s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_embd_v_s(), qnext_state_slots); + // a PLE layer keeps its convolution history in the tail of the same row, so the + // delta-net slice it opens with keeps the offsets every other layer uses + s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + hparams.n_embd_v_s() + hparams.n_embd_ple_conv(i), qnext_state_slots); auto s_name = std::string{"cache_s_l"} + std::to_string(i); ggml_set_name(s, s_name.c_str()); cache.s_l[i] = s; @@ -1517,6 +1528,28 @@ static bool llama_kv_cache_init( v = ggml_new_tensor_1d(ctx, this_type_v, v_ne); } + // a PLE layer that is not recurrent has no state row to extend, so it gets one + // holding nothing but the convolution history + if (hparams.n_embd_ple_conv(i) > 0 && cache.s_l[i] == nullptr) { + ggml_tensor * s_ple = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + hparams.n_embd_ple_conv(i), qnext_state_slots); + ggml_format_name(s_ple, "cache_s_l%d", i); + cache.s_l[i] = s_ple; + } + + if (has_qwen4exp_indexer && hparams.is_qsa(i)) { + const uint32_t ratio = hparams.dsv4_compress_ratios[i]; + ggml_tensor * idxk = ggml_new_tensor_2d(ctx, idx_type_k, hparams.indexer_head_size, cache.rows(i)); + ggml_format_name(idxk, "cache_kr_l%d", i); + cache.kr_l[i] = idxk; + + // one pooled key per block of `ratio` positions, so this costs 1/ratio of the raw cache + ggml_tensor * idxp = ggml_new_tensor_2d(ctx, idx_type_k, hparams.indexer_head_size, + (cache.rows(i) + ratio - 1)/ratio); + ggml_format_name(idxp, "cache_kp_l%d", i); + cache.kp_l[i] = idxp; + } + auto k_name = std::string{"cache_k_l"} + std::to_string(i); auto v_name = std::string{"cache_v_l"} + std::to_string(i); ggml_set_name(k, k_name.c_str()); @@ -5166,6 +5199,150 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { } } + for (auto & qsa : lctx.inp_qsa) { + // blocks are keyed by CELL index, so concurrent sequences never claim the same slot. + // Interleaved they are approximate, not broken: a block takes its first member's rope + // position, and one spanning a sequence boundary is dropped from the cut below + GGML_ASSERT(ggml_backend_buffer_is_host(qsa.bias->buffer)); + + const int32_t n_kv = qsa.cell_blk->ne[0]; + const int32_t r = qsa.ratio; + const int32_t n_tokens = qsa.bias->ne[1]; + const int32_t n_win = qsa.win_blocks->ne[0]; + const int32_t n_blocks = kv_self.kp_l.empty() ? 0 : (n_kv + r - 1)/r; + + int32_t * dst_cell_blk = (int32_t *) qsa.cell_blk->data; + int32_t * dst_win_blk = (int32_t *) qsa.win_blocks->data; + int32_t * dst_win_cells = (int32_t *) qsa.win_cells->data; + int32_t * dst_win_pos = (int32_t *) qsa.win_pos->data; + float * dst_bias = (float *) qsa.bias->data; + + // the fused indexer op takes a per-head weight; this architecture sums the heads plain + if (qsa.head_w) { + float * dst_w = (float *) qsa.head_w->data; + std::fill(dst_w, dst_w + ggml_nelements(qsa.head_w), 1.0f); + } + + // -1 where the block cannot be pooled. Block 0 stands in so the gather stays in range; + // the bias below is what keeps those out of the ranking + std::vector blk_of(n_kv, -1); + std::vector cell_pos(n_kv, -1); + std::vector filled(n_blocks, 0); + std::vector blk_seq(n_blocks, -1); + std::vector blk_mixed(n_blocks, false); + std::vector blk_cells(r*n_blocks, 0); + std::vector blk_pos(n_blocks, -1); + + for (int32_t j = 0; j < n_kv; ++j) { + const auto & cell = kv_self.cells[j]; + if (cell.is_empty()) { + continue; + } + + const llama_pos p = cell.pos; + const int32_t b = j/r; + + const llama_seq_id sid = *cell.seq_id.begin(); + if (blk_seq[b] < 0) { + blk_seq[b] = sid; + } else if (blk_seq[b] != sid) { + blk_mixed[b] = true; + } + + blk_of[j] = (int32_t) b; + cell_pos[j] = p; + blk_cells[b*r + j%r] = j; + if (blk_pos[b] < 0 || p < blk_pos[b]) { + blk_pos[b] = p; + } + filled[b]++; + } + + // only these blocks can have changed. A short window repeats its last entry, so the + // scatter writes the same correct value twice rather than an unrelated block + std::vector touched; + touched.reserve(n_win); + if (lctx.qsa_pooled_stale) { + for (int32_t b = 0; b < n_blocks; ++b) { + touched.push_back((int32_t) b); + } + } else + for (int32_t i = 0; i < n_tokens && i < batch.n_tokens; ++i) { + const int32_t b = (int32_t) ((kv_self.head + i)/r); + if (b < n_blocks && std::find(touched.begin(), touched.end(), b) == touched.end()) { + touched.push_back(b); + if ((int32_t) touched.size() == n_win) { + break; + } + } + } + if (touched.empty()) { + touched.push_back(0); + } + + for (int32_t w = 0; w < n_win; ++w) { + const int32_t b = touched[std::min(w, touched.size() - 1)]; + dst_win_blk[w] = b; + std::copy_n(blk_cells.begin() + b*r, r, dst_win_cells + w*r); + // the block's own first position, which equals b*r only while a single sequence + // fills the cache from zero + const int32_t bp = blk_pos[b] < 0 ? 0 : (int32_t) blk_pos[b]; + for (int32_t sec = 0; sec < GGML_MROPE_SECTIONS; ++sec) { + dst_win_pos[sec*n_win + w] = bp; + } + } + + for (int32_t j = 0; j < n_kv; ++j) { + const int32_t b = blk_of[j]; + if (b >= 0 && (filled[b] < r || blk_mixed[b])) { + blk_of[j] = -1; + } + dst_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; + } + + // KQ_mask already carries causality, so this holds only the per-block part: -inf for an + // unpoolable block, a boost for the incomplete block the query sits in + auto fill_bias = [&](int32_t i0, int32_t i1) { + for (int32_t i = i0; i < i1; ++i) { + // n_kv times n_tokens reaches 2^31 at the model's full context, so index wide + float * cur_bias = dst_bias + (size_t) i*n_kv; + + if (i >= batch.n_tokens) { + std::fill(cur_bias, cur_bias + n_kv, 0.0f); + continue; + } + + const llama_pos q = batch.pos[i]; + const llama_pos tail_start = (q + 1)/r*r; + + for (int32_t j = 0; j < n_kv; ++j) { + const llama_pos p = cell_pos[j]; + + // finite, so it can never meet a -inf and produce a nan + cur_bias[j] = (p >= tail_start && p <= q) ? 1e9f + : (blk_of[j] < 0 ? -INFINITY : 0.0f); + } + } + }; + if (n_kv >= 1024 && n_tokens >= 32) { + // the same threshold the KQ_mask fill below uses for the same iteration space + const int n_thread = std::max(1, int(std::thread::hardware_concurrency()/2)); + const int32_t npt = (n_tokens + n_thread - 1)/n_thread; + std::vector workers; + for (int32_t i0 = npt; i0 < n_tokens; i0 += npt) { + workers.emplace_back(fill_bias, i0, std::min(n_tokens, i0 + npt)); + } + fill_bias(0, std::min(n_tokens, npt)); + for (auto & w : workers) w.join(); + } else { + fill_bias(0, n_tokens); + } + + // a reused graph was built with the narrow window, so clearing on `touched` alone would + // leave the blocks it could not reach at zero + lctx.qsa_pooled_stale = (int32_t) touched.size() > n_win; + } + if (batch.token && lctx.inp_tokens) { #if IK_PRINT_TIMING == 2 auto tim1 = ggml_time_us(); @@ -5877,6 +6054,98 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { } } + if (lctx.inp_ple_rows) { + const auto & hp = lctx.model.hparams; + + const int32_t n_tokens = batch.n_tokens; + const int32_t n_gram = hp.ple_ngram_size; + const int32_t n_heads = hp.ple_n_heads; + const int32_t per_gram = hp.ple_heads_per_ngram; + const llama_token eos = hp.ple_eos_token_id; + + GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_ple_rows->buffer)); + int32_t * data = (int32_t *) lctx.inp_ple_rows->data; + + // an embedding ubatch has no token ids, and this feeds get_rows into a 320 M row table, + // so every position still needs a defined row. The reference hashes the image placeholder + // here; without that key EOS just makes the span a segment boundary + const llama_token img_tok = hp.ple_image_token_id != 0 + ? (llama_token) hp.ple_image_token_id + : eos; + auto tok_of = [&](int32_t k) -> llama_token { + return batch.token ? batch.token[k] : img_tok; + }; + + // snapshot before any update: one pass would let a token read an earlier token of this + // same ubatch as prior context + std::map> snap; + for (int32_t i = 0; i < n_tokens; ++i) { + const llama_seq_id seq = batch.seq_id[i][0]; + if (snap.count(seq)) { + continue; + } + auto & h = lctx.ple_hist[seq]; + if (h.next_pos != batch.pos[i]) { + h.toks.assign(n_gram - 1, eos); + } + h.toks.resize(n_gram - 1, eos); + snap[seq] = h.toks; + } + + for (int32_t i = 0; i < n_tokens; ++i) { + const llama_pos pos = batch.pos[i]; + const llama_seq_id seq = batch.seq_id[i][0]; + + const auto & hist = snap[seq]; + + // predecessor s (1-based) of this token: from the ubatch when it is there, from the + // sequence's own history when it is not, EOS past a segment boundary + auto prev = [&](int32_t s) -> llama_token { + const int32_t j = i - s; + if (j >= 0 && batch.seq_id[j][0] == seq && batch.pos[j] == pos - s) { + return tok_of(j); + } + // s - i positions before this ubatch started, most recent last + const int32_t back = s - i; + const int32_t k = (int32_t) hist.size() - back; + if (back > 0 && k >= 0 && pos - s >= 0) { + return hist[k]; + } + return eos; + }; + + std::vector ctx_toks(n_gram); + ctx_toks[0] = tok_of(i); + bool cut = false; + for (int32_t s = 1; s < n_gram; ++s) { + ctx_toks[s] = cut ? eos : prev(s); + if (ctx_toks[s] == eos) { + cut = true; + } + } + + for (int32_t n = 2; n <= n_gram; ++n) { + uint64_t mixed = (uint64_t) ctx_toks[0] * hp.ple_layer_multipliers[0]; + for (int32_t j = 1; j < n; ++j) { + mixed ^= (uint64_t) ctx_toks[j] * hp.ple_layer_multipliers[j]; + } + const int32_t base = (n - 2) * per_gram; + for (int32_t g = 0; g < per_gram; ++g) { + const int32_t h_i = base + g; + data[i * n_heads + h_i] = + (int32_t) (mixed % hp.ple_head_vocab_sizes[h_i] + hp.ple_head_offsets[h_i]); + } + } + + auto & h = lctx.ple_hist[seq]; + h.toks.push_back(tok_of(i)); + if ((int32_t) h.toks.size() > n_gram - 1) { + h.toks.erase(h.toks.begin(), h.toks.end() - (n_gram - 1)); + } + h.next_pos = pos + 1; + } + } + if (lctx.inp_pos_bucket) { const int64_t n_tokens = batch.n_tokens; @@ -7041,8 +7310,9 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) { // TODO: tmp fix https://github.com/ggerganov/llama.cpp/issues/6685#issuecomment-2057579516 // DSA: build_defrag additionally moves the indexer-key cache (kr_l), +3 tensors/layer/move // (src view, dst view, copy), so budget 9*n_layer per move when the indexer cache is present. - const bool has_dsa_indexer_defrag = - lctx.model.arch == LLM_ARCH_GLM_DSA && !kv_self.kr_l.empty(); + // build_defrag moves kr_l for whichever layer holds one, so the budget follows the cache + // rather than the architecture that allocated it. + const bool has_dsa_indexer_defrag = !kv_self.kr_l.empty(); const uint32_t tensors_per_move = has_dsa_indexer_defrag ? 9 : 6; const uint32_t max_moves = (lctx.model.max_nodes(1) - 2*n_layer)/(tensors_per_move*n_layer); @@ -7153,6 +7423,10 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) { return; } + // blocks are keyed by cell index, so moving cells changes every touched block's membership; + // the pooled block keys must be rebuilt from the moved indexer keys + lctx.qsa_pooled_stale = !kv_self.kp_l.empty(); + //LLAMA_LOG_INFO("(tmp log) KV defrag cell moves: %u\n", n_moves); //LLAMA_LOG_INFO("expected gf nodes: %u\n", 6*n_moves*n_layer); @@ -8574,6 +8848,11 @@ struct llama_context * llama_init_from_model( memory_size_k_indexer += ggml_nbytes(k); } } + for (auto & k : ctx->kv_self.kp_l) { + if (k) { + memory_size_k_indexer += ggml_nbytes(k); + } + } for (auto & v : ctx->kv_self.v_l) { if (v) { memory_size_v += ggml_nbytes(v); @@ -8952,6 +9231,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { case LLM_ARCH_QWEN3VLMOE: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN35: + case LLM_ARCH_QWEN4EXP: return LLAMA_ROPE_TYPE_IMROPE; // all model arches should be listed explicitly here @@ -10876,6 +11156,11 @@ struct llama_data_read { } } + // the pooled block keys are derived from the indexer keys restored just above, and the + // builder rebuilds only the blocks an ubatch writes into, so the next graph must pool + // every block once + ctx->qsa_pooled_stale = !kv_self.kp_l.empty(); + if (ctx->model.arch == LLM_ARCH_OPENPANGU && !read_openpangu_state(ctx, n_layer, seq_id, false)) { return false;