Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-07 20:47:30 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
eb46b8cf0a2efb5059cd77ac9601bd5d57b4b1e3
llama.cpp/src/models/qwen4exp.cpp
T

1214 lines
54 KiB
C++
Raw Blame History

#include "models.h"
#include "llama-impl.h"
#include "llama-memory-hybrid-idx.h"
#include "llama-memory-recurrent.h"
#include <algorithm>
#include <cinttypes>
void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);
ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);
ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
GGML_ASSERT(hparams.ssm_d_conv > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 &&
hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0);
// HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank)
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult);
ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0);
hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
GGML_ASSERT(hparams.indexer_n_head > 0
&& hparams.indexer_head_size > 0
&& hparams.indexer_top_k > 0);
ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false);
// PLE n-gram hash embeddings; if the key group is absent every field stays zero
hparams.is_ple_impl.reset();
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);
GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer");
for (uint32_t il : ple_layers) {
if (il >= hparams.n_layer_all) {
throw std::runtime_error(format("PLE layer %u is out of range", il));
}
hparams.is_ple_impl.set(il);
}
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);
// optional: files written before this key fall back to the EOS token
ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false);
ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);
GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0);
hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
hparams.ple_head_dim = hparams.n_embd_per_layer;
if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) {
throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size));
}
if (hparams.ple_n_heads == 0 || hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS) {
throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads));
}
ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers);
// the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses
std::array<uint64_t, LLAMA_MAX_PLE_HEADS> head_offsets = {};
std::array<uint64_t, LLAMA_MAX_PLE_HEADS> head_vocab_sizes = {};
ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, head_offsets);
ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, head_vocab_sizes);
for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {
if (head_vocab_sizes[h] == 0 ||
head_offsets[h] > INT32_MAX ||
head_vocab_sizes[h] > INT32_MAX ||
head_offsets[h] + head_vocab_sizes[h] > INT32_MAX) {
throw std::runtime_error(format("PLE head %u range does not fit the int32 row index", h));
}
hparams.ple_head_offsets[h] = (uint32_t) head_offsets[h];
hparams.ple_head_vocab_sizes[h] = (uint32_t) head_vocab_sizes[h];
}
}
// linear attention everywhere except every full_attention_interval-th layer
if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
GGML_ASSERT(full_attn_interval > 0);
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
}
}
switch (hparams.n_layer()) {
case 48: type = LLM_TYPE_A3B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t hc_dim = hc * n_embd;
const int64_t hc_lr = hparams.hc_low_rank;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
// there is no output_norm: the final hyper-connection mixer carries it
hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, 0);
hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, 0);
hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
}
// flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back
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.require_weight(ple_name.c_str());
const int64_t ple_rows = ple_w.tensor->ne[1];
// sanity check
for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {
if ((int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h] > ple_rows) {
throw std::runtime_error(format("PLE head %u range exceeds the %" PRId64 " table rows", h, ple_rows));
}
}
per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"),
{ hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY);
}
for (int il = 0; il < n_layer; ++il) {
auto & layer = layers[il];
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;
const int64_t head_k_dim = hparams.ssm_d_state;
const int64_t head_v_dim = hparams.ssm_d_state;
const int64_t n_k_heads = hparams.ssm_n_group;
const int64_t n_v_heads = hparams.ssm_dt_rank;
const int64_t key_dim = head_k_dim * n_k_heads;
const int64_t value_dim = head_v_dim * n_v_heads;
const int64_t conv_dim = key_dim * 2 + value_dim;
// two HC modules per layer: before the token mixer, before the MoE
layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, 0);
layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, 0);
layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, 0);
layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0);
layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, 0);
layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, 0);
layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, 0);
layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0);
if (!hparams.is_recr(il)) {
// full attention: wq holds [q|gate] interleaved per head
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
const int64_t idx_dim = hparams.indexer_head_size;
layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0);
layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0);
layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, 0);
layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, 0);
} else {
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, 0);
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0);
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, 0);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, 0);
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0);
layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0);
}
if (hparams.is_ple(il)) {
layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0);
layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0);
layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, 0);
layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0);
layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, 0);
layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0);
}
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);
create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// Hyper-connections keep hc parallel residual streams [n_embd, hc, T] in place of layer norms.
// Returns the mixed [n_embd, T] stream; `inject` gets 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: reduce over one stream, then scale all streams with the [hc_dim] gamma
// the converter folded each gamma to (1 + w)
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 centres the scatter weights on 1, so a zero injection is a 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);
ggml_build_forward_expand(gf, inpL);
auto * inp = build_inp_mem_hybrid();
// qwen4exp always builds llama_memory_hybrid_idx, so this downcast is safe
// the indexer cache inside it is absent when the GGUF has no indexer tensors
const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(inp->mctx);
const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
if (mctx_idx) {
GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() &&
"the indexer cache must track the attention cache cell for cell");
}
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
ggml_tensor * ple_emb = nullptr;
if (hparams.ple_n_heads > 0) {
ple_emb = build_inp_ple(mctx_hyb);
// make sure ple_emb and build_inp_embd are in the same graph split
ggml_build_forward_expand(gf, ple_emb);
}
// 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(), ple_emb, 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(), mctx_hyb, cur, inp_pos, sections, il);
}
if (il == n_layer - 1 && inp_out_ids) {
// everything below is per token, so drop the rows that produce no output
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inject = ggml_get_rows(ctx0, inject, inp_out_ids);
res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]);
res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids);
res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]);
}
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);
// "l_last" is the layer output name that build_cvec and imatrix look for
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);
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 GDN: sigmoid output gate, 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);
}
// QSA attends to a budget of whole blocks of compress_ratio tokens, plus the incomplete tail
// one mean-pooled indexer key scores each block; set_input resolves the cache layout
class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i {
public:
llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) :
mctx(mctx), ratio(ratio), blk_bias(blk_bias) {}
virtual ~llm_graph_input_qsa() = default;
void set_input(const llama_ubatch * ubatch) override {
mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);
mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias);
}
bool can_reuse(const llm_graph_params & params) override {
mctx = static_cast<const llama_memory_hybrid_idx_context *>(params.mctx);
const auto * idx = mctx->get_idx();
if (idx == nullptr) {
return false;
}
const int64_t n_kv = idx->get_n_kv();
const int64_t n_stream = mctx->get_n_stream();
const int64_t n_blocks = (n_kv + ratio - 1)/ratio;
bool res = true;
res &= params.ubatch.n_tokens % n_stream == 0;
res &= k_idxs->ne[0] == params.ubatch.n_tokens;
res &= cell_blk->ne[0] == n_kv;
res &= cell_blk->ne[1] == n_stream;
res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks;
res &= blk_pos->ne[0] == 4*n_blocks*n_stream;
res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv);
res &= bias->ne[1] == params.ubatch.n_tokens/n_stream;
return res;
}
// per stream: a cell index names a different token in each stream
ggml_tensor * k_idxs = nullptr; // I32 [n_tokens]
ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream]
ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream]
ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream]
ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream]
const llama_memory_hybrid_idx_context * mctx;
const uint32_t ratio;
// the per-cell half of the bias is the attention mask, so only the per-block half is uploaded
const bool blk_bias;
};
ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
const llama_memory_hybrid_idx_context * mctx_hyb,
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * kq_mask,
int * sections,
int il) {
const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx();
const int64_t idx_dim = hparams.indexer_head_size;
const int64_t n_idx_h = hparams.indexer_n_head;
const int64_t r = hparams.dsv4_compress_ratios[il];
const int64_t n_kv = mctx_idx->get_n_kv();
GGML_ASSERT(r > 0);
const int64_t n_blocks = (n_kv + r - 1)/r;
// build_attn_qsa and the KQ mask need the tokens to divide evenly across the streams
const int64_t n_stream = mctx_hyb->get_n_stream();
GGML_ASSERT(n_tokens % n_stream == 0);
const int64_t n_tps = n_tokens/n_stream;
// only the "which block is visible" half of the bias varies per block
// the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells
// alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out
// the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there
const bool blk_bias = kq_mask != nullptr &&
kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream &&
cparams.causal_attn && !hparams.use_alibi;
// nothing above depends on the layer, so the layers sharing a ratio share one input set
llm_graph_input_qsa * inp = nullptr;
const auto it = qsa_inps.find((uint32_t) r);
if (it != qsa_inps.end()) {
inp = it->second;
} else {
auto qsa = std::make_unique<llm_graph_input_qsa>(mctx_hyb, (uint32_t) r, blk_bias);
qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch);
qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream);
qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream);
qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream);
qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream);
ggml_set_input(qsa->cell_blk);
ggml_set_input(qsa->blk_cells);
ggml_set_input(qsa->blk_pos);
ggml_set_input(qsa->bias);
inp = qsa.get();
res->add_input(std::move(qsa));
qsa_inps.emplace((uint32_t) r, inp);
}
// cached indexer keys are raw: pooling precedes norm and rotation, so apply neither
ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur);
k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens);
cb(k_raw, "indexer_k_raw", il);
ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il));
// one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream].
ggml_tensor * k_all = mctx_idx->get_k(ctx0, il);
k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0);
// gathers per stream: blk_cells row s indexes stream s's own cells
ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells);
members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream);
// mean over the block members; r is small, so summing slices beats a transpose plus sum_rows
ggml_tensor * pooled = nullptr;
for (int64_t i = 0; i < r; ++i) {
ggml_tensor * slice = ggml_cont(ctx0,
ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream,
members->nb[2], members->nb[3], i*members->nb[1]));
pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice;
}
pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r);
cb(pooled, "indexer_k_pooled", il);
// rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after.
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream);
pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr,
n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream);
cb(pooled, "indexer_k", il);
ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur);
q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens);
q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il);
q = ggml_rope_multi(ctx0, q, inp_pos, nullptr,
n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(q, "indexer_q", il);
// rectify each head dot product before the sum, as in the DeepSeek lightning indexer
// mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s
ggml_tensor * score = ggml_mul_mat(ctx0, pooled,
ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream));
score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);
score = ggml_relu(ctx0, score);
score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
score = ggml_sum_rows(ctx0, score);
score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream);
cb(score, "indexer_score", il);
// one value per block, so it is cheaper to bias here than after the cells are expanded
if (blk_bias) {
score = ggml_add(ctx0, score, inp->bias);
}
// every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary
ggml_tensor * expanded = ggml_get_rows(ctx0,
ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk);
expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3));
if (blk_bias) {
// flash attention keeps the mask in f16; the scores are f32
ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32);
expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream));
} else {
expanded = ggml_add(ctx0, expanded, inp->bias);
}
cb(expanded, "indexer_score_tokens", il);
// the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail
const int64_t width = std::min<int64_t>(n_kv, (int64_t) hparams.indexer_top_k + r - 1);
ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width));
// build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask.
top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream);
cb(top_k, "indexer_top_k", il);
return top_k;
}
// Dense GQA self-attention restricted to the cells that top_k names.
// The mask build below copies the MLA sparse path in llm_graph_context::build_attn.
ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
llm_graph_input_attn_kv * inp,
ggml_tensor * q_cur,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
ggml_tensor * top_k,
float kq_scale,
int il) {
// rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer
// has already scored with its own query in build_qsa_top_k, so top_k is unaffected.
if (inp->self_k_rot) {
q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);
k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);
}
if (inp->self_v_rot) {
v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
}
// these nodes are added to the graph together so that they are not reordered
// by doing so, the number of splits in the graph is reduced
// expand k later to enable rope fusion which directly writes into k-v cache
ggml_build_forward_expand(gf, q_cur);
ggml_build_forward_expand(gf, v_cur);
ggml_build_forward_expand(gf, k_cur);
const auto * mctx_cur = inp->mctx;
// store to KV cache
{
const auto & k_idxs = inp->get_k_idxs();
const auto & v_idxs = inp->get_v_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}
ggml_tensor * kq_mask = inp->get_kq_mask();
// prepare new kq mask - starts filled with -INFINITY
ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
// reshape KQ mask into tensor with rows of size 1:
// [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
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);
// reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
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);
// prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
// this will be our source of zero values for unmasking top k mask elements
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]);
zeros = ggml_fill(ctx0, zeros, 0.0f);
// modify KQ mask by unmasking elements that are in top_k indices
// 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])
ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
// reshape to restore the original shape of KQ mask:
// [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
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);
// combine with the original kq mask
kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il);
cb(cur, "kqv_out", il);
// the rotation is its own inverse, so undo it on the value side of the output
if (inp->self_v_rot) {
cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
}
return cur;
}
ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn(
llm_graph_input_attn_kv * inp,
const llama_memory_hybrid_idx_context * mctx_hyb,
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());
// indexer reads the same block input as q/k/v; no cache or no ratio means dense
const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0;
ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr;
// 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);
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);
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);
const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
if (top_k) {
cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il);
} else {
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 = hparams.ssm_d_state;
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);
GGML_ASSERT(head_v_dim * num_v_heads == d_inner);
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];
// the channels must match how load_arch_tensors sizes wqkv, not ssm_d_inner
const int64_t conv_channels = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads;
ggml_tensor * conv_input = build_conv_state_at(inp, conv_states_all, qkv_mixed,
conv_kernel_size - 1, 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;
int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, conv_channels);
// 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);
// repeat to match shapes when head keys != value keys; unneeded with 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);
ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);
// gated normalization, as self.norm(core_attn_out, z) in the reference
ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);
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);
cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);
cb(cur, "linear_attn_out", il);
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) {
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);
// shared experts, as in the Qwen3Next reference
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);
// shared expert has its own sigmoided gate (ffn_gate_inp_shexp, one 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);
shared_gate = ggml_sigmoid(ctx0, shared_gate);
cb(shared_gate, "shared_expert_gate_sigmoid", il);
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;
}
// PLE n-gram hash embedding: each token gathers ple_n_heads rows of a shared table.
// mixed_n = (t[p]*m[0]) ^ ... ^ (t[p-n+1]*m[n-1]); row = mixed_n % vocab[h] + offset[h]
// The hash runs host-side because ggml has no int64 and no xor. EOS resets the window.
class llm_graph_input_ple : public llm_graph_input_i {
public:
llm_graph_input_ple(const llama_model_qwen4exp & pmodel,
const llama_kv_cache_context * mctx) : pmodel(pmodel), mctx(mctx) {}
virtual ~llm_graph_input_ple() = default;
void set_input(const llama_ubatch * ubatch) override;
bool can_reuse(const llm_graph_params & params) override {
mctx = static_cast<const llama_memory_hybrid_idx_context *>(params.mctx)->get_attn();
return rows->ne[0] == (int64_t) pmodel.hparams.ple_n_heads * params.ubatch.n_tokens;
}
ggml_tensor * rows = nullptr; // I32 [ple_n_heads * n_tokens]
const llama_model_qwen4exp & pmodel;
// the predecessor tokens live in the attention KV cells (ext.tok)
const llama_kv_cache_context * mctx;
// scratch, reused across set_input() calls
std::vector<llama_token> prev;
};
void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) {
const auto & hp = pmodel.hparams;
// an image arrives as an embd batch, so ubatch->token is null, but every position still needs a row for ggml_get_rows
// stand in the image token id that the reference hashes, or EOS if the file has no such key
// gemma3n and gemma4 do the same with a hardcoded row 0 of per_layer_token_embd.
const llama_token img_tok = hp.ple_image_token_id != 0
? (llama_token) hp.ple_image_token_id
: (llama_token) hp.ple_eos_token_id;
auto tok_of = [&](int64_t k) -> llama_token {
return ubatch->token ? ubatch->token[k] : img_tok;
};
const int64_t n_tokens = ubatch->n_tokens;
const int64_t n_gram = hp.ple_ngram_size;
const int64_t n_heads = hp.ple_n_heads;
const int64_t per_gram = hp.ple_heads_per_ngram;
const int64_t eos = hp.ple_eos_token_id;
const int64_t n_prev = n_gram - 1;
std::vector<int32_t> idx(n_heads * n_tokens);
GGML_ASSERT(mctx != nullptr);
for (int64_t i = 0; i < n_tokens; ++i) {
// the preceding tokens would be ambiguous, see get_prev_tokens()
GGML_ASSERT(ubatch->n_seq_id[i] == 1 && "PLE n-gram embeddings do not support tokens shared by multiple sequences");
}
// predecessors come from the KV cells (ext.tok); apply_ubatch() already stored this ubatch, so its own tokens count too
mctx->get_prev_tokens(*ubatch, n_prev, prev);
for (int64_t i = 0; i < n_tokens; ++i) {
// an EOS in the window resets everything at or before it
// a missing predecessor (before the sequence start, or no cached cell) reads as EOS
// the EOS of the token itself does not cut its own context, as in the reference
std::vector<int64_t> ctx(n_gram);
ctx[0] = tok_of(i);
bool cut = false;
for (int64_t s = 1; s < n_gram; ++s) {
// predecessor s positions back; prev[] is oldest-first, missing entries are LLAMA_TOKEN_NULL
const llama_token t = cut ? LLAMA_TOKEN_NULL : prev[i*n_prev + (n_prev - s)];
cut = cut || t < 0 || t == eos;
ctx[s] = cut ? eos : t;
}
for (int64_t n = 2; n <= n_gram; ++n) {
uint64_t mixed = (uint64_t) ctx[0] * hp.ple_layer_multipliers[0];
for (int64_t j = 1; j < n; ++j) {
mixed ^= (uint64_t) ctx[j] * hp.ple_layer_multipliers[j];
}
const int64_t base = (n - 2) * per_gram;
for (int64_t g = 0; g < per_gram; ++g) {
const int64_t h_i = base + g;
idx[i * n_heads + h_i] =
(int32_t) (mixed % hp.ple_head_vocab_sizes[h_i] + hp.ple_head_offsets[h_i]);
}
}
}
ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows));
}
// Read a conv history out of its own recurrent row and write the new tail back.
// The shared build_conv_state cannot do this: qwen4exp has two such rows per layer.
ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
llm_graph_input_rs * inp,
ggml_tensor * conv_states_all,
ggml_tensor * x,
int64_t state_cols,
int64_t channels,
int il) {
const auto * mctx_cur = inp->mctx;
const auto kv_head = mctx_cur->get_head();
const int64_t n_seqs = ubatch.n_seqs;
const int64_t row_total = conv_states_all->ne[0];
// the row is exactly this convolution's state, so the gather is reused as a whole
GGML_ASSERT(state_cols * channels == row_total);
auto it = rs_rows.find(conv_states_all);
if (it == rs_rows.end()) {
it = rs_rows.emplace(conv_states_all, build_rs(inp, conv_states_all, row_total, n_seqs)).first;
}
ggml_tensor * rows = it->second;
ggml_tensor * state = ggml_reshape_3d(ctx0, rows, state_cols, channels, n_seqs);
cb(state, "conv_state_at", il);
ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);
// keep the last state_cols columns for the next ubatch
const size_t row_size = ggml_row_size(conv_states_all->type, row_total);
ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
state_cols, channels, n_seqs,
conv_input->nb[1], conv_input->nb[2],
ggml_row_size(conv_input->type, conv_input->ne[0] - state_cols));
ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
state_cols * channels, n_seqs,
conv_states_all->nb[1],
kv_head * row_size);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));
return conv_input;
}
ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple(
const llama_memory_hybrid_idx_context * mctx_hyb) {
const int64_t n_heads = hparams.ple_n_heads;
// the attention cells see every ubatch regardless of the layer types
auto ple_inp = std::make_unique<llm_graph_input_ple>(
static_cast<const llama_model_qwen4exp &>(model), mctx_hyb->get_attn());
ple_inp->rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_heads * n_tokens);
ggml_set_input(ple_inp->rows);
ggml_tensor * rows = ple_inp->rows;
res->add_input(std::move(ple_inp));
// gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does
ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows);
emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens);
cb(emb, "ple_embd", -1);
return emb;
}
ggml_tensor * llama_model_qwen4exp::graph::build_ple(
llm_graph_input_rs * inp,
ggml_tensor * emb,
ggml_tensor * hidden,
int il) {
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t hc_dim = hc * n_embd;
ggml_tensor * key = build_lora_mm(model.layers[il].ple_key, emb);
ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb);
// both norms group over one hc stream, with a weight over the whole hc*n_embd layout
auto grouped_norm = [&](ggml_tensor * x, ggml_tensor * w) {
ggml_tensor * t = ggml_reshape_3d(ctx0, x, n_embd, hc, n_tokens);
t = ggml_rms_norm(ctx0, t, hparams.f_norm_rms_eps);
t = ggml_reshape_2d(ctx0, t, hc_dim, n_tokens);
t = ggml_mul(ctx0, t, w);
return ggml_reshape_3d(ctx0, t, n_embd, hc, n_tokens);
};
key = grouped_norm(key, model.layers[il].ple_norm_key);
ggml_tensor * query = grouped_norm(hidden, model.layers[il].ple_norm_query);
// per-stream dot product, then a signed square root before the sigmoid
ggml_tensor * s = ggml_sum_rows(ctx0, ggml_mul(ctx0, key, query));
s = ggml_scale(ctx0, s, 1.0f / sqrtf((float) n_embd));
ggml_tensor * mag = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, s), 1e-6f, INFINITY));
ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, ggml_sgn(ctx0, s), mag));
cb(gate, "ple_gate", il);
// [n_embd, 1, T] value broadcast across the hc streams, scaled by the gate
ggml_tensor * v3 = ggml_reshape_3d(ctx0, value, n_embd, 1, n_tokens);
v3 = ggml_repeat_4d(ctx0, v3, n_embd, hc, n_tokens, 1);
ggml_tensor * gated = ggml_mul(ctx0, v3, gate);
cb(gated, "ple_gated_value", il);
ggml_tensor * normalized = grouped_norm(
ggml_reshape_2d(ctx0, gated, hc_dim, n_tokens),
model.layers[il].ple_norm_conv);
normalized = ggml_reshape_2d(ctx0, normalized, hc_dim, n_tokens);
// depthwise causal conv, dilated by the n-gram size, as a sum of shifted copies
// ggml_conv_1d_dw is documented as unreliable:
// out[c, t] = sum_k w[k, c] * x[c, t - (K-1-k)*dilation]
// The history of the earlier ubatches is prepended, so a chunked prefill matches a single-shot one.
const int64_t kern = hparams.ple_conv_kernel;
const int64_t dil = hparams.ple_ngram_size;
const int64_t hist = (kern - 1) * dil;
// the conv history is per sequence, so the input carries the sequence axis too
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
// [hist + n_seq_tokens, hc_dim, n_seqs], tokens on ne[0]
ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_p_l(il),
ggml_reshape_3d(ctx0, normalized, hc_dim, n_seq_tokens, n_seqs),
hist, hc_dim, il);
ggml_tensor * conv_out = nullptr;
for (int64_t k = 0; k < kern; ++k) {
// tap k reads (kern-1-k)*dilation positions back
const int64_t start = hist - (kern - 1 - k) * dil;
ggml_tensor * shifted = ggml_cont(ctx0,
ggml_transpose(ctx0,
ggml_view_3d(ctx0, padded, n_seq_tokens, hc_dim, n_seqs,
padded->nb[1], padded->nb[2],
ggml_row_size(padded->type, start))));
// column k of the [kern, hc_dim] kernel is one weight per channel
ggml_tensor * wk = ggml_cont(ctx0,
ggml_view_2d(ctx0, model.layers[il].ple_conv1d, 1, hc_dim,
model.layers[il].ple_conv1d->nb[1],
k * model.layers[il].ple_conv1d->nb[0]));
// this kernel keeps the file type, so cast it before it multiplies an f32 activation
wk = ggml_reshape_1d(ctx0, wk, hc_dim);
if (wk->type != GGML_TYPE_F32) {
wk = ggml_cast(ctx0, wk, GGML_TYPE_F32);
}
ggml_tensor * term = ggml_mul(ctx0, shifted, wk);
conv_out = conv_out ? ggml_add(ctx0, conv_out, term) : term;
}
conv_out = ggml_silu(ctx0, conv_out);
conv_out = ggml_reshape_3d(ctx0, ggml_cont(ctx0, conv_out), n_embd, hc, n_tokens);
cb(conv_out, "ple_conv_out", il);
return ggml_add(ctx0, hidden, ggml_add(ctx0, gated, conv_out));
}
Reference in New Issue View Git Blame Copy Permalink