Brunobkr/llama.cpp_AlgMor24_github

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(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,  hparams.n_ff_exp);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    switch (hparams.n_layer()) {29        case 60: type = LLM_TYPE_428B_A23B; break;30        default: type = LLM_TYPE_UNKNOWN;31    }32}33 34void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) {35    LLAMA_LOAD_LOCALS;36    const int64_t n_expert_shared = hparams.n_expert_shared;37    const int64_t n_ff_exp        = hparams.n_ff_exp;38 39    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);40 41    // output42    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);43    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, 0);44 45    for (int i = 0; i < n_layer; ++i) {46        auto & layer = layers[i];47 48        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0);49        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);50 51        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);52        // per-head QK-norm: a single head_dim vector applied to every head53        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);54        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);55 56        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);57 58        if (i < (int) hparams.n_layer_dense_lead) {59            // leading dense layers60            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);61            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);62            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);63        } else {64            // routed experts65            layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,    "weight", i), {n_embd, n_expert}, 0);66            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias",   i), {n_expert}, 0);67            layer.ffn_gate_exps   = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS,   "weight", i), {n_embd, n_ff_exp, n_expert}, 0);68            layer.ffn_down_exps   = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS,   "weight", i), {n_ff_exp, n_embd, n_expert}, 0);69            layer.ffn_up_exps     = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,     "weight", i), {n_embd, n_ff_exp, n_expert}, 0);70 71            // shared expert72            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);73            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {        n_ff_exp * n_expert_shared, n_embd}, 0);74            layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);75 76            // indexer77            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);78            layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, 0);79            layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, 0);80            layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, 0);81        }82    }83}84 85std::unique_ptr<llm_graph_context> llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const {86    return std::make_unique<graph>(*this, params);87}88 89class llm_graph_input_msa : public llm_graph_input_i {90public:91    llm_graph_input_msa(const llama_kv_cache_msa_context * mctx, int blk, int local) :92        mctx(mctx), blk(blk), local(local) {}93 94    void set_input(const llama_ubatch * ubatch) override {95        if (pos_slot_i) { mctx->set_input_pos_slot(pos_slot_i, ubatch); }96        if (pos_slot_f) { mctx->set_input_pos_slot(pos_slot_f, ubatch); }97        if (cell_blk)   { mctx->set_input_cell_pos(cell_blk, ubatch, blk); }98        if (pos_mask)   { mctx->set_input_pos_mask(pos_mask, ubatch); }99 100        // local-force bias over position blocks101        if (bias && ubatch->pos) {102            const int64_t n_tokens = ubatch->n_tokens;103            const int64_t nblk     = bias->ne[0];104            std::vector<float> data((size_t) nblk * n_tokens, 0.0f);105            for (int64_t i = 0; i < n_tokens; ++i) {106                const int64_t L = ubatch->pos[i] / blk;107                for (int l = 0; l < local && L - l >= 0; ++l) {108                    if (L - l < nblk) {109                        data[(size_t) i * nblk + (L - l)] = 1e30f;110                    }111                }112            }113            ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float));114        }115    }116 117    // valid as long as the tensor dims still match the new ubatch/cache window and the118    // ubatch is in the same regime (decode graphs have pos_slot_f, batch graphs cell_blk)119    bool can_reuse(const llm_graph_params & params) override {120        const auto * mctx_new = static_cast<const llama_kv_cache_msa_context *>(params.mctx);121 122        this->mctx = mctx_new;123 124        const int64_t n_ps = GGML_PAD((int64_t) mctx_new->get_n_pos(), blk);125        const int64_t ns   = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq;126 127        const bool decode = params.ubatch.n_tokens == ns;   // one token per stream128 129        bool res = true;130 131        res &= bias->ne[0] * blk == n_ps;132        res &= bias->ne[1]       == params.ubatch.n_tokens;133 134        res &= pos_mask->ne[0] == n_ps;135        res &= pos_mask->ne[1] == params.ubatch.n_tokens;136 137        res &= pos_slot_i->ne[0] == n_ps;138        res &= pos_slot_i->ne[1] == ns;139 140        res &= decode == (pos_slot_f != nullptr);141        res &= decode == (cell_blk   == nullptr);142 143        if (pos_slot_f) {144            res &= pos_slot_f->ne[0] == n_ps;145            res &= pos_slot_f->ne[1] == ns;146        }147 148        if (cell_blk) {149            res &= cell_blk->ne[0] == (int64_t) mctx_new->get_base()->get_n_kv();150            res &= cell_blk->ne[1] == ns;151        }152 153        return res;154    }155 156    ggml_tensor * bias       = nullptr; // F32 [nblk, n_tokens] local-force bias (position blocks)157    ggml_tensor * pos_mask   = nullptr; // F32 [n_ps, n_tokens] 0/-inf visibility, by position158    ggml_tensor * pos_slot_i = nullptr; // I32 [n_ps, ns]       pos -> cell (get_rows index)159    ggml_tensor * pos_slot_f = nullptr; // F32 [n_ps, ns]       pos -> cell (gatherable values, decode)160    ggml_tensor * cell_blk   = nullptr; // I32 [n_kv, ns]       cell -> position block (batch)161 162    const llama_kv_cache_msa_context * mctx;163 164    int blk;165    int local;166};167 168// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3])169ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa(170        ggml_tensor * q_cur,   // [D, HQ, T]171        ggml_tensor * k,       // [D, n_keys, 1, C]172        ggml_tensor * v,       // [D, n_keys, 1, C]173        ggml_tensor * mask,    // [n_keys, R, 1, C] f16, contiguous174        int64_t Gp, float kq_scale, int il) const {175 176    const int64_t D  = q_cur->ne[0];177    const int64_t HQ = q_cur->ne[1];178    const int64_t T  = q_cur->ne[2];179    const int64_t C  = k->ne[3];180    const int64_t R  = HQ*T/(Gp*C);181    GGML_ASSERT(Gp*C*R == HQ*T);182    GGML_ASSERT(mask->type == GGML_TYPE_F16);183 184    // [D, HQ, T] -> [D, Gp, C, R] -> [D, R, Gp, C]185    // batch  (C=HKV,   R=T): channel = group186    // decode (C=HKV*ns, R=1): channel = (group, stream), group innermost187    ggml_tensor * q = ggml_reshape_4d(ctx0, q_cur, D, Gp, C, R);188    q = ggml_permute(ctx0, q, 0, 2, 3, 1);189 190    ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale,191                                          hparams.f_max_alibi_bias, 0.0f);192    ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32);193    cb(o, "msa_fattn", il);194 195    // [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T]196    o = ggml_permute(ctx0, o, 0, 1, 3, 2);197    if (!ggml_is_contiguous(o)) {198        o = ggml_cont(ctx0, o);   // no-op layout at decode (R == 1), copy at batch199    }200    return ggml_reshape_2d(ctx0, o, D*HQ, T);201}202 203llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {204    const int64_t n_embd_head = hparams.n_embd_head_v();205    const auto & mm = static_cast<const llama_model_minimax_m3 &>(model);206 207    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());208    // partial rotary: head_dim != n_rot, so don't assert n_embd_head == n_rot209 210    ggml_tensor * cur;211    ggml_tensor * inpL;212 213    inpL = build_inp_embd(model.tok_embd);214 215    ggml_tensor * inp_pos = build_inp_pos();216 217    // ==========================================218    // TODO: avoid such kind of complexity in the model graphs219 220    // MSA calls ggml_flash_attn_ext directly and assumes the non-transposed V layout that221    // llama.cpp only provides when flash attention is enabled. Block selection is anchored222    // to absolute KV cache slots, which equal positions only for append-only per-stream223    // caches either a single sequence, or multiple sequences with kv_unified == false (each224    // stream then has its own slot space). A unified cache with multiple sequences225    // interleaves slots and would silently break block anchoring so it falls back to dense.226    const bool fa_on       = cparams.flash_attn;227    const bool streams_ok  = cparams.n_seq_max == 1 || !cparams.kv_unified;228    const bool msa_enabled = fa_on && streams_ok;229 230    auto * inp_attn = build_attn_inp_kv_msa(msa_enabled);231 232    static bool warned_no_fa = false;233    if (!fa_on && !warned_no_fa) {234        LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention "235                       "(output may be degraded). Enable flash attention for MSA.\n", __func__);236        warned_no_fa = true;237    }238    static bool warned_unified = false;239    if (fa_on && !streams_ok && !warned_unified) {240        LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; MSA needs per-sequence streams "241                       "-> running DENSE attention. Output may be degraded. Drop --kv-unified to enable MSA.\n", __func__);242        warned_unified = true;243    }244    // ==========================================245 246    // hoisted per-graph MSA state (shared by every sparse layer)247    llm_graph_input_msa * msa = nullptr;248    ggml_tensor * msa_kqm = nullptr;249    ggml_tensor * msa_mf  = nullptr;   // F32 copy of the FA mask for the final mask add250    int64_t n_kv = 0, n_ps = 0, nblk = 0, ns = 1, n_tps = 0;251    bool msa_decode = false;           // gather (1 token per stream) vs mask252    const int     blk = mm.msa_p.blk;253    const int64_t Hd  = hparams.indexer_n_head;   // one indexer head per GQA group254 255    if (msa_enabled) {256        const auto * mctx_msa = static_cast<const llama_kv_cache_msa_context *>(mctx);257 258        msa_kqm = inp_attn->get_kq_mask();259        n_kv  = msa_kqm->ne[0];260        n_tps = msa_kqm->ne[1];        // tokens per stream261        ns    = msa_kqm->ne[3];        // streams in this ubatch262        GGML_ASSERT(msa_kqm->type == GGML_TYPE_F16 && "MSA requires the FA (f16) mask");263        GGML_ASSERT(n_tps*ns == n_tokens);264 265        // the position axis covers every position currently in the cache and is padded to whole blocks266        n_ps = GGML_PAD((int64_t) mctx_msa->get_n_pos(), blk);267        nblk = n_ps / blk;268        msa_decode = n_tps == 1;269 270        auto inp = std::make_unique<llm_graph_input_msa>(mctx_msa, blk, mm.msa_p.local);271 272        inp->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk, n_tokens);  // stream-grouped tokens273        ggml_set_input(inp->bias);274 275        inp->pos_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, n_tokens);276        ggml_set_input(inp->pos_mask);277 278        inp->pos_slot_i = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_ps, ns);279        ggml_set_input(inp->pos_slot_i);280 281        if (msa_decode) {282            inp->pos_slot_f = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_ps, ns);283            ggml_set_input(inp->pos_slot_f);284        } else {285            inp->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns);286            ggml_set_input(inp->cell_blk);287 288            msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32);289        }290 291        msa = (llm_graph_input_msa *) res->add_input(std::move(inp));292    }293 294    ggml_tensor * inp_out_ids = build_inp_out_ids();295 296    for (int il = 0; il < n_layer; ++il) {297        ggml_tensor * inpSA = inpL;298 299        // self-attention300        {301            cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);302            cb(cur, "attn_norm", il);303 304            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,305                    n_embd_head, n_head, n_head_kv, il);306 307            // per-head QK RMSNorm (weights already include Gemma's +1)308            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);309            cb(Qcur, "Qcur_normed", il);310            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);311            cb(Kcur, "Kcur_normed", il);312 313            // partial rotary: only the first n_rot dims are rotated314            Qcur = ggml_rope_ext(315                ctx0, Qcur, inp_pos, nullptr,316                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,317                ext_factor, attn_factor, beta_fast, beta_slow);318            Kcur = ggml_rope_ext(319                ctx0, Kcur, inp_pos, nullptr,320                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,321                ext_factor, attn_factor, beta_fast, beta_slow);322 323            cb(Qcur, "Qcur", il);324            cb(Kcur, "Kcur", il);325            cb(Vcur, "Vcur", il);326 327            const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead;328 329            if (!is_sparse) {330                cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s,331                        Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,332                        1.0f/sqrtf(float(n_embd_head)), il);333            } else {334                const int64_t n_idx_dim = hparams.indexer_head_size;   // 128335 336                // Index Branch, project, norm, partial RoPE, cache337                ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur);338                ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur);339                iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, Hd, n_tokens);340                ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1,  n_tokens);341                iq = build_norm(iq, model.layers[il].index_q_norm, NULL, LLM_NORM_RMS, il);  // +1 baked342                ik = build_norm(ik, model.layers[il].index_k_norm, NULL, LLM_NORM_RMS, il);343                iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,344                                   freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);345                ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig,346                                   freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);347 348                const auto * mctx_msa_l = static_cast<const llama_kv_cache_msa_context *>(mctx);349                const auto * mctx_cur = mctx_msa_l->get_base();350                const auto * mctx_idx = mctx_msa_l->get_idx();351                ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, ik, inp_attn->get_k_idxs_idx(), il));352                ggml_tensor * ik_kv = mctx_idx->get_k(ctx0, il);353 354                if (inp_attn->self_k_rot) {355                    Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);356                    Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);357                }358                if (inp_attn->self_v_rot) {359                    Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);360                }361 362                // Main branch: store K/V, take cache views363                ggml_build_forward_expand(gf, Qcur);364                ggml_build_forward_expand(gf, Kcur);365                ggml_build_forward_expand(gf, Vcur);366                ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));367                ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));368                ggml_tensor * k = mctx_cur->get_k(ctx0, il);369                ggml_tensor * v = mctx_cur->get_v(ctx0, il);370                GGML_ASSERT(!(v->nb[1] > v->nb[2]) && "MSA assumes v_trans=false (FA on)");371 372                const int64_t D   = k->ne[0];373                const int64_t HKV = k->ne[1];374                const int64_t Gp  = n_head/HKV;375                GGML_ASSERT(HKV == Hd && "MSA: one indexer head per GQA group");376                GGML_ASSERT(k->ne[3] == ns);377                const int K = mm.msa_p.topk_blocks < (int) nblk ? mm.msa_p.topk_blocks : (int) nblk;378 379                const float kq_scale = 1.0f/sqrtf(float(n_embd_head));380 381                if (msa_decode) {382                    // decode: batched over streams top-k + gather, one grouped FA383                    // gather the indexer keys through the pos -> cell map384                    ggml_tensor * ik3 = ggml_view_3d(ctx0, ik_kv, n_idx_dim, n_kv, ns,385                            ik_kv->nb[2], ik_kv->nb[3], 0);386                    ggml_tensor * ikp = ggml_get_rows(ctx0, ik3, msa->pos_slot_i);   // [n_idx_dim, n_ps, ns]387                    ggml_tensor * iq4 = ggml_reshape_4d(ctx0, iq, n_idx_dim, Hd, 1, ns);388                    ggml_tensor * sc  = ggml_mul_mat(ctx0,389                            ggml_reshape_4d(ctx0, ikp, n_idx_dim, n_ps, 1, ns), iq4);390                    ggml_mul_mat_set_prec(sc, GGML_PREC_F32);391                    // unmapped positions come out -inf, so they can never rank into the top-k392                    sc = ggml_add_inplace(ctx0, sc,393                            ggml_reshape_4d(ctx0, msa->pos_mask, n_ps, 1, 1, ns));394                    ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);395                    cb(bs, "msa_bs", il);396 397                    ggml_tensor * bsf = ggml_add(ctx0, bs,398                            ggml_reshape_4d(ctx0, msa->bias, nblk, 1, 1, ns));399                    ggml_tensor * idx = ggml_top_k(ctx0, bsf, K);   // position blocks400 401                    // pos idx:  tj[t,k,h,s] = blk*idx[k,h,s] + t   (positions - mask gather)402                    // cell idx: cs[t,k,h,s] = pos_slot[tj]         (pos -> cell translation)403                    // row idx:  tr[t,k,h,s] = cs*HKV + h           (per-stream K/V gather)404                    ggml_tensor * a = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk);405                    a = ggml_reshape_4d(ctx0, a, 1, K, Hd, ns);406                    ggml_tensor * tj = ggml_add(ctx0,407                            ggml_repeat_4d(ctx0, a, blk, K, Hd, ns),408                            ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) blk, 1.0f), blk, 1, 1));409 410                    ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);411 412                    ggml_tensor * cs = ggml_get_rows(ctx0,413                            ggml_reshape_3d(ctx0, msa->pos_slot_f, 1, n_ps, ns), tokj);   // [1, blk*K*Hd, ns]414                    cs = ggml_reshape_4d(ctx0, cs, blk, K, Hd, ns);415 416                    ggml_tensor * tr = ggml_add(ctx0,417                            ggml_scale(ctx0, cs, (float) HKV),418                            ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) HKV, 1.0f), 1, 1, Hd));419 420                    ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, (int64_t) blk*K*Hd, ns), GGML_TYPE_I32);421 422                    ggml_tensor * k3 = ggml_view_3d(ctx0, k, D, HKV*n_kv, ns, k->nb[1], k->nb[3], 0);423                    ggml_tensor * v3 = ggml_view_3d(ctx0, v, D, HKV*n_kv, ns, v->nb[1], v->nb[3], 0);424                    ggml_tensor * mp = ggml_reshape_3d(ctx0, msa->pos_mask, 1, n_ps, ns);425 426                    ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr);427                    ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr);428                    ggml_tensor * mg = ggml_get_rows(ctx0, mp, tokj);429 430                    // fold (group, stream) onto the FA channel dim431                    const ggml_type kt = ggml_is_quantized(k->type) ? GGML_TYPE_F16 : k->type;432                    const ggml_type vt = ggml_is_quantized(v->type) ? GGML_TYPE_F16 : v->type;433                    ggml_tensor * kfa = ggml_reshape_4d(ctx0, kg, D, (int64_t) blk*K, 1, Hd*ns);434                    ggml_tensor * vfa = ggml_reshape_4d(ctx0, vg, D, (int64_t) blk*K, 1, Hd*ns);435                    if (kfa->type != kt) { kfa = ggml_cast(ctx0, kfa, kt); }436                    if (vfa->type != vt) { vfa = ggml_cast(ctx0, vfa, vt); }437                    // the FA mask must be F16438                    ggml_tensor * mfa = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, (int64_t) blk*K, 1, 1, Hd*ns), GGML_TYPE_F16);439 440                    cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il);441                } else {442                    // batch: per-stream loop443                    std::vector<ggml_tensor *> outs(ns);444                    for (int64_t st = 0; st < ns; ++st) {445                        ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, n_tps,446                                iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]);447                        ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv,448                                ik_kv->nb[2], st*ik_kv->nb[3]);449                        ggml_tensor * psl_s = ggml_view_1d(ctx0, msa->pos_slot_i, n_ps,450                                st*msa->pos_slot_i->nb[1]);451                        ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, n_tps,452                                msa->pos_mask->nb[1], msa->pos_mask->nb[1], st*n_tps*msa->pos_mask->nb[1]);453                        ggml_tensor * cb_s = ggml_view_1d(ctx0, msa->cell_blk, n_kv,454                                st*msa->cell_blk->nb[1]);455                        ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, n_tps, 1,456                                msa_mf->nb[1], msa_mf->nb[3], st*msa_mf->nb[3]);457                        ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, n_tps,458                                msa->bias->nb[1], msa->bias->nb[1], st*n_tps*msa->bias->nb[1]);459                        ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps,460                                Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]);461                        ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1,462                                k->nb[1], k->nb[2], k->nb[3], st*k->nb[3]);463                        ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1,464                                v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]);465 466                        // block scores: the indexer keys are gathered through the pos -> cell map first467                        // scores are unscaled, only the top-k ordering matters468                        ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s);   // [n_idx_dim, n_ps]469                        ggml_tensor * sc = ggml_mul_mat(ctx0, ikp,470                                ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps));471                        // indexer scores run in F32472                        ggml_mul_mat_set_prec(sc, GGML_PREC_F32);473                        sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps);474                        // unmapped positions (holes, padding, empty cells) come out -inf475                        sc = ggml_add_inplace(ctx0, sc, pm_s);476                        ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);477                        cb(bs, "msa_bs", il);478 479                        // bias the scores so locally-forced blocks always rank first480                        ggml_tensor * bsf = ggml_add(ctx0, bs, bias_s);   // [nblk, Hd, n_tps]481                        cb(bsf, "msa_bsf", il);482 483                        ggml_tensor * idx = ggml_top_k(ctx0, bsf, K);   // [K, Hd, n_tps] i32484 485                        ggml_tensor * ninf = ggml_cast(ctx0,486                                ggml_scale_bias(ctx0, bias_s, 0.0f, -1e30f),487                                GGML_TYPE_F16);                              // [nblk, 1, n_tps]488                        ninf = ggml_repeat_4d(ctx0, ninf, nblk, Hd, n_tps, 1);489                        ggml_tensor * zero = ggml_scale(ctx0,490                                ggml_cast(ctx0, idx, GGML_TYPE_F32), 0.0f);491                        ggml_tensor * bm = ggml_set_rows(ctx0,492                                ggml_reshape_3d(ctx0, ninf, 1, nblk, Hd*n_tps),493                                ggml_reshape_3d(ctx0, zero, 1, K,    Hd*n_tps),494                                ggml_reshape_2d(ctx0, idx,     K,    Hd*n_tps));495                        bm = ggml_reshape_3d(ctx0, bm, nblk, Hd, n_tps);496                        bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, n_tps, Hd]497                        cb(bm, "msa_block_mask", il);498 499                        // expand block -> cell granularity through the cell -> position block500                        // map, then combine with the causal mask. empty cells are masked by the causal mask.501                        ggml_tensor * bm2 = ggml_cont(ctx0, ggml_transpose(ctx0,502                                ggml_reshape_2d(ctx0, bm, nblk, n_tps*Hd)));       // [n_tps*Hd, nblk]503                        ggml_tensor * bmc = ggml_get_rows(ctx0, bm2, cb_s);        // [n_tps*Hd, n_kv] F32504                        ggml_tensor * bmx = ggml_cont(ctx0, ggml_transpose(ctx0, bmc));505                        bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd);506                        ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, mf_s);507                        mask4 = ggml_cast(ctx0,508                                ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd), GGML_TYPE_F16);509                        cb(mask4, "msa_mask4", il);510 511                        // cache views with groups on ne[3];512                        ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2);513                        ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2);514 515                        outs[st] = build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il);516                    }517                    cur = outs[0];518                    for (int64_t st = 1; st < ns; ++st) {519                        cur = ggml_concat(ctx0, cur, outs[st], 1);520                    }521                }522                if (inp_attn->self_v_rot) {523                    cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot);524                }525                cb(cur, "kqv_out", il);526                if (model.layers[il].wo) {527                    cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);528                }529            }530        }531 532        if (il == n_layer - 1 && inp_out_ids) {533            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);534            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);535        }536 537        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);538        cb(ffn_inp, "ffn_inp", il);539 540        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);541        cb(cur, "ffn_norm", il);542 543        if ((uint32_t) il < hparams.n_layer_dense_lead) {544            // leading dense FFN (swigluoai)545            cur = build_ffn(cur,546                    model.layers[il].ffn_up,   NULL, NULL,547                    model.layers[il].ffn_gate, NULL, NULL,548                    model.layers[il].ffn_down, NULL, NULL,549                    NULL,550                    LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);551            cb(cur, "ffn_out", il);552        } else {553            // routed experts (swigluoai MoE)554            ggml_tensor * moe_out = build_moe_ffn(cur,555                    model.layers[il].ffn_gate_inp,556                    model.layers[il].ffn_up_exps,557                    model.layers[il].ffn_gate_exps,558                    model.layers[il].ffn_down_exps,559                    model.layers[il].ffn_exp_probs_b,560                    n_expert, n_expert_used,561                    LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm,562                    hparams.expert_weights_scale,563                    (llama_expert_gating_func_type) hparams.expert_gating_func,564                    il);565            cb(moe_out, "ffn_moe_out", il);566 567            // shared expert (swigluoai)568            ggml_tensor * ffn_shexp = build_ffn(cur,569                    model.layers[il].ffn_up_shexp,   NULL, NULL,570                    model.layers[il].ffn_gate_shexp, NULL, NULL,571                    model.layers[il].ffn_down_shexp, NULL, NULL,572                    NULL,573                    LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il);574            cb(ffn_shexp, "ffn_shexp", il);575 576            cur = ggml_add(ctx0, moe_out, ffn_shexp);577            cb(cur, "ffn_out", il);578        }579 580        cur = ggml_add(ctx0, cur, ffn_inp);581 582        cur = build_cvec(cur, il);583        cb(cur, "l_out", il);584 585        // input for next layer586        inpL = cur;587    }588 589    cur = inpL;590 591    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);592    cb(cur, "result_norm", -1);593    res->t_embd = cur;594 595    // lm_head596    cur = build_lora_mm(model.output, cur, model.output_s);597    cb(cur, "result_output", -1);598    res->t_logits = cur;599 600    ggml_build_forward_expand(gf, cur);601}602