diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py index 1174edcb7b..0b0ee4ac6d 100644 --- a/conversion/qwen4exp.py +++ b/conversion/qwen4exp.py @@ -65,8 +65,12 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] ) - # ple_layer_ids is 1-based in the HF config + # ple_layer_ids is 1-based in the HF config. An empty list means the + # checkpoint carries no n-gram table at all, so emit no PLE keys either + # rather than keys the loader would then have to treat as optional. ple_layers = [i - 1 for i in hp["ple_layer_ids"]] + if not ple_layers: + return self.gguf_writer.add_ple_layers(ple_layers) self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"]) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index d503e1f1e2..5f04618d90 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -26,33 +26,35 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false); - 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_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id); - ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); - - hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; - hparams.ple_head_dim = hparams.n_embd_per_layer; - GGML_ASSERT(hparams.ple_ngram_size >= 2 && hparams.ple_ngram_size <= LLAMA_MAX_PLE_NGRAM); - GGML_ASSERT(hparams.ple_n_heads > 0 && hparams.ple_n_heads <= LLAMA_MAX_PLE_HEADS); - - ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); - ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, hparams.ple_head_offsets); - ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, hparams.ple_head_vocab_sizes); - + // PLE n-gram hash embeddings. A checkpoint may carry none, in which case the + // whole key group is absent and every PLE field stays zeroed. std::fill(hparams.is_ple_impl.begin(), hparams.is_ple_impl.end(), 0); - { - uint32_t n_ple = 0; - ml.get_arr_n(LLM_KV_PLE_LAYERS, n_ple); - GGML_ASSERT(n_ple > 0 && "qwen4exp needs at least one PLE layer"); + hparams.ple_n_heads = 0; + uint32_t n_ple = 0; + ml.get_arr_n(LLM_KV_PLE_LAYERS, n_ple, false); + if (n_ple > 0) { std::vector ple_layers; ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers); for (uint32_t il : ple_layers) { GGML_ASSERT(il < hparams.n_layer_all); hparams.is_ple_impl[il] = 1; } + + 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_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id); + ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); + + hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; + hparams.ple_head_dim = hparams.n_embd_per_layer; + GGML_ASSERT(hparams.ple_ngram_size >= 2 && hparams.ple_ngram_size <= LLAMA_MAX_PLE_NGRAM); + GGML_ASSERT(hparams.ple_n_heads > 0 && hparams.ple_n_heads <= LLAMA_MAX_PLE_HEADS); + + ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); + ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, hparams.ple_head_offsets); + ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, hparams.ple_head_vocab_sizes); } // linear attention everywhere except every full_attention_interval-th layer @@ -90,12 +92,14 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { } // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back - const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); - const auto * ple_w = ml.get_weight(ple_name.c_str()); - GGML_ASSERT(ple_w != nullptr && "qwen4exp is missing the PLE n-gram table"); - const int64_t ple_rows = ple_w->tensor->ne[1]; - per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), - { hparams.ple_head_dim, ple_rows }, 0); + if (hparams.ple_n_heads > 0) { + const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); + const auto * ple_w = ml.get_weight(ple_name.c_str()); + GGML_ASSERT(ple_w != nullptr && "qwen4exp is missing the PLE n-gram table"); + const int64_t ple_rows = ple_w->tensor->ne[1]; + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), + { hparams.ple_head_dim, ple_rows }, 0); + } for (int il = 0; il < n_layer; ++il) { auto & layer = layers[il]; @@ -167,6 +171,486 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { } std::unique_ptr llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const { - GGML_UNUSED(params); - throw std::runtime_error("qwen4exp: graph not implemented yet"); + return std::make_unique(*this, params); +} + +// Hyper-connections replace every layer norm in this architecture. The state +// carried between blocks is `hc` parallel residual streams, [n_embd, hc, T]. +// Each block reads a single mixed [n_embd, T] view of them and writes its output +// back into all of them through per-stream injection weights. +// +// This is deliberately *not* shared with deepseek4.cpp. The two formulations +// agree on the layout and on nothing else: DSV4 mixes with a full-rank +// projection and Sinkhorn-normalises it, here the mix is a low-rank +// down/silu/up gate and the collapse is a plain mean over streams. See +// plans/playful-prancing-snowflake.md for the comparison that settled this. + +// The mix output is [n_embd, T]; `inject` receives the [hc, T] scatter weights. +ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix( + ggml_tensor * x, + ggml_tensor * w_norm, + ggml_tensor * w_down, + ggml_tensor * w_up, + ggml_tensor * w_inject, + ggml_tensor ** inject, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t nt = x->ne[2]; + + // grouped RMSNorm: ggml_rms_norm reduces over ne[0], which is exactly one + // residual stream, then the [hc_dim] gamma scales all streams at once. + // The gammas were already folded to (1 + w) by the converter. + ggml_tensor * xn = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps); + xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt); + xn = ggml_mul(ctx0, xn, w_norm); + cb(xn, "hc_norm", il); + + ggml_tensor * lo = build_lora_mm(w_down, xn); + lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc)); + ggml_tensor * gate = ggml_sigmoid(ctx0, build_lora_mm(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); + + // collapse the streams by their mean + ggml_tensor * mixed = ggml_view_2d(ctx0, gated, n_embd, nt, + ggml_row_size(gated->type, n_embd) * hc, 0); + mixed = ggml_cont(ctx0, mixed); + for (int64_t c = 1; c < hc; ++c) { + ggml_tensor * s = ggml_view_2d(ctx0, gated, n_embd, nt, + ggml_row_size(gated->type, n_embd) * hc, + ggml_row_size(gated->type, n_embd) * c); + mixed = ggml_add(ctx0, mixed, s); + } + mixed = ggml_scale(ctx0, mixed, 1.0f / (float) hc); + cb(mixed, "hc_mixed", il); + + if (inject) { + *inject = build_lora_mm(w_inject, xn); + cb(*inject, "hc_inject", il); + } + + return mixed; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine( + ggml_tensor * residual, + ggml_tensor * block_out, + ggml_tensor * inject, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = residual->ne[2]; + + // 2*sigmoid keeps the scatter weights centred on 1, so an untrained + // injection matrix reproduces the plain residual add + 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; +} + +llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + const int64_t hc = hparams.dsv4_hc_mult; + + GGML_ASSERT(hparams.n_embd_head_v() == hparams.n_embd_head_k()); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + + auto * inp = build_inp_mem_hybrid(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // 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_init", -1); + + for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = res_hc; + + if (hparams.is_ple(il)) { + res_hc = build_ple(inp->get_recr(), res_hc, il); + } + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_mix(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, il); + + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_layer_attn_linear(inp->get_recr(), cur, il); + } else { + cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il); + } + + res_hc = build_hc_combine(res_hc, cur, inject, il); + + cur = build_hc_mix(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, il); + + cur = build_layer_ffn(cur, il); + cb(cur, "ffn_out", il); + + res_hc = build_hc_combine(res_hc, cur, inject, il); + + // build_cvec expects [n_embd, T], so steer the mean of the streams and + // let the next mix carry it. Tagged "l_last" because that is the layer + // output name imatrix_FIXED.cpp knows how to parse, same as deepseek4. + cb(res_hc, "l_last", il); + } + + // the final mixer is the output norm: there is no separate one + ggml_tensor * cur = build_hc_mix(res_hc, + model.hc_head_norm, model.hc_head_down, model.hc_head_up, + nullptr, nullptr, -1); + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +std::pair llama_model_qwen4exp::graph::build_qkvz( + ggml_tensor * input, + int il) { + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s); + qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs); + cb(qkv_mixed, "linear_attn_qkv_mixed", il); + + ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, input, model.layers[il].wqkv_gate_s); + cb(z, "z", il); + + return { qkv_mixed, z }; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( + ggml_tensor * input, + ggml_tensor * weights, + ggml_tensor * gate, + int layer) { + // the one numerical difference from Qwen3.5's gated delta net: this model + // gates the normalised output with sigmoid, not silu + ggml_tensor * normalized = build_norm(input, weights, nullptr, LLM_NORM_RMS, layer); + ggml_tensor * gated = ggml_sigmoid(ctx0, gate); + + return ggml_mul(ctx0, normalized, gated); +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( + llm_graph_input_attn_kv * inp, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention + + // Qwen3Next uses a single Q projection that outputs query + gate + ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + cb(Qcur_full, "Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0); + cb(Qcur, "Qcur_reshaped", il); + + // Apply Q normalization + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + + ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + cb(Kcur, "Kcur", il); + + ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); + cb(Vcur, "Vcur", il); + + // Apply K normalization + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "gate_reshaped", il); + + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + // Apply IMRoPE + Qcur = ggml_rope_multi( + ctx0, Qcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_multi( + ctx0, Kcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + // Attention computation + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_pregate", il); + + ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate); + cb(gate_sigmoid, "gate_sigmoid", il); + + cur = ggml_mul(ctx0, cur, gate_sigmoid); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_output", il); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il) { + const auto * mctx_cur = inp->mctx; + + const int64_t d_inner = hparams.ssm_d_inner; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t num_k_heads = hparams.ssm_n_group; + const int64_t num_v_heads = hparams.ssm_dt_rank; + const int64_t head_v_dim = d_inner / num_v_heads; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + // Input projections + auto qkvz = build_qkvz(cur, il); + ggml_tensor * qkv_mixed = qkvz.first; + ggml_tensor * z = qkvz.second; + + ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s); + beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs); + cb(beta, "beta", il); + + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "beta_sigmoid", il); + + ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s); + alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs); + cb(alpha, "alpha", il); + + ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt); + ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased); + cb(alpha_softplus, "a_softplus", il); + + ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus + cb(gate, "gate", il); + + gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + + ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; + const int64_t conv_kernel_size = conv_kernel->ne[0]; + const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state; + + ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il); + + ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + cb(state, "state_predelta", il); + + ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + cb(conv_output_proper, "conv_output_raw", il); + + ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); + cb(conv_output_silu, "conv_output_silu", il); + + ggml_tensor * conv_qkv_mix = conv_output_silu; + + // Calculate the total conv dimension + int64_t qkv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads; + int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim); + + // Extract the convolved Q, K, V from conv_output + ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + 0); + + ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); + + ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_v_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads)); + + cb(q_conv, "q_conv", il); + cb(k_conv, "k_conv", il); + cb(v_conv, "v_conv", il); + + const float eps_norm = hparams.f_norm_rms_eps; + + q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); + k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + + //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); + //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); + //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); + + // if head keys and value keys are different, repeat to force tensors into matching shapes + // note: need explicit repeat only if we are not using the fused GDN. + if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { + GGML_ASSERT(num_v_heads % num_k_heads == 0); + q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + } + + cb(q_conv, "q_conv_predelta", il); + cb(k_conv, "k_conv_predelta", il); + cb(v_conv, "v_conv_predelta", il); + + ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + + // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim] + ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); + + // Apply gated normalization: self.norm(core_attn_out, z) + ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il); + + // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim] + ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs); + cb(final_output, "final_output", il); + + // Output projection + cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s); + cb(cur, "linear_attn_out", il); + + // Reshape back to original dimensions + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) { + // Check if this is an MoE layer + GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr); + + ggml_tensor * moe_out = + build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il, + nullptr, model.layers[il].ffn_gate_up_exps, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); + cb(moe_out, "ffn_moe_out", il); + + // Add shared experts if present - following Qwen3Next reference implementation + if (model.layers[il].ffn_up_shexp != nullptr) { + ggml_tensor * ffn_shexp = + build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, + model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, + model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + // Apply shared expert gating as in the reference implementation + // The shared expert has its own gate that is sigmoided + // Note: ffn_gate_inp_shexp is the shared expert gate (outputs 1 value per token) + ggml_tensor * shared_gate = build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur); + cb(shared_gate, "shared_expert_gate", il); + + // Apply sigmoid to the gate + shared_gate = ggml_sigmoid(ctx0, shared_gate); + cb(shared_gate, "shared_expert_gate_sigmoid", il); + + + // Apply the gate to the shared expert output + ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate); + cb(ffn_shexp, "ffn_shexp_gated", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } else { + cur = moe_out; + } + + return cur; +} + +// The PLE n-gram hash embedding lands in the follow-up commit; a checkpoint +// without ple_layer_ids never reaches this path. +ggml_tensor * llama_model_qwen4exp::graph::build_ple( + llm_graph_input_rs * inp, + ggml_tensor * hidden, + int il) { + GGML_UNUSED(inp); + GGML_UNUSED(hidden); + GGML_UNUSED(il); + throw std::runtime_error("qwen4exp: PLE not implemented yet"); }