Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2#include "llama-impl.h"3#include "llama-memory-hybrid-idx.h"4#include "llama-memory-recurrent.h"5 6#include <algorithm>7#include <cinttypes>8 9// bad metadata must be catchable: GGML_ASSERT aborts the whole process10static void qwen4exp_require_nonzero(const llama_model_loader & ml, llm_kv kid, uint32_t value) {11 if (value == 0) {12 throw std::runtime_error(format("%s must be greater than zero, got %u", ml.llm_kv(kid).c_str(), value));13 }14}15 16// get_arr() copies a short array as-is, leaving a zero tail the n-gram hash silently drops17static void qwen4exp_require_arr_len(llama_model_loader & ml, llm_kv kid, uint32_t n_min) {18 uint32_t n_arr = 0;19 ml.get_arr_n(kid, n_arr, true);20 if (n_arr < n_min) {21 throw std::runtime_error(format("%s has %u entries, but at least %u are required",22 ml.llm_kv(kid).c_str(), n_arr, n_min));23 }24}25 26void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {27 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all, false);28 ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);29 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);30 31 ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);32 33 ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);34 ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);35 ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);36 ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);37 ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);38 qwen4exp_require_nonzero(ml, LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);39 qwen4exp_require_nonzero(ml, LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);40 qwen4exp_require_nonzero(ml, LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);41 qwen4exp_require_nonzero(ml, LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);42 qwen4exp_require_nonzero(ml, LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);43 44 // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank)45 ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult);46 ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);47 // a count of 1 has nothing to mix: transformers configuration_qwen4_exp.py:196, vLLM48 // config.py:49 and SGLang configs/qwen4_exp.py:38 all raise on hc_count <= 149 if (hparams.dsv4_hc_mult <= 1) {50 throw std::runtime_error(format("%s must be greater than one, got %u",51 ml.llm_kv(LLM_KV_HYPER_CONNECTION_COUNT).c_str(), hparams.dsv4_hc_mult));52 }53 qwen4exp_require_nonzero(ml, LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);54 hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;55 56 ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);57 ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);58 ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);59 qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);60 qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);61 qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);62 ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false);63 64 // PLE n-gram hash embeddings; if the key group is absent every field stays zero65 hparams.is_ple_impl.reset();66 hparams.ple_n_heads = 0;67 68 uint32_t n_ple = 0;69 ml.get_arr_n(LLM_KV_PLE_LAYERS, n_ple, false);70 if (n_ple > 0) {71 std::vector<uint32_t> ple_layers;72 ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers);73 if (n_ple != 1) {74 // hparams holds one set of hash constants, so several PLE modules cannot be represented75 throw std::runtime_error(format("%s lists %u layers, but only one PLE layer is supported",76 ml.llm_kv(LLM_KV_PLE_LAYERS).c_str(), n_ple));77 }78 for (uint32_t il : ple_layers) {79 if (il >= hparams.n_layer_all) {80 throw std::runtime_error(format("PLE layer %u is out of range", il));81 }82 hparams.is_ple_impl.set(il);83 }84 85 ml.get_key(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size);86 ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram);87 ml.get_key(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel);88 ml.get_key(LLM_KV_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id);89 // optional: files written before this key fall back to the EOS token90 ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false);91 ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);92 qwen4exp_require_nonzero(ml, LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel);93 qwen4exp_require_nonzero(ml, LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);94 95 hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;96 hparams.ple_head_dim = hparams.n_embd_per_layer;97 if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) {98 throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size));99 }100 if (hparams.ple_n_heads == 0 || hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS) {101 throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads));102 }103 104 qwen4exp_require_arr_len(ml, LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_ngram_size);105 qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_OFFSETS, hparams.ple_n_heads);106 qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_VOCAB_SIZES, hparams.ple_n_heads);107 108 ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers);109 110 // the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses111 std::array<uint64_t, LLAMA_MAX_PLE_HEADS> head_offsets = {};112 std::array<uint64_t, LLAMA_MAX_PLE_HEADS> head_vocab_sizes = {};113 ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, head_offsets);114 ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, head_vocab_sizes);115 for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {116 if (head_vocab_sizes[h] == 0 ||117 head_offsets[h] > INT32_MAX ||118 head_vocab_sizes[h] > INT32_MAX ||119 head_offsets[h] + head_vocab_sizes[h] > INT32_MAX) {120 throw std::runtime_error(format("PLE head %u range does not fit the int32 row index", h));121 }122 hparams.ple_head_offsets[h] = (uint32_t) head_offsets[h];123 hparams.ple_head_vocab_sizes[h] = (uint32_t) head_vocab_sizes[h];124 }125 }126 127 // linear attention everywhere except every full_attention_interval-th layer128 if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {129 uint32_t full_attn_interval = 4;130 ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);131 qwen4exp_require_nonzero(ml, LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval);132 for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {133 hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);134 }135 }136 137 // the PLE conv history is a row of the recurrent cache, which linear layers alone have138 for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {139 if (hparams.is_ple(i) && !hparams.is_recr(i)) {140 throw std::runtime_error(format("PLE layer %u is not a linear attention layer", i));141 }142 }143 144 switch (hparams.n_layer()) {145 case 48: type = LLM_TYPE_A3B; break;146 default: type = LLM_TYPE_UNKNOWN;147 }148}149 150void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {151 LLAMA_LOAD_LOCALS;152 153 const int64_t hc = hparams.dsv4_hc_mult;154 const int64_t hc_dim = hc * n_embd;155 const int64_t hc_lr = hparams.hc_low_rank;156 157 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);158 159 // there is no output_norm: the final hyper-connection mixer carries it160 // the gammas load as [n_embd, hc] so the grouped norm multiplies them without a graph reshape161 hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { n_embd, hc }, TENSOR_ALLOW_RESHAPE);162 hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, 0);163 hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, 0);164 165 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);166 if (output == NULL) {167 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);168 }169 170 // flat [ple_head_dim, n_rows] gather target171 if (hparams.ple_n_heads > 0) {172 // the head ranges are what the gather indexes, so they set the minimum row count173 int64_t ple_rows = 0;174 for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {175 ple_rows = std::max(ple_rows, (int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h]);176 }177 178 // the converter pads the table; a model synthesised from metadata has no tensor to ask179 const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();180 if (const auto * ple_w = ml.get_weight(ple_name.c_str())) {181 if (ple_w->tensor->ne[1] < ple_rows) {182 throw std::runtime_error(format("%s has %" PRId64 " rows, too few for the PLE head ranges (%" PRId64 ")",183 ple_name.c_str(), ple_w->tensor->ne[1], ple_rows));184 }185 ple_rows = ple_w->tensor->ne[1];186 }187 188 per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"),189 { hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY);190 }191 192 for (int il = 0; il < n_layer; ++il) {193 auto & layer = layers[il];194 195 const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used;196 const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;197 198 const int64_t head_k_dim = hparams.ssm_d_state;199 const int64_t head_v_dim = hparams.ssm_d_state;200 const int64_t n_k_heads = hparams.ssm_n_group;201 const int64_t n_v_heads = hparams.ssm_dt_rank;202 const int64_t key_dim = head_k_dim * n_k_heads;203 const int64_t value_dim = head_v_dim * n_v_heads;204 const int64_t conv_dim = key_dim * 2 + value_dim;205 206 // two HC modules per layer: before the token mixer, before the MoE207 layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { n_embd, hc }, TENSOR_ALLOW_RESHAPE);208 layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, 0);209 layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, 0);210 layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0);211 layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { n_embd, hc }, TENSOR_ALLOW_RESHAPE);212 layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, 0);213 layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, 0);214 layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0);215 216 if (!hparams.is_recr(il)) {217 // full attention: wq holds [q|gate] interleaved per head218 create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);219 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0);220 221 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);222 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);223 224 const int64_t idx_dim = hparams.indexer_head_size;225 layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0);226 layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0);227 layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, 0);228 layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, 0);229 } else {230 layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, 0);231 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0);232 layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0);233 layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, 0);234 layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, 0);235 layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0);236 layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0);237 layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0);238 layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0);239 }240 241 if (hparams.is_ple(il)) {242 layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0);243 layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0);244 layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { n_embd, hc }, TENSOR_ALLOW_RESHAPE);245 layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { n_embd, hc }, TENSOR_ALLOW_RESHAPE);246 layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { n_embd, hc }, TENSOR_ALLOW_RESHAPE);247 layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0);248 }249 250 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0);251 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);252 create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);253 254 layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0);255 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);256 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);257 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0);258 }259}260 261std::unique_ptr<llm_graph_context> llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const {262 return std::make_unique<graph>(*this, params);263}264 265// Hyper-connections keep hc parallel residual streams [n_embd, hc, T] in place of layer norms.266// Returns the mixed [n_embd, T] stream; `inject` gets the [hc, T] scatter weights.267ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix(268 ggml_tensor * x,269 ggml_tensor * w_norm,270 ggml_tensor * w_down,271 ggml_tensor * w_up,272 ggml_tensor * w_inject,273 ggml_tensor ** inject,274 int il) {275 const int64_t hc = hparams.dsv4_hc_mult;276 const int64_t hc_dim = hc * n_embd;277 const int64_t nt = x->ne[2];278 279 // grouped RMSNorm: reduce over one stream, then scale all streams with the [n_embd, hc] gamma280 // the converter folded each gamma to (1 + w)281 ggml_tensor * xn = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), w_norm);282 xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt);283 cb(xn, "hc_norm", il);284 285 ggml_tensor * lo = build_lora_mm(w_down, xn);286 lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc));287 ggml_tensor * gate = build_lora_mm(w_up, lo);288 cb(gate, "hc_gate", il);289 290 ggml_tensor * mixed = nullptr;291 if (cparams.fused_dsv4_hc_pre && il >= 0) {292 // sigmoid gate and mean over the streams in one op293 mixed = ggml_dsv4_hc_pre_gated(ctx0,294 ggml_reshape_3d(ctx0, xn, n_embd, hc, nt),295 ggml_reshape_3d(ctx0, gate, n_embd, hc, nt), 1.0f / (float) hc);296 res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, mixed, il});297 } else {298 ggml_tensor * gated = ggml_mul(ctx0, xn, ggml_sigmoid(ctx0, gate));299 gated = ggml_reshape_3d(ctx0, gated, n_embd, hc, nt);300 301 // collapse the streams by their mean302 mixed = ggml_view_2d(ctx0, gated, n_embd, nt,303 ggml_row_size(gated->type, n_embd) * hc, 0);304 mixed = ggml_cont(ctx0, mixed);305 for (int64_t c = 1; c < hc; ++c) {306 ggml_tensor * s = ggml_view_2d(ctx0, gated, n_embd, nt,307 ggml_row_size(gated->type, n_embd) * hc,308 ggml_row_size(gated->type, n_embd) * c);309 mixed = ggml_add(ctx0, mixed, s);310 }311 mixed = ggml_scale(ctx0, mixed, 1.0f / (float) hc);312 }313 cb(mixed, "hc_mixed", il);314 315 if (inject) {316 *inject = build_lora_mm(w_inject, xn);317 cb(*inject, "hc_inject", il);318 }319 320 return mixed;321}322 323ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine(324 ggml_tensor * residual,325 ggml_tensor * block_out,326 ggml_tensor * inject,327 int il) {328 const int64_t hc = hparams.dsv4_hc_mult;329 const int64_t nt = residual->ne[2];330 331 // 2*sigmoid centres the scatter weights on 1, so a zero injection is a plain residual add332 ggml_tensor * w = ggml_sigmoid(ctx0, ggml_scale(ctx0, inject, 1.0f / (float) hc));333 w = ggml_scale(ctx0, w, 2.0f);334 335 ggml_tensor * cur = nullptr;336 if (cparams.fused_dsv4_hc_post && il >= 0) {337 // identity comb: every stream adds the same block output, scaled by its own weight338 cur = ggml_dsv4_hc_post(ctx0, block_out, residual, w, nullptr);339 res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, cur, il});340 } else {341 w = ggml_reshape_3d(ctx0, w, 1, hc, nt);342 343 ggml_tensor * b = ggml_reshape_3d(ctx0, block_out, n_embd, 1, nt);344 b = ggml_repeat_4d(ctx0, b, n_embd, hc, nt, 1);345 346 cur = ggml_add(ctx0, residual, ggml_mul(ctx0, b, w));347 }348 cb(cur, "hc_combine", il);349 350 return cur;351}352 353llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_params & params) :354 llm_build_delta_net_base(params), model(model) {355 const int64_t hc = hparams.dsv4_hc_mult;356 357 GGML_ASSERT(hparams.n_embd_head_v() == hparams.n_embd_head_k());358 359 int sections[4];360 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);361 362 ggml_tensor * inpL = build_inp_embd(model.tok_embd);363 cb(inpL, "model.input_embed", -1);364 ggml_build_forward_expand(gf, inpL);365 366 auto * inp = build_inp_mem_hybrid();367 368 // qwen4exp always builds llama_memory_hybrid_idx, so this downcast is safe369 // the indexer cache inside it is absent when the GGUF has no indexer tensors370 const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(inp->mctx);371 372 const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();373 if (mctx_idx) {374 GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() &&375 "the indexer cache must track the attention cache cell for cell");376 }377 378 ggml_tensor * inp_pos = build_inp_pos();379 ggml_tensor * inp_out_ids = build_inp_out_ids();380 381 ggml_tensor * ple_emb = nullptr;382 if (hparams.ple_n_heads > 0) {383 ple_emb = build_inp_ple(mctx_hyb);384 // make sure ple_emb and build_inp_embd are in the same graph split385 ggml_build_forward_expand(gf, ple_emb);386 }387 388 // the wide residual starts as hc identical copies of the embedding389 ggml_tensor * res_hc = ggml_repeat_4d(ctx0,390 ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens),391 n_embd, hc, n_tokens, 1);392 cb(res_hc, "hc_init", -1);393 394 for (int il = 0; il < n_layer; ++il) {395 res->t_layer_inp[il] = res_hc;396 397 if (hparams.is_ple(il)) {398 res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il);399 }400 401 ggml_tensor * inject = nullptr;402 ggml_tensor * cur = build_hc_mix(res_hc,403 model.layers[il].hc_attn_norm,404 model.layers[il].hc_attn_down,405 model.layers[il].hc_attn_up,406 model.layers[il].hc_attn_inject,407 &inject, il);408 409 ggml_build_forward_expand(gf, cur);410 411 if (hparams.is_recr(il)) {412 cur = build_layer_attn_linear(inp->get_recr(), cur, il);413 } else {414 cur = build_layer_attn(inp->get_attn(), mctx_hyb, cur, inp_pos, sections, il);415 }416 417 if (il == n_layer - 1 && inp_out_ids) {418 // everything below is per token, so drop the rows that produce no output419 cur = ggml_get_rows(ctx0, cur, inp_out_ids);420 inject = ggml_get_rows(ctx0, inject, inp_out_ids);421 422 res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]);423 res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids);424 res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]);425 }426 427 res_hc = build_hc_combine(res_hc, cur, inject, il);428 429 cur = build_hc_mix(res_hc,430 model.layers[il].hc_ffn_norm,431 model.layers[il].hc_ffn_down,432 model.layers[il].hc_ffn_up,433 model.layers[il].hc_ffn_inject,434 &inject, il);435 436 cur = build_layer_ffn(cur, il);437 cb(cur, "ffn_out", il);438 439 res_hc = build_hc_combine(res_hc, cur, inject, il);440 441 // "l_last" is the layer output name that build_cvec and imatrix look for442 cb(res_hc, "l_last", il);443 }444 445 // the final mixer is the output norm: there is no separate one446 ggml_tensor * cur = build_hc_mix(res_hc,447 model.hc_head_norm, model.hc_head_down, model.hc_head_up,448 nullptr, nullptr, -1);449 450 cb(cur, "result_norm", -1);451 res->t_embd = cur;452 453 cur = build_lora_mm(model.output, cur, model.output_s);454 cb(cur, "result_output", -1);455 res->t_logits = cur;456 457 ggml_build_forward_expand(gf, cur);458}459 460std::pair<ggml_tensor *, ggml_tensor *> llama_model_qwen4exp::graph::build_qkvz(461 ggml_tensor * input,462 int il) {463 const int64_t n_seqs = ubatch.n_seqs;464 const int64_t n_seq_tokens = ubatch.n_seq_tokens;465 466 ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s);467 qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs);468 cb(qkv_mixed, "linear_attn_qkv_mixed", il);469 470 ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, input, model.layers[il].wqkv_gate_s);471 cb(z, "z", il);472 473 return { qkv_mixed, z };474}475 476ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated(477 ggml_tensor * input,478 ggml_tensor * weights,479 ggml_tensor * gate,480 int layer) {481 // the one numerical difference from Qwen3.5's GDN: sigmoid output gate, not silu482 ggml_tensor * normalized = build_norm(input, weights, nullptr, LLM_NORM_RMS, layer);483 ggml_tensor * gated = ggml_sigmoid(ctx0, gate);484 485 return ggml_mul(ctx0, normalized, gated);486}487 488// QSA attends to a budget of whole blocks of compress_ratio tokens, plus the incomplete tail489// one mean-pooled indexer key scores each block; set_input resolves the cache layout490class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i {491public:492 llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) :493 mctx(mctx), ratio(ratio), blk_bias(blk_bias) {}494 virtual ~llm_graph_input_qsa() = default;495 496 void set_input(const llama_ubatch * ubatch) override {497 mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);498 mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias);499 }500 501 bool can_reuse(const llm_graph_params & params) override {502 mctx = static_cast<const llama_memory_hybrid_idx_context *>(params.mctx);503 504 const auto * idx = mctx->get_idx();505 if (idx == nullptr) {506 return false;507 }508 509 const int64_t n_kv = idx->get_n_kv();510 const int64_t n_stream = mctx->get_n_stream();511 const int64_t n_blocks = (n_kv + ratio - 1)/ratio;512 513 bool res = true;514 515 res &= params.ubatch.n_tokens % n_stream == 0;516 517 res &= k_idxs->ne[0] == params.ubatch.n_tokens;518 res &= cell_blk->ne[0] == n_kv;519 res &= cell_blk->ne[1] == n_stream;520 res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks;521 res &= blk_pos->ne[0] == 4*n_blocks*n_stream;522 res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv);523 res &= bias->ne[1] == params.ubatch.n_tokens/n_stream;524 525 return res;526 }527 528 // per stream: a cell index names a different token in each stream529 ggml_tensor * k_idxs = nullptr; // I32 [n_tokens]530 ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream]531 ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream]532 ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream]533 ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream]534 535 const llama_memory_hybrid_idx_context * mctx;536 const uint32_t ratio;537 538 // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded539 const bool blk_bias;540};541 542ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(543 const llama_memory_hybrid_idx_context * mctx_hyb,544 ggml_tensor * cur,545 ggml_tensor * inp_pos,546 ggml_tensor * kq_mask,547 int * sections,548 int il) {549 const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();550 551 const int64_t idx_dim = hparams.indexer_head_size;552 const int64_t n_idx_h = hparams.indexer_n_head;553 const int64_t r = hparams.dsv4_compress_ratios[il];554 const int64_t n_kv = mctx_idx->get_n_kv();555 556 GGML_ASSERT(r > 0);557 558 const int64_t n_blocks = (n_kv + r - 1)/r;559 560 // build_attn_qsa and the KQ mask need the tokens to divide evenly across the streams561 const int64_t n_stream = mctx_hyb->get_n_stream();562 GGML_ASSERT(n_tokens % n_stream == 0);563 const int64_t n_tps = n_tokens/n_stream;564 565 // only the "which block is visible" half of the bias varies per block566 // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells567 // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out568 // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there569 const bool blk_bias = kq_mask != nullptr &&570 kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream &&571 cparams.causal_attn && !hparams.use_alibi;572 573 // nothing above depends on the layer, so the layers sharing a ratio share one input set574 llm_graph_input_qsa * inp = nullptr;575 576 const auto it = qsa_inps.find((uint32_t) r);577 if (it != qsa_inps.end()) {578 inp = it->second;579 } else {580 auto qsa = std::make_unique<llm_graph_input_qsa>(mctx_hyb, (uint32_t) r, blk_bias);581 582 qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch);583 qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream);584 qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream);585 qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream);586 qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream);587 588 ggml_set_input(qsa->cell_blk);589 ggml_set_input(qsa->blk_cells);590 ggml_set_input(qsa->blk_pos);591 ggml_set_input(qsa->bias);592 593 inp = qsa.get();594 res->add_input(std::move(qsa));595 qsa_inps.emplace((uint32_t) r, inp);596 }597 598 // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither599 ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur);600 k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens);601 cb(k_raw, "indexer_k_raw", il);602 603 ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il));604 605 // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream].606 ggml_tensor * k_all = mctx_idx->get_k(ctx0, il);607 k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0);608 609 // gathers per stream: blk_cells row s indexes stream s's own cells610 ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells);611 members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream);612 613 // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows614 ggml_tensor * pooled = nullptr;615 for (int64_t i = 0; i < r; ++i) {616 ggml_tensor * slice = ggml_cont(ctx0,617 ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream,618 members->nb[2], members->nb[3], i*members->nb[1]));619 pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice;620 }621 pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r);622 cb(pooled, "indexer_k_pooled", il);623 624 // count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536625 pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks*n_stream, 1);626 pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);627 628 // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after.629 pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream);630 pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr,631 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,632 ext_factor, attn_factor, beta_fast, beta_slow);633 pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream);634 cb(pooled, "indexer_k", il);635 636 ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur);637 q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens);638 q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il);639 q = ggml_rope_multi(ctx0, q, inp_pos, nullptr,640 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,641 ext_factor, attn_factor, beta_fast, beta_slow);642 cb(q, "indexer_q", il);643 644 // rectify each head dot product before the sum, as in the DeepSeek lightning indexer645 // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s646 ggml_tensor * score = ggml_mul_mat(ctx0, pooled,647 ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h*n_tps, n_stream));648 score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);649 score = ggml_relu(ctx0, score);650 651 // the heads sit side by side on ne[1] and there are only a few of them652 ggml_tensor * summed = nullptr;653 for (int64_t h = 0; h < n_idx_h; ++h) {654 ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream,655 score->nb[2], score->nb[3], h*score->nb[1]);656 summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice);657 }658 659 score = summed;660 cb(score, "indexer_score", il);661 662 // one value per block, so it is cheaper to bias here than after the cells are expanded663 if (blk_bias) {664 score = ggml_add(ctx0, score, inp->bias);665 }666 667 // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary668 ggml_tensor * expanded = ggml_get_rows(ctx0,669 ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk);670 expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3));671 672 if (blk_bias) {673 // flash attention keeps the mask in f16; the scores are f32674 ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32);675 expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream));676 } else {677 expanded = ggml_add(ctx0, expanded, inp->bias);678 }679 cb(expanded, "indexer_score_tokens", il);680 681 // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail682 const int64_t width = std::min<int64_t>(n_kv, (int64_t) hparams.indexer_top_k + r - 1);683 684 ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width));685 686 // build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask.687 top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream);688 cb(top_k, "indexer_top_k", il);689 690 return top_k;691}692 693// Dense GQA self-attention restricted to the cells that top_k names.694// The mask build below copies the MLA sparse path in llm_graph_context::build_attn.695ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(696 llm_graph_input_attn_kv * inp,697 ggml_tensor * q_cur,698 ggml_tensor * k_cur,699 ggml_tensor * v_cur,700 ggml_tensor * top_k,701 float kq_scale,702 int il) {703 // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer704 // has already scored with its own query in build_qsa_top_k, so top_k is unaffected.705 if (inp->self_k_rot) {706 q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);707 k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);708 }709 710 if (inp->self_v_rot) {711 v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);712 }713 714 // these nodes are added to the graph together so that they are not reordered715 // by doing so, the number of splits in the graph is reduced716 // expand k later to enable rope fusion which directly writes into k-v cache717 ggml_build_forward_expand(gf, q_cur);718 ggml_build_forward_expand(gf, v_cur);719 ggml_build_forward_expand(gf, k_cur);720 721 const auto * mctx_cur = inp->mctx;722 723 // store to KV cache724 {725 const auto & k_idxs = inp->get_k_idxs();726 const auto & v_idxs = inp->get_v_idxs();727 728 ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));729 ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));730 }731 732 ggml_tensor * kq_mask = inp->get_kq_mask();733 734 // prepare new kq mask - starts filled with -INFINITY735 ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);736 737 // reshape KQ mask into tensor with rows of size 1:738 // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]739 kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);740 741 // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]742 ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);743 744 // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]745 // this will be our source of zero values for unmasking top k mask elements746 ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);747 zeros = ggml_fill(ctx0, zeros, 0.0f);748 749 // modify KQ mask by unmasking elements that are in top_k indices750 // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])751 ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);752 753 // reshape to restore the original shape of KQ mask:754 // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]755 kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);756 757 // combine with the original kq mask758 kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);759 760 ggml_tensor * q = q_cur;761 ggml_tensor * k = mctx_cur->get_k(ctx0, il);762 ggml_tensor * v = mctx_cur->get_v(ctx0, il);763 764 // TODO: enable sparse attention when we are ready765 // ref: https://github.com/ggml-org/llama.cpp/pull/27970766 //ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, top_k->ne[0], kq_scale, il);767 ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il);768 cb(cur, "kqv_out", il);769 770 // the rotation is its own inverse, so undo it on the value side of the output771 if (inp->self_v_rot) {772 cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);773 }774 775 return cur;776}777 778ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(779 llm_graph_input_attn_kv * inp,780 const llama_memory_hybrid_idx_context * mctx_hyb,781 ggml_tensor * cur,782 ggml_tensor * inp_pos,783 int * sections,784 int il) {785 const int64_t n_embd_head = hparams.n_embd_head_v();786 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());787 788 // indexer reads the same block input as q/k/v; no cache or no ratio means dense789 const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0;790 791 ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr;792 793 // Qwen3Next uses a single Q projection that outputs query + gate794 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 ]795 cb(Qcur_full, "Qcur_full", il);796 797 ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens,798 ggml_element_size(Qcur_full) * n_embd_head * 2,799 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0);800 cb(Qcur, "Qcur_reshaped", il);801 802 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);803 cb(Qcur, "Qcur_normed", il);804 805 ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s);806 cb(Kcur, "Kcur", il);807 808 ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s);809 cb(Vcur, "Vcur", il);810 811 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);812 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);813 cb(Kcur, "Kcur_normed", il);814 815 ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens,816 ggml_element_size(Qcur_full) * n_embd_head * 2,817 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,818 ggml_element_size(Qcur_full) * n_embd_head);819 gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);820 cb(gate, "gate_reshaped", il);821 822 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);823 824 // Apply IMRoPE825 Qcur = ggml_rope_multi(826 ctx0, Qcur, inp_pos, nullptr,827 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,828 ext_factor, attn_factor, beta_fast, beta_slow829 );830 831 Kcur = ggml_rope_multi(832 ctx0, Kcur, inp_pos, nullptr,833 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,834 ext_factor, attn_factor, beta_fast, beta_slow835 );836 837 cb(Qcur, "Qcur", il);838 cb(Kcur, "Kcur", il);839 cb(Vcur, "Vcur", il);840 841 const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;842 843 if (top_k) {844 cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il);845 } else {846 cur = build_attn(inp,847 nullptr, nullptr, nullptr,848 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);849 }850 cb(cur, "attn_pregate", il);851 852 ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate);853 cb(gate_sigmoid, "gate_sigmoid", il);854 855 cur = ggml_mul(ctx0, cur, gate_sigmoid);856 cb(cur, "attn_gated", il);857 858 cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);859 cb(cur, "attn_output", il);860 861 return cur;862}863 864ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear(865 llm_graph_input_rs * inp,866 ggml_tensor * cur,867 int il) {868 const auto * mctx_cur = inp->mctx;869 870 const int64_t d_inner = hparams.ssm_d_inner;871 const int64_t n_seqs = ubatch.n_seqs;872 const int64_t head_k_dim = hparams.ssm_d_state;873 const int64_t num_k_heads = hparams.ssm_n_group;874 const int64_t num_v_heads = hparams.ssm_dt_rank;875 const int64_t head_v_dim = hparams.ssm_d_state;876 const int64_t n_seq_tokens = ubatch.n_seq_tokens;877 878 GGML_ASSERT(n_seqs != 0);879 GGML_ASSERT(ubatch.equal_seqs());880 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);881 GGML_ASSERT(head_v_dim * num_v_heads == d_inner);882 883 auto qkvz = build_qkvz(cur, il);884 ggml_tensor * qkv_mixed = qkvz.first;885 ggml_tensor * z = qkvz.second;886 887 ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);888 beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);889 cb(beta, "beta", il);890 891 beta = ggml_sigmoid(ctx0, beta);892 cb(beta, "beta_sigmoid", il);893 894 ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);895 alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);896 cb(alpha, "alpha", il);897 898 ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt);899 ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased);900 cb(alpha_softplus, "a_softplus", il);901 902 ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus903 cb(gate, "gate", il);904 905 gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs);906 907 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);908 ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);909 910 ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d;911 const int64_t conv_kernel_size = conv_kernel->ne[0];912 913 // the channels must match how load_arch_tensors sizes wqkv, not ssm_d_inner914 const int64_t conv_channels = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads;915 916 ggml_tensor * conv_input = build_conv_state_at(inp, conv_states_all, qkv_mixed,917 conv_kernel_size - 1, conv_channels, il);918 919 ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs);920 state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs);921 cb(state, "state_predelta", il);922 923 ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel);924 cb(conv_output_proper, "conv_output_raw", il);925 926 ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper);927 cb(conv_output_silu, "conv_output_silu", il);928 929 ggml_tensor * conv_qkv_mix = conv_output_silu;930 931 int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, conv_channels);932 933 // Extract the convolved Q, K, V from conv_output934 ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs,935 ggml_row_size(conv_qkv_mix->type, head_k_dim),936 nb1_qkv,937 nb1_qkv * n_seq_tokens,938 0);939 940 ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs,941 ggml_row_size(conv_qkv_mix->type, head_k_dim),942 nb1_qkv,943 nb1_qkv * n_seq_tokens,944 head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix));945 946 ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs,947 ggml_row_size(conv_qkv_mix->type, head_v_dim),948 nb1_qkv,949 nb1_qkv * n_seq_tokens,950 ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads));951 952 cb(q_conv, "q_conv", il);953 cb(k_conv, "k_conv", il);954 cb(v_conv, "v_conv", il);955 956 957 const float eps_norm = hparams.f_norm_rms_eps;958 959 q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm);960 k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm);961 962 // repeat to match shapes when head keys != value keys; unneeded with the fused GDN963 if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) {964 GGML_ASSERT(num_v_heads % num_k_heads == 0);965 q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs);966 k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs);967 }968 969 cb(q_conv, "q_conv_predelta", il);970 cb(k_conv, "k_conv_predelta", il);971 cb(v_conv, "v_conv_predelta", il);972 973 ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il);974 975 ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);976 977 // gated normalization, as self.norm(core_attn_out, z) in the reference978 ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);979 980 ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);981 cb(final_output, "final_output", il);982 983 cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);984 cb(cur, "linear_attn_out", il);985 986 cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);987 988 return cur;989}990 991ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) {992 GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr);993 994 ggml_tensor * moe_out =995 build_moe_ffn(cur,996 model.layers[il].ffn_gate_inp,997 model.layers[il].ffn_up_exps,998 model.layers[il].ffn_gate_exps,999 model.layers[il].ffn_down_exps,1000 nullptr,1001 n_expert, n_expert_used,1002 LLM_FFN_SILU, true,1003 hparams.expert_weights_scale,1004 LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,1005 nullptr, model.layers[il].ffn_gate_up_exps,1006 model.layers[il].ffn_up_exps_s,1007 model.layers[il].ffn_gate_exps_s,1008 model.layers[il].ffn_down_exps_s);1009 cb(moe_out, "ffn_moe_out", il);1010 1011 // shared experts, as in the Qwen3Next reference1012 if (model.layers[il].ffn_up_shexp != nullptr) {1013 ggml_tensor * ffn_shexp =1014 build_ffn(cur,1015 model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,1016 model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,1017 model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,1018 NULL,1019 LLM_FFN_SILU, LLM_FFN_PAR, il);1020 cb(ffn_shexp, "ffn_shexp", il);1021 1022 // shared expert has its own sigmoided gate (ffn_gate_inp_shexp, one value per token)1023 ggml_tensor * shared_gate = build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur);1024 cb(shared_gate, "shared_expert_gate", il);1025 1026 shared_gate = ggml_sigmoid(ctx0, shared_gate);1027 cb(shared_gate, "shared_expert_gate_sigmoid", il);1028 1029 ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);1030 cb(ffn_shexp, "ffn_shexp_gated", il);1031 1032 cur = ggml_add(ctx0, moe_out, ffn_shexp);1033 cb(cur, "ffn_out", il);1034 } else {1035 cur = moe_out;1036 }1037 1038 return cur;1039}1040 1041// PLE n-gram hash embedding: each token gathers ple_n_heads rows of a shared table.1042// mixed_n = (t[p]*m[0]) ^ ... ^ (t[p-n+1]*m[n-1]); row = mixed_n % vocab[h] + offset[h]1043// The hash runs host-side because ggml has no int64 and no xor. EOS resets the window.1044 1045class llm_graph_input_ple : public llm_graph_input_i {1046public:1047 llm_graph_input_ple(const llama_model_qwen4exp & pmodel,1048 const llama_kv_cache_context * mctx) : pmodel(pmodel), mctx(mctx) {}1049 virtual ~llm_graph_input_ple() = default;1050 1051 void set_input(const llama_ubatch * ubatch) override;1052 1053 bool can_reuse(const llm_graph_params & params) override {1054 mctx = static_cast<const llama_memory_hybrid_idx_context *>(params.mctx)->get_attn();1055 return rows->ne[0] == (int64_t) pmodel.hparams.ple_n_heads * params.ubatch.n_tokens;1056 }1057 1058 ggml_tensor * rows = nullptr; // I32 [ple_n_heads * n_tokens]1059 1060 const llama_model_qwen4exp & pmodel;1061 1062 // the predecessor tokens live in the attention KV cells (ext.tok)1063 const llama_kv_cache_context * mctx;1064 1065 // scratch, reused across set_input() calls1066 std::vector<llama_token> prev;1067};1068 1069void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) {1070 const auto & hp = pmodel.hparams;1071 1072 // an image arrives as an embd batch, so ubatch->token is null, but every position still needs a row for ggml_get_rows1073 // stand in the image token id that the reference hashes, or EOS if the file has no such key1074 // gemma3n and gemma4 do the same with a hardcoded row 0 of per_layer_token_embd.1075 const llama_token img_tok = hp.ple_image_token_id != 01076 ? (llama_token) hp.ple_image_token_id1077 : (llama_token) hp.ple_eos_token_id;1078 auto tok_of = [&](int64_t k) -> llama_token {1079 return ubatch->token ? ubatch->token[k] : img_tok;1080 };1081 1082 const int64_t n_tokens = ubatch->n_tokens;1083 const int64_t n_gram = hp.ple_ngram_size;1084 const int64_t n_heads = hp.ple_n_heads;1085 const int64_t per_gram = hp.ple_heads_per_ngram;1086 const int64_t eos = hp.ple_eos_token_id;1087 const int64_t n_prev = n_gram - 1;1088 1089 std::vector<int32_t> idx(n_heads * n_tokens);1090 1091 GGML_ASSERT(mctx != nullptr);1092 1093 for (int64_t i = 0; i < n_tokens; ++i) {1094 // the preceding tokens would be ambiguous, see get_prev_tokens()1095 GGML_ASSERT(ubatch->n_seq_id[i] == 1 && "PLE n-gram embeddings do not support tokens shared by multiple sequences");1096 }1097 1098 // predecessors come from the KV cells (ext.tok); apply_ubatch() already stored this ubatch, so its own tokens count too1099 mctx->get_prev_tokens(*ubatch, n_prev, prev);1100 1101 for (int64_t i = 0; i < n_tokens; ++i) {1102 // an EOS in the window resets everything at or before it1103 // a missing predecessor (before the sequence start, or no cached cell) reads as EOS1104 // the EOS of the token itself does not cut its own context, as in the reference1105 std::vector<int64_t> ctx(n_gram);1106 ctx[0] = tok_of(i);1107 bool cut = false;1108 for (int64_t s = 1; s < n_gram; ++s) {1109 // predecessor s positions back; prev[] is oldest-first, missing entries are LLAMA_TOKEN_NULL1110 const llama_token t = cut ? LLAMA_TOKEN_NULL : prev[i*n_prev + (n_prev - s)];1111 cut = cut || t < 0 || t == eos;1112 ctx[s] = cut ? eos : t;1113 }1114 1115 for (int64_t n = 2; n <= n_gram; ++n) {1116 uint64_t mixed = (uint64_t) ctx[0] * hp.ple_layer_multipliers[0];1117 for (int64_t j = 1; j < n; ++j) {1118 mixed ^= (uint64_t) ctx[j] * hp.ple_layer_multipliers[j];1119 }1120 const int64_t base = (n - 2) * per_gram;1121 for (int64_t g = 0; g < per_gram; ++g) {1122 const int64_t h_i = base + g;1123 idx[i * n_heads + h_i] =1124 (int32_t) (mixed % hp.ple_head_vocab_sizes[h_i] + hp.ple_head_offsets[h_i]);1125 }1126 }1127 }1128 1129 ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows));1130}1131 1132// Read a conv history out of its own recurrent row and write the new tail back.1133// The shared build_conv_state cannot do this: qwen4exp has two such rows per layer.1134ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(1135 llm_graph_input_rs * inp,1136 ggml_tensor * conv_states_all,1137 ggml_tensor * x,1138 int64_t state_cols,1139 int64_t channels,1140 int il) {1141 const auto * mctx_cur = inp->mctx;1142 1143 const auto kv_head = mctx_cur->get_head();1144 1145 const int64_t n_seqs = ubatch.n_seqs;1146 const int64_t row_total = conv_states_all->ne[0];1147 1148 // the row is exactly this convolution's state, so the gather is reused as a whole1149 GGML_ASSERT(state_cols * channels == row_total);1150 1151 auto it = rs_rows.find(conv_states_all);1152 if (it == rs_rows.end()) {1153 it = rs_rows.emplace(conv_states_all, build_rs(inp, conv_states_all, row_total, n_seqs)).first;1154 }1155 ggml_tensor * rows = it->second;1156 1157 ggml_tensor * state = ggml_reshape_3d(ctx0, rows, state_cols, channels, n_seqs);1158 cb(state, "conv_state_at", il);1159 1160 ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);1161 1162 // [TAG_RECURRENT_ROLLBACK_SPLITS] keep the last state_cols columns once per rollback slot,1163 // slot s ending s tokens earlier so a rollback of s tokens reads a history that never saw them1164 const size_t row_size = ggml_row_size(conv_states_all->type, row_total);1165 const uint32_t mem_size = mctx_cur->get_size();1166 1167 const int64_t n_slots = (int64_t) cparams.n_rs_seq + 1;1168 1169 for (int64_t slot = 0; slot < n_slots; ++slot) {1170 const int64_t s_idx = std::max<int64_t>(0, conv_input->ne[0] - state_cols - slot);1171 1172 ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,1173 state_cols, channels, n_seqs,1174 conv_input->nb[1], conv_input->nb[2],1175 ggml_row_size(conv_input->type, s_idx));1176 1177 ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,1178 state_cols * channels, n_seqs,1179 conv_states_all->nb[1],1180 (slot * mem_size + kv_head) * row_size);1181 1182 ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));1183 }1184 1185 return conv_input;1186}1187 1188ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple(1189 const llama_memory_hybrid_idx_context * mctx_hyb) {1190 const int64_t n_heads = hparams.ple_n_heads;1191 1192 // the attention cells see every ubatch regardless of the layer types1193 auto ple_inp = std::make_unique<llm_graph_input_ple>(1194 static_cast<const llama_model_qwen4exp &>(model), mctx_hyb->get_attn());1195 1196 ple_inp->rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_heads * n_tokens);1197 ggml_set_input(ple_inp->rows);1198 ggml_tensor * rows = ple_inp->rows;1199 res->add_input(std::move(ple_inp));1200