mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-19 17:24:57 +02:00
llama: qwen4exp text graph with hyper-connections, GDN and MoE
Implements the decode graph for Qwen3.8-Flash-Next: the hyper-connection residual stream, gated delta net layers, the MoE block with its gated shared expert, and dense full attention. The QSA indexer and the PLE n-gram embedding are not wired up yet and land in later commits. Hyper-connections are implemented here rather than shared with deepseek4.cpp. The two formulations agree on the [n_embd, hc, n_tokens] layout and little else: DeepSeek-V4 mixes with a full-rank projection and Sinkhorn-normalises it, whereas this model uses a low-rank down/silu/up sigmoid gate and collapses by a plain mean. Only the ~10 line stream mean is genuinely common, so sharing would mean touching DSV4's hot path and its three fused CUDA ops to reuse very little. What is reused is the substantive part: the LLM_KV_HYPER_CONNECTION_* keys, the n_embd_out_impl wide-residual support already in the loader, and the layout convention. Also allows a checkpoint to carry no PLE layers at all, which makes it possible to bring the graph up and validate it in stages. Validated against vLLM, the only working reference implementation. On a scaled-down model with an init scale large enough to give non-uniform logits, agreement with vLLM sits at the numerical noise floor: llama.cpp f32 against its own bf16 gives 84.3% top-1 agreement over 255 positions, and this graph against vLLM gives 85.1%. The comparison was calibrated by seeding three deliberate bugs (silu instead of sigmoid on the delta net gate, dropping the 1/hc scale in the mix, dropping the 2x in the combine); each drops top-1 to between 0% and 11%, an order of magnitude below the floor.
This commit is contained in:
Daniel Han
@@ -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"])
|
||||
|
||||
+511
-27
@@ -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<uint32_t> 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<llm_graph_context> 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<graph>(*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<ggml_tensor *, ggml_tensor *> 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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user