Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2#include "llama-kv-cache-msa.h"3#include <cmath>4#include <vector>5#include <cstdint>6 7// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with8// DeepSeek-V3 leading-dense + routed/shared experts (sigmoid gating, routed scaling),9// swigluoai activation, and MiniMax Sparse Attention (MSA). MTP is not in released model weights.10// MSA blocks are defined over token positions. The graph translates between position space (block11// selection) and cell space (K/V/indexer storage) via per-ubatch pos<->cell maps populated from llama_kv_cells12 13void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) {14 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);15 ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);16 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);17 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);18 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);19 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);20 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);21 ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);22 ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);23 ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);24 ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size);25 ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);26 msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks };27 28 GGML_ASSERT(hparams.indexer_block_size > 0); // avoid div by zero29 30 switch (hparams.n_layer()) {31 case 60: type = LLM_TYPE_428B_A23B; break;32 default: type = LLM_TYPE_UNKNOWN;33 }34}35 36void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) {37 LLAMA_LOAD_LOCALS;38 const int64_t n_expert_shared = hparams.n_expert_shared;39 const int64_t n_ff_exp = hparams.n_ff_exp();40 41 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);42 43 // output44 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);45 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);46 47 for (int i = 0; i < n_layer; ++i) {48 auto & layer = layers[i];49 50 create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0);51 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);52 53 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);54 // per-head QK-norm: a single head_dim vector applied to every head55 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);56 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);57 58 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);59 60 if (i < (int) hparams.n_layer_dense_lead) {61 // leading dense layers62 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);63 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);64 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);65 } else {66 // routed experts67 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);68 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);69 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);70 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);71 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);72 73 // shared expert74 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);75 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, 0);76 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);77 78 // indexer79 layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, 0);80 layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, 0);81 layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, 0);82 layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, 0);83 }84 }85}86 87std::unique_ptr<llm_graph_context> llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const {88 return std::make_unique<graph>(*this, params);89}90 91class llm_graph_input_msa : public llm_graph_input_i {92public:93 llm_graph_input_msa(const llama_kv_cache_msa_context * mctx, int blk, int local) :94 mctx(mctx), blk(blk), local(local) {}95 96 void set_input(const llama_ubatch * ubatch) override {97 if (pos_slot_i) { mctx->set_input_pos_slot(pos_slot_i, ubatch); }98 if (pos_slot_f) { mctx->set_input_pos_slot(pos_slot_f, ubatch); }99 if (cell_blk) { mctx->set_input_cell_pos(cell_blk, ubatch, blk); }100 if (pos_mask) { mctx->set_input_pos_mask(pos_mask, ubatch); }101 102 // local-force bias over position blocks103 if (bias && ubatch->pos) {104 const int64_t n_tokens = ubatch->n_tokens;105 const int64_t nblk = bias->ne[0];106 std::vector<float> data((size_t) nblk * n_tokens, 0.0f);107 for (int64_t i = 0; i < n_tokens; ++i) {108 const int64_t L = ubatch->pos[i] / blk;109 for (int l = 0; l < local && L - l >= 0; ++l) {110 if (L - l < nblk) {111 data[(size_t) i * nblk + (L - l)] = 1e30f;112 }113 }114 }115 ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float));116 }117 }118 119 // valid as long as the tensor dims still match the new ubatch/cache window and the120 // ubatch is in the same regime (decode graphs have pos_slot_f, batch graphs cell_blk)121 bool can_reuse(const llm_graph_params & params) override {122 const auto * mctx_new = static_cast<const llama_kv_cache_msa_context *>(params.mctx);123 124 this->mctx = mctx_new;125 126 const int64_t n_ps = GGML_PAD((int64_t) mctx_new->get_n_pos(), blk);127 const int64_t ns = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq;128 129 const bool decode = params.ubatch.n_tokens == ns; // one token per stream130 131 bool res = true;132 133 res &= bias->ne[0] * blk == n_ps;134 res &= bias->ne[1] == params.ubatch.n_tokens;135 136 res &= pos_mask->ne[0] == n_ps;137 res &= pos_mask->ne[1] == params.ubatch.n_tokens;138 139 res &= pos_slot_i->ne[0] == n_ps;140 res &= pos_slot_i->ne[1] == ns;141 142 res &= decode == (pos_slot_f != nullptr);143 res &= decode == (cell_blk == nullptr);144 145 if (pos_slot_f) {146 res &= pos_slot_f->ne[0] == n_ps;147 res &= pos_slot_f->ne[1] == ns;148 }149 150 if (cell_blk) {151 res &= cell_blk->ne[0] == (int64_t) mctx_new->get_base()->get_n_kv();152 res &= cell_blk->ne[1] == ns;153 }154 155 return res;156 }157 158 ggml_tensor * bias = nullptr; // F32 [nblk, n_tokens] local-force bias (position blocks)159 ggml_tensor * pos_mask = nullptr; // F32 [n_ps, n_tokens] 0/-inf visibility, by position160 ggml_tensor * pos_slot_i = nullptr; // I32 [n_ps, ns] pos -> cell (get_rows index)161 ggml_tensor * pos_slot_f = nullptr; // F32 [n_ps, ns] pos -> cell (gatherable values, decode)162 ggml_tensor * cell_blk = nullptr; // I32 [n_kv, ns] cell -> position block (batch)163 164 const llama_kv_cache_msa_context * mctx;165 166 int blk;167 int local;168};169 170// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3])171ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa(172 ggml_tensor * q_cur, // [D, HQ, T]173 ggml_tensor * k, // [D, n_keys, 1, C]174 ggml_tensor * v, // [D, n_keys, 1, C]175 ggml_tensor * mask, // [n_keys, R, 1, C] f16, contiguous176 int64_t Gp, float kq_scale, int il) const {177 178 const int64_t D = q_cur->ne[0];179 const int64_t HQ = q_cur->ne[1];180 const int64_t T = q_cur->ne[2];181 const int64_t C = k->ne[3];182 const int64_t R = HQ*T/(Gp*C);183 GGML_ASSERT(Gp*C*R == HQ*T);184 GGML_ASSERT(mask->type == GGML_TYPE_F16);185 186 // [D, HQ, T] -> [D, Gp, C, R] -> [D, R, Gp, C]187 // batch (C=HKV, R=T): channel = group188 // decode (C=HKV*ns, R=1): channel = (group, stream), group innermost189 ggml_tensor * q = ggml_reshape_4d(ctx0, q_cur, D, Gp, C, R);190 q = ggml_permute(ctx0, q, 0, 2, 3, 1);191 192 ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale,193 hparams.f_max_alibi_bias, 0.0f);194 ggml_prec_set_acc(o, GGML_PREC_F32);195 cb(o, "msa_fattn", il);196 197 // [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T]198 o = ggml_permute(ctx0, o, 0, 1, 3, 2);199 if (!ggml_is_contiguous(o)) {200 o = ggml_cont(ctx0, o); // no-op layout at decode (R == 1), copy at batch201 }202 return ggml_reshape_2d(ctx0, o, D*HQ, T);203}204 205llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {206 const int64_t n_embd_head = hparams.n_embd_head_v();207 const auto & mm = static_cast<const llama_model_minimax_m3 &>(model);208 209 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());210 // partial rotary: head_dim != n_rot, so don't assert n_embd_head == n_rot211 212 ggml_tensor * cur;213 ggml_tensor * inpL;214 215 inpL = build_inp_embd(model.tok_embd);216 217 ggml_tensor * inp_pos = build_inp_pos();218 219 // ==========================================220 // TODO: avoid such kind of complexity in the model graphs221 222 // MSA calls ggml_flash_attn_ext directly and assumes the non-transposed V layout that223 // llama.cpp only provides when flash attention is enabled. Block selection is anchored224 // to absolute KV cache slots, which equal positions only for append-only per-stream225 // caches either a single sequence, or multiple sequences with kv_unified == false (each226 // stream then has its own slot space). A unified cache with multiple sequences227 // interleaves slots and would silently break block anchoring so it falls back to dense.228 const bool fa_on = cparams.flash_attn;229 const bool streams_ok = cparams.n_seq_max == 1 || !cparams.kv_unified;230 const bool msa_enabled = fa_on && streams_ok;231 232 auto * inp_attn = build_attn_inp_kv_msa(msa_enabled);233 234 static bool warned_no_fa = false;235 if (!fa_on && !warned_no_fa) {236 LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention "237 "(output may be degraded). Enable flash attention for MSA.\n", __func__);238 warned_no_fa = true;239 }240 static bool warned_unified = false;241 if (fa_on && !streams_ok && !warned_unified) {242 LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; MSA needs per-sequence streams "243 "-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__);244 warned_unified = true;245 }246 // ==========================================247 248 // hoisted per-graph MSA state (shared by every sparse layer)249 llm_graph_input_msa * msa = nullptr;250 ggml_tensor * msa_kqm = nullptr;251 ggml_tensor * msa_mf = nullptr; // F32 copy of the FA mask for the final mask add252 int64_t n_kv = 0, n_ps = 0, nblk = 0, ns = 1, n_tps = 0;253 bool msa_decode = false; // gather (1 token per stream) vs mask254 const int blk = mm.msa_p.blk;255 const int64_t Hd = hparams.indexer_n_head; // one indexer head per GQA group256 257 if (msa_enabled) {258 const auto * mctx_msa = static_cast<const llama_kv_cache_msa_context *>(mctx);259 260 msa_kqm = inp_attn->get_kq_mask();261 n_kv = msa_kqm->ne[0];262 n_tps = msa_kqm->ne[1]; // tokens per stream263 ns = msa_kqm->ne[3]; // streams in this ubatch264 GGML_ASSERT(msa_kqm->type == GGML_TYPE_F16 && "MSA requires the FA (f16) mask");265 GGML_ASSERT(n_tps*ns == n_tokens);266 267 // the position axis covers every position currently in the cache and is padded to whole blocks268 n_ps = GGML_PAD((int64_t) mctx_msa->get_n_pos(), blk);269 nblk = n_ps / blk;270 msa_decode = n_tps == 1;271 272 auto inp = std::make_unique<llm_graph_input_msa>(mctx_msa, blk, mm.msa_p.local);273 274 inp->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens); // stream-grouped tokens275 ggml_set_input(inp->bias);276 277 inp->pos_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, n_tokens);278 ggml_set_input(inp->pos_mask);279 280 inp->pos_slot_i = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_ps, ns);281 ggml_set_input(inp->pos_slot_i);282 283 if (msa_decode) {284 inp->pos_slot_f = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, ns);285 ggml_set_input(inp->pos_slot_f);286 } else {287 inp->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns);288 ggml_set_input(inp->cell_blk);289 290 msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32);291 }292 293 msa = (llm_graph_input_msa *) res->add_input(std::move(inp));294 }295 296 ggml_tensor * inp_out_ids = build_inp_out_ids();297 298 for (int il = 0; il < n_layer; ++il) {299 ggml_tensor * inpSA = inpL;300 301 // self-attention302 {303 cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);304 cb(cur, "attn_norm", il);305 306 auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,307 n_embd_head, n_head, n_head_kv, il);308 309 // per-head QK RMSNorm (weights already include Gemma's +1)310 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);311 cb(Qcur, "Qcur_normed", il);312 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);313 cb(Kcur, "Kcur_normed", il);314 315 // partial rotary: only the first n_rot dims are rotated316 Qcur = ggml_rope_ext(317 ctx0, Qcur, inp_pos, nullptr,318 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,319 ext_factor, attn_factor, beta_fast, beta_slow);320 Kcur = ggml_rope_ext(321 ctx0, Kcur, inp_pos, nullptr,322 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,323 ext_factor, attn_factor, beta_fast, beta_slow);324 325 cb(Qcur, "Qcur", il);326 cb(Kcur, "Kcur", il);327 cb(Vcur, "Vcur", il);328 329 const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead;330 331 if (!is_sparse) {332 cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s,333 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,334 1.0f/sqrtf(float(n_embd_head)), il);335 } else {336 const int64_t n_idx_dim = hparams.indexer_head_size; // 128337 338 // Index Branch, project, norm, partial RoPE, cache339 ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur);340 ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur);341 iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, Hd, n_tokens);342 ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1, n_tokens);343 iq = build_norm(iq, model.layers[il].index_q_norm, NULL, LLM_NORM_RMS, il); // +1 baked344 ik = build_norm(ik, model.layers[il].index_k_norm, NULL, LLM_NORM_RMS, il);345 iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,346 freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);347 ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,348 freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);349 350 const auto * mctx_msa_l = static_cast<const llama_kv_cache_msa_context *>(mctx);351 const auto * mctx_cur = mctx_msa_l->get_base();352 const auto * mctx_idx = mctx_msa_l->get_idx();353 ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, ik, inp_attn->get_k_idxs_idx(), il));354 ggml_tensor * ik_kv = mctx_idx->get_k(ctx0, il);355 356 if (inp_attn->self_k_rot) {357 Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);358 Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);359 }360 if (inp_attn->self_v_rot) {361 Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);362 }363 364 // Main branch: store K/V, take cache views365 ggml_build_forward_expand(gf, Qcur);366 ggml_build_forward_expand(gf, Kcur);367 ggml_build_forward_expand(gf, Vcur);368 ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));369 ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));370 ggml_tensor * k = mctx_cur->get_k(ctx0, il);371 ggml_tensor * v = mctx_cur->get_v(ctx0, il);372 GGML_ASSERT(!(v->nb[1] > v->nb[2]) && "MSA assumes v_trans=false (FA on)");373 374 const int64_t D = k->ne[0];375 const int64_t HKV = k->ne[1];376 const int64_t Gp = n_head/HKV;377 GGML_ASSERT(HKV == Hd && "MSA: one indexer head per GQA group");378 GGML_ASSERT(k->ne[3] == ns);379 const int K = mm.msa_p.topk_blocks < (int) nblk ? mm.msa_p.topk_blocks : (int) nblk;380 381 const float kq_scale = 1.0f/sqrtf(float(n_embd_head));382 383 if (msa_decode) {384 // decode: batched over streams top-k + gather, one grouped FA385 // gather the indexer keys through the pos -> cell map386 ggml_tensor * ik3 = ggml_view_3d(ctx0, ik_kv, n_idx_dim, n_kv, ns,387 ik_kv->nb[2], ik_kv->nb[3], 0);388 ggml_tensor * ikp = ggml_get_rows(ctx0, ik3, msa->pos_slot_i); // [n_idx_dim, n_ps, ns]389 ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns);390 ggml_tensor * sc = ggml_mul_mat(ctx0,391 ggml_reshape_4d(ctx0, ikp, n_idx_dim, n_ps, 1, ns), iq4);392 ggml_prec_set_acc(sc, GGML_PREC_F32);393 // unmapped positions come out -inf, so they can never rank into the top-k394 sc = ggml_add_inplace(ctx0, sc,395 ggml_reshape_4d(ctx0, msa->pos_mask, n_ps, 1, 1, ns));396 ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);397 cb(bs, "msa_bs", il);398 399 ggml_tensor * bsf = ggml_add(ctx0, bs,400 ggml_reshape_4d(ctx0, msa->bias, nblk, 1, 1, ns));401 ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // position blocks402 403 // pos idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (positions - mask gather)404 // cell idx: cs[t,k,h,s] = pos_slot[tj] (pos -> cell translation)405 // row idx: tr[t,k,h,s] = cs*HKV + h (per-stream K/V gather)406 ggml_tensor * a = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk);407 a = ggml_reshape_4d(ctx0, a, 1, K, Hd, ns);408 ggml_tensor * tj = ggml_add(ctx0,409 ggml_repeat_4d(ctx0, a, blk, K, Hd, ns),410 ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) blk, 1.0f), blk, 1, 1));411 412 ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);413 414 ggml_tensor * cs = ggml_get_rows(ctx0,415 ggml_reshape_3d(ctx0, msa->pos_slot_f, 1, n_ps, ns), tokj); // [1, blk*K*Hd, ns]416 cs = ggml_reshape_4d(ctx0, cs, blk, K, Hd, ns);417 418 ggml_tensor * tr = ggml_add(ctx0,419 ggml_scale(ctx0, cs, (float) HKV),420 ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd));421 422 ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);423 424 ggml_tensor * k3 = ggml_view_3d(ctx0, k, D, HKV*n_kv, ns, k->nb[1], k->nb[3], 0);425 ggml_tensor * v3 = ggml_view_3d(ctx0, v, D, HKV*n_kv, ns, v->nb[1], v->nb[3], 0);426 ggml_tensor * mp = ggml_reshape_3d(ctx0, msa->pos_mask, 1, n_ps, ns);427 428 ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr);429 ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr);430 ggml_tensor * mg = ggml_get_rows(ctx0, mp, tokj);431 432 // fold (group, stream) onto the FA channel dim433 const ggml_type kt = ggml_is_quantized(k->type) ? GGML_TYPE_F16 : k->type;434 const ggml_type vt = ggml_is_quantized(v->type) ? GGML_TYPE_F16 : v->type;435 ggml_tensor * kfa = ggml_reshape_4d(ctx0, kg, D, (int64_t) blk*K, 1, Hd*ns);436 ggml_tensor * vfa = ggml_reshape_4d(ctx0, vg, D, (int64_t) blk*K, 1, Hd*ns);437 if (kfa->type != kt) { kfa = ggml_cast(ctx0, kfa, kt); }438 if (vfa->type != vt) { vfa = ggml_cast(ctx0, vfa, vt); }439 // the FA mask must be F16440 ggml_tensor * mfa = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, (int64_t) blk*K, 1, 1, Hd*ns), GGML_TYPE_F16);441 442 cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il);443 } else {444 // batch: per-stream loop445 std::vector<ggml_tensor *> outs(ns);446 for (int64_t st = 0; st < ns; ++st) {447 ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, n_tps,448 iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]);449 ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv,450 ik_kv->nb[2], st*ik_kv->nb[3]);451 ggml_tensor * psl_s = ggml_view_1d(ctx0, msa->pos_slot_i, n_ps,452 st*msa->pos_slot_i->nb[1]);453 ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, n_tps,454 msa->pos_mask->nb[1], msa->pos_mask->nb[1], st*n_tps*msa->pos_mask->nb[1]);455 ggml_tensor * cb_s = ggml_view_1d(ctx0, msa->cell_blk, n_kv,456 st*msa->cell_blk->nb[1]);457 ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, n_tps, 1,458 msa_mf->nb[1], msa_mf->nb[3], st*msa_mf->nb[3]);459 ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, n_tps,460 msa->bias->nb[1], msa->bias->nb[1], st*n_tps*msa->bias->nb[1]);461 ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps,462 Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]);463 ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1,464 k->nb[1], k->nb[2], k->nb[3], st*k->nb[3]);465 ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1,466 v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]);467 468 // block scores: the indexer keys are gathered through the pos -> cell map first469 // scores are unscaled, only the top-k ordering matters470 ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s); // [n_idx_dim, n_ps]471 ggml_tensor * sc = ggml_mul_mat(ctx0, ikp,472 ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps));473 // indexer scores run in F32474 ggml_prec_set_acc(sc, GGML_PREC_F32);475 sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps);476 // unmapped positions (holes, padding, empty cells) come out -inf477 sc = ggml_add_inplace(ctx0, sc, pm_s);478 ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);479 cb(bs, "msa_bs", il);480 481 // bias the scores so locally-forced blocks always rank first482 ggml_tensor * bsf = ggml_add(ctx0, bs, bias_s); // [nblk, Hd, n_tps]483 cb(bsf, "msa_bsf", il);484 485 ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // [K, Hd, n_tps] i32486 487 ggml_tensor * ninf = ggml_cast(ctx0,488 ggml_scale_bias(ctx0, bias_s, 0.0f, -1e30f),489 GGML_TYPE_F16); // [nblk, 1, n_tps]490 ninf = ggml_repeat_4d(ctx0, ninf, nblk, Hd, n_tps, 1);491 ggml_tensor * zero = ggml_scale(ctx0,492 ggml_cast(ctx0, idx, GGML_TYPE_F32), 0.0f);493 ggml_tensor * bm = ggml_set_rows(ctx0,494 ggml_reshape_3d(ctx0, ninf, 1, nblk, Hd*n_tps),495 ggml_reshape_3d(ctx0, zero, 1, K, Hd*n_tps),496 ggml_reshape_2d(ctx0, idx, K, Hd*n_tps));497 bm = ggml_reshape_3d(ctx0, bm, nblk, Hd, n_tps);498 bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, n_tps, Hd]499 cb(bm, "msa_block_mask", il);500 501 // expand block -> cell granularity through the cell -> position block502 // map, then combine with the causal mask. empty cells are masked by the causal mask.503 ggml_tensor * bm2 = ggml_cont(ctx0, ggml_transpose(ctx0,504 ggml_reshape_2d(ctx0, bm, nblk, n_tps*Hd))); // [n_tps*Hd, nblk]505 ggml_tensor * bmc = ggml_get_rows(ctx0, bm2, cb_s); // [n_tps*Hd, n_kv] F32506 ggml_tensor * bmx = ggml_cont(ctx0, ggml_transpose(ctx0, bmc));507 bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd);508 ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, mf_s);509 mask4 = ggml_cast(ctx0,510 ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd), GGML_TYPE_F16);511 cb(mask4, "msa_mask4", il);512 513 // cache views with groups on ne[3];514 ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2);515 ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2);516 517 outs[st] = build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il);518 }519 cur = outs[0];520 for (int64_t st = 1; st < ns; ++st) {521 cur = ggml_concat(ctx0, cur, outs[st], 1);522 }523 }524 if (inp_attn->self_v_rot) {525 cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot);526 }527 cb(cur, "kqv_out", il);528 if (model.layers[il].wo) {529 cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);530 }531 }532 }533 534 if (il == n_layer - 1 && inp_out_ids) {535 cur = ggml_get_rows(ctx0, cur, inp_out_ids);536 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);537 }538 539 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);540 cb(ffn_inp, "ffn_inp", il);541 542 cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);543 cb(cur, "ffn_norm", il);544 545 if ((uint32_t) il < hparams.n_layer_dense_lead) {546 // leading dense FFN (swigluoai)547 cur = build_ffn(cur,548 model.layers[il].ffn_up, NULL, NULL,549 model.layers[il].ffn_gate, NULL, NULL,550 model.layers[il].ffn_down, NULL, NULL,551 NULL,552 LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);553 cb(cur, "ffn_out", il);554 } else {555 // routed experts (swigluoai MoE)556 ggml_tensor * moe_out = build_moe_ffn(cur,557 model.layers[il].ffn_gate_inp,558 model.layers[il].ffn_up_exps,559 model.layers[il].ffn_gate_exps,560 model.layers[il].ffn_down_exps,561 model.layers[il].ffn_exp_probs_b,562 n_expert, n_expert_used,563 LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm,564 hparams.expert_weights_scale,565 (llama_expert_gating_func_type) hparams.expert_gating_func,566 il);567 cb(moe_out, "ffn_moe_out", il);568 569 // shared expert (swigluoai)570 ggml_tensor * ffn_shexp = build_ffn(cur,571 model.layers[il].ffn_up_shexp, NULL, NULL,572 model.layers[il].ffn_gate_shexp, NULL, NULL,573 model.layers[il].ffn_down_shexp, NULL, NULL,574 NULL,575 LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);576 cb(ffn_shexp, "ffn_shexp", il);577 578 cur = ggml_add(ctx0, moe_out, ffn_shexp);579 cb(cur, "ffn_out", il);580 }581 582 cur = ggml_add(ctx0, cur, ffn_inp);583 584 cur = build_cvec(cur, il);585 cb(cur, "l_out", il);586 587 // input for next layer588 inpL = cur;589 }590 591 cur = inpL;592 593 cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);594 cb(cur, "result_norm", -1);595 res->t_embd = cur;596 597 // lm_head598 cur = build_lora_mm(model.output, cur, model.output_s);599 cb(cur, "result_output", -1);600 res->t_logits = cur;601 602 ggml_build_forward_expand(gf, cur);603}604