hack/tools/scripts/llm-key-hunter/results/_cache/content/canahmetozguven_iosnoteapp.json

{"repo": "canahmetozguven/iosnoteapp", "files": {"ios/llama/src/llama-graph.cpp": {"3b38a88567de9c9ef34225bb9e8fbfa035c66ac4": {"t": 1786068559.0075927, "n": 76786, "h": "c203955ae6cf8467", "c": "#include \"llama-graph.h\"\n\n#include \"llama-impl.h\"\n#include \"llama-batch.h\"\n#include \"llama-cparams.h\"\n\n#include \"llama-kv-cache.h\"\n#include \"llama-kv-cache-iswa.h\"\n#include \"llama-memory-hybrid.h\"\n#include \"llama-memory-recurrent.h\"\n\n#include <cassert>\n#include <cmath>\n#include <cstring>\n#include <unordered_set>\n\nvoid llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {\n    if (ubatch->token) {\n        const int64_t n_tokens = ubatch->n_tokens;\n\n        ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens));\n    }\n\n    if (ubatch->embd) {\n        const int64_t n_embd   = embd->ne[0];\n        const int64_t n_tokens = ubatch->n_tokens;\n\n        ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(embd));\n    }\n}\n\nbool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {\n    bool res = true;\n\n    res &= (!tokens && !params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens);\n    res &= (!embd   && !params.ubatch.embd)  || (embd   &&   embd->ne[1] == params.ubatch.n_tokens);\n\n    return res;\n}\n\nvoid llm_graph_input_pos::set_input(const llama_ubatch * ubatch) {\n    if (ubatch->pos && pos) {\n        const int64_t n_tokens = ubatch->n_tokens;\n\n        if (ubatch->token && n_pos_per_embd == 4) {\n            // in case we're using M-RoPE with text tokens, convert the 1D positions to 4D\n            // the 3 first dims are the same, and 4th dim is all 0\n            std::vector<llama_pos> pos_data(n_tokens*n_pos_per_embd);\n            // copy the first dimension\n            for (int i = 0; i < n_tokens; ++i) {\n                pos_data[               i] = ubatch->pos[i];\n                pos_data[    n_tokens + i] = ubatch->pos[i];\n                pos_data[2 * n_tokens + i] = ubatch->pos[i];\n                pos_data[3 * n_tokens + i] = 0; // 4th dim is 0\n            }\n            ggml_backend_tensor_set(pos, pos_data.data(), 0, pos_data.size()*ggml_element_size(pos));\n        } else {\n            ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos));\n        }\n    }\n}\n\nbool llm_graph_input_pos::can_reuse(const llm_graph_params & params) {\n    bool res = true;\n\n    res &= pos->ne[0] == params.ubatch.n_tokens*n_pos_per_embd;\n\n    return res;\n}\n\nvoid llm_graph_input_attn_temp::set_input(const llama_ubatch * ubatch) {\n    if (ubatch->pos && attn_scale) {\n        const int64_t n_tokens = ubatch->n_tokens;\n\n        GGML_ASSERT(f_attn_temp_scale != 0.0f);\n        GGML_ASSERT(n_attn_temp_floor_scale != 0);\n\n        std::vector<float> attn_scale_data(n_tokens, 0.0f);\n        for (int i = 0; i < n_tokens; ++i) {\n            const float pos = ubatch->pos[i];\n            attn_scale_data[i] = std::log(\n                std::floor((pos + f_attn_temp_offset) / n_attn_temp_floor_scale) + 1.0\n            ) * f_attn_temp_scale + 1.0;\n        }\n\n        ggml_backend_tensor_set(attn_scale, attn_scale_data.data(), 0, n_tokens*ggml_element_size(attn_scale));\n    }\n}\n\nvoid llm_graph_input_pos_bucket::set_input(const llama_ubatch * ubatch) {\n    if (pos_bucket) {\n        const int64_t n_tokens = ubatch->n_tokens;\n\n        GGML_ASSERT(ggml_backend_buffer_is_host(pos_bucket->buffer));\n        GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing\n\n        int32_t * data = (int32_t *) pos_bucket->data;\n\n        for (int j = 0; j < n_tokens; ++j) {\n            for (int i = 0; i < n_tokens; ++i) {\n                data[j*n_tokens + i] = llama_relative_position_bucket(ubatch->pos[i], ubatch->pos[j], hparams.n_rel_attn_bkts, true);\n            }\n        }\n    }\n}\n\nvoid llm_graph_input_pos_bucket_kv::set_input(const llama_ubatch * ubatch) {\n    if (pos_bucket) {\n        mctx->set_input_pos_bucket(pos_bucket, ubatch);\n    }\n}\n\nvoid llm_graph_input_out_ids::set_input(const llama_ubatch * ubatch) {\n    GGML_ASSERT(out_ids);\n\n    const int64_t n_tokens = ubatch->n_tokens;\n\n    GGML_ASSERT(ggml_backend_buffer_is_host(out_ids->buffer));\n    int32_t * data = (int32_t *) out_ids->data;\n\n    if (n_outputs == n_tokens) {\n        for (int i = 0; i < n_tokens; ++i) {\n            data[i] = i;\n        }\n\n        return;\n    }\n\n    GGML_ASSERT(ubatch->output);\n\n    int n_outputs = 0;\n\n    for (int i = 0; i < n_tokens; ++i) {\n        if (ubatch->output[i]) {\n            data[n_outputs++] = i;\n        }\n    }\n}\n\nbool llm_graph_input_out_ids::can_reuse(const llm_graph_params & params) {\n    bool res = true;\n\n    res &= n_outputs == params.n_outputs;\n\n    return res;\n}\n\nvoid llm_graph_input_mean::set_input(const llama_ubatch * ubatch) {\n    if (cparams.embeddings && cparams.pooling_type == LLAMA_POOLING_TYPE_MEAN) {\n        const int64_t n_tokens     = ubatch->n_tokens;\n        const int64_t n_seq_tokens = ubatch->n_seq_tokens;\n        const int64_t n_seqs_unq   = ubatch->n_seqs_unq;\n\n        GGML_ASSERT(mean);\n        GGML_ASSERT(ggml_backend_buffer_is_host(mean->buffer));\n\n        float * data = (float *) mean->data;\n        memset(mean->data, 0, n_tokens*n_seqs_unq*ggml_element_size(mean));\n\n        std::vector<uint64_t> sums(n_seqs_unq, 0);\n        for (int i = 0; i < n_tokens; i += n_seq_tokens) {\n            for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {\n                const llama_seq_id seq_id  = ubatch->seq_id[i][s];\n                const int32_t      seq_idx = ubatch->seq_idx[seq_id];\n\n                sums[seq_idx] += ubatch->n_seq_tokens;\n            }\n        }\n\n        std::vector<float> div(n_seqs_unq, 0.0f);\n        for (int s = 0; s < n_seqs_unq; ++s) {\n            const uint64_t sum = sums[s];\n            if (sum > 0) {\n                div[s] = 1.0f/float(sum);\n            }\n        }\n\n        for (int i = 0; i < n_tokens; i += n_seq_tokens) {\n            for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {\n                const llama_seq_id seq_id  = ubatch->seq_id[i][s];\n                const int32_t      seq_idx = ubatch->seq_idx[seq_id];\n\n                for (int j = 0; j < n_seq_tokens; ++j) {\n                    data[seq_idx*n_tokens + i + j] = div[seq_idx];\n                }\n            }\n        }\n    }\n}\n\nvoid llm_graph_input_cls::set_input(const llama_ubatch * ubatch) {\n    const int64_t n_tokens     = ubatch->n_tokens;\n    const int64_t n_seqs_unq   = ubatch->n_seqs_unq;\n\n    if (cparams.embeddings && (\n        cparams.pooling_type == LLAMA_POOLING_TYPE_CLS  ||\n        cparams.pooling_type == LLAMA_POOLING_TYPE_RANK ||\n        cparams.pooling_type == LLAMA_POOLING_TYPE_LAST\n    )) {\n        GGML_ASSERT(cls);\n        GGML_ASSERT(ggml_backend_buffer_is_host(cls->buffer));\n\n        uint32_t * data = (uint32_t *) cls->data;\n        memset(cls->data, 0, n_seqs_unq*ggml_element_size(cls));\n\n        std::vector<int> target_pos(n_seqs_unq, -1);\n        std::vector<int> target_row(n_seqs_unq, -1);\n\n        const bool last = (\n             cparams.pooling_type == LLAMA_POOLING_TYPE_LAST ||\n            (cparams.pooling_type == LLAMA_POOLING_TYPE_RANK && arch == LLM_ARCH_QWEN3) // qwen3 reranking & embedding models use last token\n        );\n\n        for (int i = 0; i < n_tokens; ++i) {\n            const llama_pos pos = ubatch->pos[i];\n\n            for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {\n                const llama_seq_id seq_id  = ubatch->seq_id[i][s];\n                const int32_t      seq_idx = ubatch->seq_idx[seq_id];\n\n                if (\n                    (target_pos[seq_idx] == -1) ||\n                    ( last && pos >= target_pos[seq_idx]) ||\n                    (!last && pos <  target_pos[seq_idx])\n                ) {\n                    target_pos[seq_idx] = pos;\n                    target_row[seq_idx] = i;\n                }\n            }\n        }\n\n        for (int s = 0; s < n_seqs_unq; ++s) {\n            if (target_row[s] >= 0) {\n                data[s] = target_row[s];\n            }\n        }\n    }\n}\n\nvoid llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {\n    GGML_UNUSED(ubatch);\n\n    const int64_t n_rs = mctx->get_n_rs();\n\n    if (s_copy) {\n        GGML_ASSERT(ggml_backend_buffer_is_host(s_copy->buffer));\n        int32_t * data = (int32_t *) s_copy->data;\n\n        // assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n\n        for (uint32_t i = 0; i < n_rs; ++i) {\n            data[i] = mctx->s_copy(i);\n        }\n    }\n}\n\nbool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {\n    const auto * mctx = static_cast<const llama_memory_recurrent_context *>(params.mctx);\n\n    this->mctx = mctx;\n\n    bool res = true;\n\n    res &= s_copy->ne[0] == mctx->get_n_rs();\n\n    res &= s_copy_main->ne[0]  == params.ubatch.n_seqs;\n    res &= s_copy_extra->ne[0] == mctx->get_n_rs() - params.ubatch.n_seqs;\n\n    res &= head == mctx->get_head();\n    res &= rs_z == mctx->get_rs_z();\n\n    return res;\n}\n\nvoid llm_graph_input_cross_embd::set_input(const llama_ubatch * ubatch) {\n    GGML_UNUSED(ubatch);\n\n    if (cross_embd && !cross->v_embd.empty()) {\n        assert(cross_embd->type == GGML_TYPE_F32);\n\n        ggml_backend_tensor_set(cross_embd, cross->v_embd.data(), 0, ggml_nbytes(cross_embd));\n    }\n}\n\nstatic void print_mask(const float * data, int64_t n_tokens, int64_t n_kv, int64_t n_swa, llama_swa_type swa_type) {\n    LLAMA_LOG_DEBUG(\"%s: === Attention mask ===\\n\", __func__);\n    const char * swa_type_str = \"unknown\";\n\n    switch (swa_type) {\n        case LLAMA_SWA_TYPE_NONE:      swa_type_str = \"LLAMA_SWA_TYPE_NONE\"; break;\n        case LLAMA_SWA_TYPE_STANDARD:  swa_type_str = \"LLAMA_SWA_TYPE_STANDARD\"; break;\n        case LLAMA_SWA_TYPE_CHUNKED:   swa_type_str = \"LLAMA_SWA_TYPE_CHUNKED\"; break;\n        case LLAMA_SWA_TYPE_SYMMETRIC: swa_type_str = \"LLAMA_SWA_TYPE_SYMMETRIC\"; break;\n    };\n\n    LLAMA_LOG_DEBUG(\"%s: n_swa : %d, n_kv: %d, swq_type: %s\\n\", __func__, (int)n_swa, (int)n_kv, swa_type_str);\n    LLAMA_LOG_DEBUG(\"%s: '0' = can attend, '∞' = masked\\n\", __func__);\n    LLAMA_LOG_DEBUG(\"%s: Rows = query tokens, Columns = key/value tokens\\n\\n\", __func__);\n\n    LLAMA_LOG_DEBUG(\"    \");\n    for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) {\n        LLAMA_LOG_DEBUG(\"%2d\", j);\n    }\n    LLAMA_LOG_DEBUG(\"\\n\");\n\n    for (int i = 0; i < std::min((int64_t)20, n_tokens); ++i) {\n        LLAMA_LOG_DEBUG(\" %2d \", i);\n        for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) {\n            float val = data[i * n_kv + j];\n            if (val == -INFINITY) {\n                LLAMA_LOG_DEBUG(\" ∞\");\n            } else {\n                LLAMA_LOG_DEBUG(\" 0\");\n            }\n        }\n        LLAMA_LOG_DEBUG(\"\\n\");\n    }\n}\n\nvoid llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {\n    const int64_t n_kv     = ubatch->n_tokens;\n    const int64_t n_tokens = ubatch->n_tokens;\n\n    const auto fill_mask = [&](float * data, int n_swa, llama_swa_type swa_type) {\n        for (int i1 = 0; i1 < n_tokens; ++i1) {\n            const llama_seq_id s1 = ubatch->seq_id[i1][0];\n            const llama_pos    p1 = ubatch->pos[i1];\n\n            const uint64_t idst = i1*n_kv;\n\n            for (int i0 = 0; i0 < n_tokens; ++i0) {\n                const llama_seq_id s0 = ubatch->seq_id[i0][0];\n                const llama_pos p0    = ubatch->pos[i0];\n\n                // mask different sequences\n                if (s0 != s1) {\n                    continue;\n                }\n\n                // mask future tokens\n                if (cparams.causal_attn && p0 > p1) {\n                    continue;\n                }\n\n                // apply SWA if any\n                if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {\n                    continue;\n                }\n\n                data[idst + i0] = hparams.use_alibi ? -std::abs(p0 - p1) : 0.0f;\n            }\n        }\n    };\n\n    {\n        GGML_ASSERT(self_kq_mask);\n        GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));\n\n        float * data = (float *) self_kq_mask->data;\n\n        std::fill(data, data + ggml_nelements(self_kq_mask), -INFINITY);\n\n        fill_mask(data, 0, LLAMA_SWA_TYPE_NONE);\n\n        if (debug) {\n            print_mask(data, n_tokens, n_kv, 0, LLAMA_SWA_TYPE_NONE);\n        }\n    }\n\n    if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {\n        GGML_ASSERT(self_kq_mask_swa);\n        GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer));\n\n        float * data = (float *) self_kq_mask_swa->data;\n\n        std::fill(data, data + ggml_nelements(self_kq_mask_swa), -INFINITY);\n\n        fill_mask(data, hparams.n_swa, hparams.swa_type);\n\n        if (debug) {\n            print_mask(data, n_tokens, n_kv, hparams.n_swa, hparams.swa_type);\n        }\n    }\n}\n\nvoid llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {\n    mctx->set_input_k_idxs(self_k_idxs, ubatch);\n    mctx->set_input_v_idxs(self_v_idxs, ubatch);\n\n    mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);\n}\n\nbool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) {\n    const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);\n\n    this->mctx = mctx;\n\n    bool res = true;\n\n    res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;\n  //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there\n\n    res &= self_kq_mask->ne[0] == mctx->get_n_kv();\n    res &= self_kq_mask->ne[1] == params.ubatch.n_tokens;\n\n    return res;\n}\n\nvoid llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {\n    mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch);\n    mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch);\n\n    mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);\n\n    mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch);\n    mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch);\n\n    mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);\n}\n\nbool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {\n    const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);\n\n    this->mctx = mctx;\n\n    bool res = true;\n\n    res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;\n  //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there\n\n    res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;\n  //res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there\n\n    res &= self_kq_mask->ne[0] == mctx->get_base()->get_n_kv();\n    res &= self_kq_mask->ne[1] == params.ubatch.n_tokens;\n\n    res &= self_kq_mask_swa->ne[0] == mctx->get_swa()->get_n_kv();\n    res &= self_kq_mask_swa->ne[1] == params.ubatch.n_tokens;\n\n    return res;\n}\n\nvoid llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) {\n    GGML_ASSERT(cross_kq_mask);\n\n    const int64_t n_enc    = cross_kq_mask->ne[0];\n    const int64_t n_tokens = ubatch->n_tokens;\n\n    GGML_ASSERT(ggml_backend_buffer_is_host(cross_kq_mask->buffer));\n    GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing\n\n    float * data = (float *) cross_kq_mask->data;\n\n    for (int i = 0; i < n_tokens; ++i) {\n        for (int j = 0; j < n_enc; ++j) {\n            float f = -INFINITY;\n\n            for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {\n                const llama_seq_id seq_id = ubatch->seq_id[i][s];\n\n                if (cross->seq_ids_enc[j].find(seq_id) != cross->seq_ids_enc[j].end()) {\n                    f = 0.0f;\n                }\n            }\n\n            data[i*n_enc + j] = f;\n        }\n    }\n}\n\nvoid llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {\n    mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);\n    mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);\n\n    mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);\n\n    const int64_t n_rs = mctx->get_recr()->get_n_rs();\n\n    if (inp_rs->s_copy) {\n        GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer));\n        int32_t * data = (int32_t *) inp_rs->s_copy->data;\n\n        // assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n\n        for (uint32_t i = 0; i < n_rs; ++i) {\n            data[i] = mctx->get_recr()->s_copy(i);\n        }\n    }\n}\n\nbool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {\n    const auto * mctx = static_cast<const llama_memory_hybrid_context *>(params.mctx);\n\n    this->mctx = mctx;\n\n    bool res = true;\n\n    res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;\n  //res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there\n\n    res &= inp_attn->self_kq_mask->ne[0] == mctx->get_attn()->get_n_kv();\n    res &= inp_attn->self_kq_mask->ne[1] == params.ubatch.n_tokens;\n\n    res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();\n\n    res &= inp_rs->s_copy_main->ne[0]  == params.ubatch.n_seqs;\n    res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs;\n\n    res &= inp_rs->head == mctx->get_recr()->get_head();\n    res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();\n\n    return res;\n}\n\nvoid llm_graph_input_sampling::set_input(const llama_ubatch * ubatch) {\n    // set the inputs only for the active samplers in the current ubatch\n    std::unordered_set<llama_seq_id> active_samplers;\n    for (uint32_t i = 0; i < ubatch->n_tokens; i++) {\n        if (ubatch->output[i]) {\n            llama_seq_id seq_id = ubatch->seq_id[i][0];\n            active_samplers.insert(seq_id);\n        }\n    }\n\n    for (auto seq_id : active_samplers) {\n        if (samplers.find(seq_id) == samplers.end()) {\n            continue;\n        }\n\n        auto & sampler = samplers[seq_id];\n\n        if (sampler->iface->backend_set_input) {\n            sampler->iface->backend_set_input(sampler);\n        }\n    }\n}\n\nbool llm_graph_input_sampling::can_reuse(const llm_graph_params & params) {\n    if (samplers.size() != params.samplers.size()) {\n        return false;\n    }\n\n    for (const auto & [seq_id, sampler] : params.samplers) {\n        if (samplers[seq_id] != sampler) {\n            return false;\n        }\n    }\n\n    return true;\n}\n\n//\n// llm_graph_result\n//\n\nllm_graph_result::llm_graph_result(int64_t max_nodes) : max_nodes(max_nodes) {\n    reset();\n\n    const char * LLAMA_GRAPH_RESULT_DEBUG = getenv(\"LLAMA_GRAPH_RESULT_DEBUG\");\n    debug = LLAMA_GRAPH_RESULT_DEBUG ? atoi(LLAMA_GRAPH_RESULT_DEBUG) : 0;\n}\n\nint64_t llm_graph_result::get_max_nodes() const {\n    return max_nodes;\n}\n\nvoid llm_graph_result::reset() {\n    t_tokens      = nullptr;\n    t_logits      = nullptr;\n    t_embd        = nullptr;\n    t_embd_pooled = nullptr;\n    t_sampled.clear();\n    t_sampled_probs.clear();\n    t_sampled_logits.clear();\n    t_candidates.clear();\n\n    params = {};\n\n    inputs.clear();\n\n    buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));\n\n    ggml_init_params params = {\n        /*.mem_size   =*/ buf_compute_meta.size(),\n        /*.mem_buffer =*/ buf_compute_meta.data(),\n        /*.no_alloc   =*/ true,\n    };\n\n    ctx_compute.reset(ggml_init(params));\n\n    gf = ggml_new_graph_custom(ctx_compute.get(), max_nodes, false);\n}\n\nvoid llm_graph_result::set_inputs(const llama_ubatch * ubatch) {\n    for (auto & input : inputs) {\n        input->set_input(ubatch);\n    }\n}\n\nvoid llm_graph_result::set_outputs() {\n    if (t_logits != nullptr) {\n        ggml_set_output(t_logits);\n    }\n    if (t_embd != nullptr) {\n        ggml_set_output(t_embd);\n    }\n    if (t_embd_pooled != nullptr) {\n        ggml_set_output(t_embd_pooled);\n    }\n    for (auto & [seq_id, t] : t_sampled) {\n        if (t != nullptr) {\n            ggml_set_output(t);\n        }\n    }\n    for (auto & [seq_id, t] : t_sampled_probs) {\n        if (t != nullptr) {\n            ggml_set_output(t);\n        }\n    }\n    for (auto & [seq_id, t] : t_sampled_logits) {\n        if (t != nullptr) {\n            ggml_set_output(t);\n        }\n    }\n    for (auto & [seq_id, t] : t_candidates) {\n        if (t != nullptr) {\n            ggml_set_output(t);\n        }\n    }\n}\n\nbool llm_graph_result::can_reuse(const llm_graph_params & params) {\n    if (!this->params.allow_reuse(params)) {\n        if (debug > 1) {\n            LLAMA_LOG_DEBUG(\"%s: cannot reuse graph due to incompatible graph parameters\\n\", __func__);\n        }\n\n        return false;\n    }\n\n    if (debug > 1) {\n        LLAMA_LOG_DEBUG(\"%s: checking compatibility of %d inputs:\\n\", __func__, (int) inputs.size());\n    }\n\n    bool res = true;\n\n    for (auto & input : inputs) {\n        const bool cur = input->can_reuse(params);\n\n        if (debug > 1) {\n            LLAMA_LOG_DEBUG(\"%s: can_reuse = %d\\n\", \"placeholder\", cur);\n        }\n\n        res = res && cur;\n    }\n\n    if (debug > 0) {\n        LLAMA_LOG_DEBUG(\"%s: can reuse graph = %d\\n\", __func__, res);\n    }\n\n    return res;\n}\n\nllm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) {\n    inputs.emplace_back(std::move(input));\n    return inputs.back().get();\n}\n\nvoid llm_graph_result::set_params(const llm_graph_params & params) {\n    this->params = params;\n}\n\n//\n// llm_graph_context\n//\n\nllm_graph_context::llm_graph_context(const llm_graph_params & params) :\n    arch             (params.arch),\n    hparams          (params.hparams),\n    cparams          (params.cparams),\n    ubatch           (params.ubatch),\n    n_embd           (hparams.n_embd),\n    n_layer          (hparams.n_layer),\n    n_rot            (hparams.n_rot),\n    n_ctx            (cparams.n_ctx),\n    n_head           (hparams.n_head()),\n    n_head_kv        (hparams.n_head_kv()),\n    n_embd_head_k    (hparams.n_embd_head_k),\n    n_embd_k_gqa     (hparams.n_embd_k_gqa()),\n    n_embd_head_v    (hparams.n_embd_head_v),\n    n_embd_v_gqa     (hparams.n_embd_v_gqa()),\n    n_expert         (hparams.n_expert),\n    n_expert_used    (cparams.warmup ? hparams.n_expert : hparams.n_expert_used),\n    freq_base        (cparams.rope_freq_base),\n    freq_scale       (cparams.rope_freq_scale),\n    ext_factor       (cparams.yarn_ext_factor),\n    attn_factor      (cparams.yarn_attn_factor),\n    beta_fast        (cparams.yarn_beta_fast),\n    beta_slow        (cparams.yarn_beta_slow),\n    norm_eps         (hparams.f_norm_eps),\n    norm_rms_eps     (hparams.f_norm_rms_eps),\n    n_tokens         (ubatch.n_tokens),\n    n_outputs        (params.n_outputs),\n    n_ctx_orig       (cparams.n_ctx_orig_yarn),\n    pooling_type     (cparams.pooling_type),\n    rope_type        (hparams.rope_type),\n    sched            (params.sched),\n    backend_cpu      (params.backend_cpu),\n    cvec             (params.cvec),\n    loras            (params.loras),\n    mctx             (params.mctx),\n    cross            (params.cross),\n    samplers         (params.samplers),\n    cb_func          (params.cb),\n    res              (params.res),\n    ctx0             (res->get_ctx()),\n    gf               (res->get_gf()) {\n        res->set_params(params);\n    }\n\nvoid llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const {\n    if (cb_func) {\n        cb_func(ubatch, cur, name, il);\n    }\n}\n\nggml_tensor * llm_graph_context::build_cvec(\n         ggml_tensor * cur,\n                 int   il) const {\n    return cvec->apply_to(ctx0, cur, il);\n}\n\nggml_tensor * llm_graph_context::build_lora_mm(\n          ggml_tensor * w,\n          ggml_tensor * cur) const {\n    ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);\n\n    for (const auto & lora : *loras) {\n        llama_adapter_lora_weight * lw = lora.first->get_weight(w);\n        if (lw == nullptr) {\n            continue;\n        }\n\n        const float adapter_scale = lora.second;\n        const float scale = lw->get_scale(lora.first->alpha, adapter_scale);\n\n        ggml_tensor * ab_cur = ggml_mul_mat(\n                ctx0, lw->b,\n                ggml_mul_mat(ctx0, lw->a, cur)\n                );\n\n        ab_cur = ggml_scale(ctx0, ab_cur, scale);\n        res = ggml_add(ctx0, res, ab_cur);\n    }\n\n    return res;\n}\n\nggml_tensor * llm_graph_context::build_lora_mm_id(\n          ggml_tensor * w,   // ggml_tensor * as\n          ggml_tensor * cur, // ggml_tensor * b\n          ggml_tensor * ids) const {\n    ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);\n    for (const auto & lora : *loras) {\n        llama_adapter_lora_weight * lw = lora.first->get_weight(w);\n        if (lw == nullptr) {\n            continue;\n        }\n\n        const float alpha = lora.first->alpha;\n        const float rank  = (float) lw->b->ne[0];\n        const float scale = alpha ? lora.second * alpha / rank : lora.second;\n\n        ggml_tensor * ab_cur = ggml_mul_mat_id(\n                ctx0, lw->b,\n                ggml_mul_mat_id(ctx0, lw->a, cur, ids),\n                ids\n                );\n\n        ab_cur = ggml_scale(ctx0, ab_cur, scale);\n        res = ggml_add(ctx0, res, ab_cur);\n    }\n\n    return res;\n}\n\nggml_tensor * llm_graph_context::build_norm(\n         ggml_tensor * cur,\n         ggml_tensor * mw,\n         ggml_tensor * mb,\n       llm_norm_type   type,\n                 int   il) const {\n    switch (type) {\n        case LLM_NORM:       cur = ggml_norm    (ctx0, cur, hparams.f_norm_eps);     break;\n        case LLM_NORM_RMS:   cur = ggml_rms_norm(ctx0, cur, hparams.f_norm_rms_eps); break;\n        case LLM_NORM_GROUP:\n            {\n                cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, cur->ne[1]);\n                cur = ggml_group_norm(ctx0, cur, hparams.n_norm_groups, hparams.f_norm_group_eps);\n                cur = ggml_reshape_2d(ctx0, cur, cur->ne[0],    cur->ne[2]);\n            } break;\n    }\n\n    if (mw || mb) {\n        cb(cur, \"norm\", il);\n    }\n\n    if (mw) {\n        cur = ggml_mul(ctx0, cur, mw);\n        if (mb) {\n            cb(cur, \"norm_w\", il);\n        }\n    }\n\n    if (mb) {\n        cur = ggml_add(ctx0, cur, mb);\n    }\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_ffn(\n         ggml_tensor * cur,\n         ggml_tensor * up,\n         ggml_tensor * up_b,\n         ggml_tensor * up_s,\n         ggml_tensor * gate,\n         ggml_tensor * gate_b,\n         ggml_tensor * gate_s,\n         ggml_tensor * down,\n         ggml_tensor * down_b,\n         ggml_tensor * down_s,\n         ggml_tensor * act_scales,\n     llm_ffn_op_type   type_op,\n   llm_ffn_gate_type   type_gate,\n                 int   il) const {\n    ggml_tensor * tmp = up ? build_lora_mm(up, cur) : cur;\n    cb(tmp, \"ffn_up\", il);\n\n    if (up_b) {\n        tmp = ggml_add(ctx0, tmp, up_b);\n        cb(tmp, \"ffn_up_b\", il);\n    }\n\n    if (up_s) {\n        tmp = ggml_mul(ctx0, tmp, up_s);\n        cb(tmp, \"ffn_up_s\", il);\n    }\n\n    if (gate) {\n        switch (type_gate) {\n            case LLM_FFN_SEQ:\n                {\n                    cur = build_lora_mm(gate, tmp);\n                    cb(cur, \"ffn_gate\", il);\n                } break;\n            case LLM_FFN_PAR:\n                {\n                    cur = build_lora_mm(gate, cur);\n                    cb(cur, \"ffn_gate\", il);\n                } break;\n        }\n\n        if (gate_b) {\n            cur = ggml_add(ctx0, cur, gate_b);\n            cb(cur, \"ffn_gate_b\", il);\n        }\n\n        if (gate_s) {\n            cur = ggml_mul(ctx0, cur, gate_s);\n            cb(cur, \"ffn_gate_s\", il);\n        }\n\n    } else {\n        cur = tmp;\n    }\n\n    switch (type_op) {\n        case LLM_FFN_SILU:\n            if (gate && type_gate == LLM_FFN_PAR) {\n                cur = ggml_swiglu_split(ctx0, cur, tmp);\n                cb(cur, \"ffn_swiglu\", il);\n                type_gate = LLM_FFN_SEQ;\n            } else {\n                cur = ggml_silu(ctx0, cur);\n                cb(cur, \"ffn_silu\", il);\n            } break;\n        case LLM_FFN_GELU:\n            if (gate && type_gate == LLM_FFN_PAR) {\n                cur = ggml_geglu_split(ctx0, cur, tmp);\n                cb(cur, \"ffn_geglu\", il);\n                type_gate = LLM_FFN_SEQ;\n            } else {\n                cur = ggml_gelu(ctx0, cur);\n                cb(cur, \"ffn_gelu\", il);\n                if (act_scales != NULL) {\n                    cur = ggml_div(ctx0, cur, act_scales);\n                    cb(cur, \"ffn_act\", il);\n                }\n            } break;\n        case LLM_FFN_RELU:\n            if (gate && type_gate == LLM_FFN_PAR) {\n                cur = ggml_reglu_split(ctx0, cur, tmp);\n                cb(cur, \"ffn_reglu\", il);\n                type_gate = LLM_FFN_SEQ;\n            } else {\n                cur = ggml_relu(ctx0, cur);\n                cb(cur, \"ffn_relu\", il);\n            } break;\n        case LLM_FFN_RELU_SQR:\n            {\n                cur = ggml_relu(ctx0, cur);\n                cb(cur, \"ffn_relu\", il);\n\n                cur = ggml_sqr(ctx0, cur);\n                cb(cur, \"ffn_sqr(relu)\", il);\n            } break;\n        case LLM_FFN_SWIGLU:\n            {\n                cur = ggml_swiglu(ctx0, cur);\n                cb(cur, \"ffn_swiglu\", il);\n            } break;\n        case LLM_FFN_GEGLU:\n            {\n                cur = ggml_geglu(ctx0, cur);\n                cb(cur, \"ffn_geglu\", il);\n            } break;\n        case LLM_FFN_REGLU:\n            {\n                cur = ggml_reglu(ctx0, cur);\n                cb(cur, \"ffn_reglu\", il);\n            } break;\n        default:\n            GGML_ABORT(\"fatal error\");\n    }\n\n    if (gate && type_gate == LLM_FFN_PAR) {\n        cur = ggml_mul(ctx0, cur, tmp);\n        cb(cur, \"ffn_gate_par\", il);\n    }\n\n    if (down) {\n        cur = build_lora_mm(down, cur);\n        if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) {\n            // GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators\n            ggml_mul_mat_set_prec(cur, GGML_PREC_F32);\n        }\n    }\n\n    if (down_b) {\n        cb(cur, \"ffn_down\", il);\n    }\n\n    if (down_b) {\n        cur = ggml_add(ctx0, cur, down_b);\n    }\n\n    if (down_s) {\n        cur = ggml_mul(ctx0, cur, down_s);\n        cb(cur, \"ffn_down_s\", il);\n    }\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_moe_ffn(\n         ggml_tensor * cur,\n         ggml_tensor * gate_inp,\n         ggml_tensor * up_exps,\n         ggml_tensor * gate_exps,\n         ggml_tensor * down_exps,\n         ggml_tensor * exp_probs_b,\n             int64_t   n_expert,\n             int64_t   n_expert_used,\n     llm_ffn_op_type   type_op,\n                bool   norm_w,\n                bool   scale_w,\n               float   w_scale,\n         llama_expert_gating_func_type gating_op,\n                 int   il,\n         ggml_tensor * probs_in) const {\n    return build_moe_ffn(\n        cur,\n        gate_inp,  /* gate_inp_b  */ nullptr,\n        up_exps,   /* up_exps_b   */ nullptr,\n        gate_exps, /* gate_exps_b */ nullptr,\n        down_exps, /* down_exps_b */ nullptr,\n        exp_probs_b,\n        n_expert,\n        n_expert_used,\n        type_op,\n        norm_w,\n        scale_w,\n        w_scale,\n        gating_op,\n        il,\n        probs_in\n    );\n}\n\nggml_tensor * llm_graph_context::build_moe_ffn(\n         ggml_tensor * cur,\n         ggml_tensor * gate_inp,\n         ggml_tensor * gate_inp_b,\n         ggml_tensor * up_exps,\n         ggml_tensor * up_exps_b,\n         ggml_tensor * gate_exps,\n         ggml_tensor * gate_exps_b,\n         ggml_tensor * down_exps,\n         ggml_tensor * down_exps_b,\n         ggml_tensor * exp_probs_b,\n             int64_t   n_expert,\n             int64_t   n_expert_used,\n     llm_ffn_op_type   type_op,\n                bool   norm_w,\n                bool   scale_w,\n               float   w_scale,\n        llama_expert_gating_func_type gating_op,\n                 int   il,\n         ggml_tensor * probs_in) const {\n    const int64_t n_embd   = cur->ne[0];\n    const int64_t n_tokens = cur->ne[1];\n    const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN\n\n    ggml_tensor * logits = nullptr;\n\n    if (probs_in == nullptr) {\n        logits = build_lora_mm(gate_inp, cur); // [n_expert, n_tokens]\n        cb(logits, \"ffn_moe_logits\", il);\n    } else {\n        logits = probs_in;\n    }\n\n    if (gate_inp_b) {\n        logits = ggml_add(ctx0, logits, gate_inp_b);\n        cb(logits, \"ffn_moe_logits_biased\", il);\n    }\n\n    ggml_tensor * probs = nullptr;\n    switch (gating_op) {\n        case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX:\n            {\n                probs = ggml_soft_max(ctx0, logits); // [n_expert, n_tokens]\n            } break;\n        case LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID:\n            {\n                probs = ggml_sigmoid(ctx0, logits); // [n_expert, n_tokens]\n            } break;\n        case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT:\n            {\n                probs = logits; // [n_expert, n_tokens]\n            } break;\n        default:\n            GGML_ABORT(\"fatal error\");\n    }\n    cb(probs, \"ffn_moe_probs\", il);\n\n    // add experts selection bias - introduced in DeepSeek V3\n    // leave probs unbiased as it's later used to get expert weights\n    ggml_tensor * selection_probs = probs;\n    if (exp_probs_b != nullptr) {\n        selection_probs = ggml_add(ctx0, probs, exp_probs_b);\n        cb(selection_probs, \"ffn_moe_probs_biased\", il);\n    }\n\n    // llama4 doesn't have exp_probs_b, and sigmoid is only used after top_k\n    // see: https://github.com/meta-llama/llama-models/blob/699a02993512fb36936b1b0741e13c06790bcf98/models/llama4/moe.py#L183-L198\n    if (arch == LLM_ARCH_LLAMA4) {\n        selection_probs = logits;\n    }\n\n    if (arch == LLM_ARCH_GROVEMOE) {\n        selection_probs = ggml_sigmoid(ctx0, logits); // [n_expert, n_tokens]\n        cb(selection_probs, \"ffn_moe_probs_biased\", il);\n    }\n\n    // select top n_group_used expert groups\n    // https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/e815299b0bcbac849fa540c768ef21845365c9eb/modeling_deepseek.py#L440-L457\n    if (hparams.n_expert_groups > 1 && n_tokens > 0) {\n        const int64_t n_exp_per_group = n_expert / hparams.n_expert_groups;\n\n        // organize experts into n_expert_groups\n        ggml_tensor * selection_groups = ggml_reshape_3d(ctx0, selection_probs, n_exp_per_group, hparams.n_expert_groups, n_tokens); // [n_exp_per_group, n_expert_groups, n_tokens]\n\n        ggml_tensor * group_scores = ggml_argsort_top_k(ctx0, selection_groups, 2); // [2, n_expert_groups, n_tokens]\n        group_scores = ggml_get_rows(ctx0, ggml_reshape_4d(ctx0, selection_groups, 1, selection_groups->ne[0], selection_groups->ne[1], selection_groups->ne[2]), group_scores); // [1, 2, n_expert_groups, n_tokens]\n\n        // get top n_group_used expert groups\n        group_scores = ggml_sum_rows(ctx0, ggml_reshape_3d(ctx0, group_scores, group_scores->ne[1], group_scores->ne[2], group_scores->ne[3])); // [1, n_expert_groups, n_tokens]\n        group_scores = ggml_reshape_2d(ctx0, group_scores, group_scores->ne[1], group_scores->ne[2]); // [n_expert_groups, n_tokens]\n\n        ggml_tensor * expert_groups = ggml_argsort_top_k(ctx0, group_scores, hparams.n_group_used); // [n_group_used, n_tokens]\n        cb(expert_groups, \"ffn_moe_group_topk\", il);\n\n        // mask out the other groups\n        selection_probs = ggml_get_rows(ctx0, selection_groups, expert_groups); // [n_exp_per_group, n_group_used, n_tokens]\n        selection_probs = ggml_set_rows(ctx0, ggml_fill(ctx0, selection_groups, -INFINITY), selection_probs, expert_groups); // [n_exp_per_group, n_expert_groups, n_tokens]\n        selection_probs = ggml_reshape_2d(ctx0, selection_probs, n_expert, n_tokens); // [n_expert, n_tokens]\n        cb(selection_probs, \"ffn_moe_probs_masked\", il);\n    }\n\n    // select experts\n    ggml_tensor * selected_experts = ggml_argsort_top_k(ctx0, selection_probs, n_expert_used); // [n_expert_used, n_tokens]\n    cb(selected_experts->src[0], \"ffn_moe_argsort\", il);\n    cb(selected_experts, \"ffn_moe_topk\", il);\n\n    if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) {\n        // TODO: Use scalar div instead when/if implemented\n        ggml_tensor * f_sel = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32);\n        selected_experts = ggml_cast(ctx0, ggml_scale(ctx0, f_sel, 1.0f / float(hparams.n_group_experts)), GGML_TYPE_I32);\n        probs = ggml_reshape_3d(ctx0, probs, 1, hparams.n_expert, n_tokens);\n    } else {\n        probs = ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens);\n    }\n\n    ggml_tensor * weights = ggml_get_rows(ctx0, probs, selected_experts); // [1, n_expert_used, n_tokens]\n    cb(weights, \"ffn_moe_weights\", il);\n\n\n    if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT) {\n        weights = ggml_reshape_2d(ctx0, weights, n_expert_used, n_tokens);\n        weights = ggml_soft_max(ctx0, weights); // [n_expert_used, n_tokens]\n        weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens);\n        cb(weights, \"ffn_moe_weights_softmax\", il);\n    }\n\n    if (norm_w) {\n        weights = ggml_reshape_2d(ctx0, weights, n_expert_used, n_tokens);\n\n        ggml_tensor * weights_sum = ggml_sum_rows(ctx0, weights); // [1, n_tokens]\n        cb(weights_sum, \"ffn_moe_weights_sum\", il);\n\n        // Avoid division by zero, clamp to smallest number representable by F16\n        weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY);\n        cb(weights_sum, \"ffn_moe_weights_sum_clamped\", il);\n\n        weights = ggml_div(ctx0, weights, weights_sum); // [n_expert_used, n_tokens]\n        cb(weights, \"ffn_moe_weights_norm\", il);\n\n        weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens);\n    }\n    if (scale_w) {\n        weights = ggml_scale(ctx0, weights, w_scale);\n        cb(weights, \"ffn_moe_weights_scaled\", il);\n    }\n\n    //call early so that topk-moe can be used\n    ggml_build_forward_expand(gf, weights);\n\n    cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);\n\n    if (weight_before_ffn) {\n        // repeat cur to [n_embd, n_expert_used, n_tokens]\n        ggml_tensor * repeated = ggml_repeat_4d(ctx0, cur, n_embd, n_expert_used, n_tokens, 1);\n        cur = ggml_mul(ctx0, repeated, weights);\n        cb(cur, \"ffn_moe_weighted\", il);\n    }\n\n    ggml_tensor * up = build_lora_mm_id(up_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]\n    cb(up, \"ffn_moe_up\", il);\n\n    if (up_exps_b) {\n        up = ggml_add_id(ctx0, up, up_exps_b, selected_experts);\n        cb(up, \"ffn_moe_up_biased\", il);\n    }\n\n    ggml_tensor * experts = nullptr;\n    if (gate_exps) {\n        cur = build_lora_mm_id(gate_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]\n        cb(cur, \"ffn_moe_gate\", il);\n    } else {\n        cur = up;\n    }\n\n    if (gate_exps_b) {\n        cur = ggml_add_id(ctx0, cur, gate_exps_b, selected_experts);\n        cb(cur, \"ffn_moe_gate_biased\", il);\n    }\n\n    switch (type_op) {\n        case LLM_FFN_SILU:\n            if (gate_exps) {\n                cur = ggml_swiglu_split(ctx0, cur, up);\n                cb(cur, \"ffn_moe_swiglu\", il);\n            } else {\n                cur = ggml_silu(ctx0, cur);\n                cb(cur, \"ffn_moe_silu\", il);\n            } break;\n        case LLM_FFN_GELU:\n            if (gate_exps) {\n                cur = ggml_geglu_split(ctx0, cur, up);\n                cb(cur, \"ffn_moe_geglu\", il);\n            } else {\n                cur = ggml_gelu(ctx0, cur);\n                cb(cur, \"ffn_moe_gelu\", il);\n            } break;\n        case LLM_FFN_SWIGLU_OAI_MOE:\n            {\n                // TODO: move to hparams?\n                constexpr float alpha = 1.702f;\n                constexpr float limit = 7.0f;\n                cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit);\n                cb(cur, \"ffn_moe_swiglu_oai\", il);\n            } break;\n        case LLM_FFN_RELU:\n            if (gate_exps) {\n                cur = ggml_reglu_split(ctx0, cur, up);\n                cb(cur, \"ffn_moe_reglu\", il);\n            } else {\n                cur = ggml_relu(ctx0, cur);\n                cb(cur, \"ffn_moe_relu\", il);\n            } break;\n        case LLM_FFN_RELU_SQR:\n            if (gate_exps) {\n                // TODO: add support for gated squared relu\n                GGML_ABORT(\"fatal error: gated squared relu not implemented\");\n            } else {\n                cur = ggml_relu(ctx0, cur);\n                cur = ggml_sqr(ctx0, cur);\n                cb(cur, \"ffn_moe_relu_sqr\", il);\n            } break;\n        default:\n            GGML_ABORT(\"fatal error\");\n    }\n\n    experts = build_lora_mm_id(down_exps, cur, selected_experts); // [n_embd, n_expert_used, n_tokens]\n    cb(experts, \"ffn_moe_down\", il);\n\n    if (down_exps_b) {\n        experts = ggml_add_id(ctx0, experts, down_exps_b, selected_experts);\n        cb(experts, \"ffn_moe_down_biased\", il);\n    }\n\n    if (!weight_before_ffn) {\n        experts = ggml_mul(ctx0, experts, weights);\n        cb(cur, \"ffn_moe_weighted\", il);\n    }\n\n    ggml_tensor * cur_experts[LLAMA_MAX_EXPERTS] = { nullptr };\n\n    assert(n_expert_used > 0);\n\n    // order the views before the adds\n    for (uint32_t i = 0; i < hparams.n_expert_used; ++i) {\n        cur_experts[i] = ggml_view_2d(ctx0, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]);\n\n        ggml_build_forward_expand(gf, cur_experts[i]);\n    }\n\n    // aggregate experts\n    // note: here we explicitly use hparams.n_expert_used instead of n_expert_used\n    //       to avoid potentially a large number of add nodes during warmup\n    //       ref: https://github.com/ggml-org/llama.cpp/pull/14753\n    ggml_tensor * moe_out = cur_experts[0];\n\n    for (uint32_t i = 1; i < hparams.n_expert_used; ++i) {\n        moe_out = ggml_add(ctx0, moe_out, cur_experts[i]);\n    }\n\n    if (hparams.n_expert_used == 1) {\n        // avoid returning a non-contiguous tensor\n        moe_out = ggml_cont(ctx0, moe_out);\n    }\n\n    cb(moe_out, \"ffn_moe_out\", il);\n\n    return moe_out;\n}\n\n// input embeddings with optional lora\nggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {\n    const int64_t n_embd = hparams.n_embd_inp();\n\n    auto inp = std::make_unique<llm_graph_input_embd>();\n\n    ggml_tensor * cur = nullptr;\n\n    if (ubatch.token) {\n        inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);\n        //cb(inp->tokens, \"inp_tokens\", -1);\n        ggml_set_input(inp->tokens);\n        res->t_tokens = inp->tokens;\n\n        cur = ggml_get_rows(ctx0, tok_embd, inp->tokens);\n\n        // apply lora for embedding tokens if needed\n        for (const auto & lora : *loras) {\n            llama_adapter_lora_weight * lw = lora.first->get_weight(tok_embd);\n            if (lw == nullptr) {\n                continue;\n            }\n\n            const float adapter_scale = lora.second;\n            const float scale = lw->get_scale(lora.first->alpha, adapter_scale);\n\n            ggml_tensor * inpL_delta = ggml_scale(ctx0, ggml_mul_mat(\n                        ctx0, lw->b, // non-transposed lora_b\n                        ggml_get_rows(ctx0, lw->a, inp->tokens)\n                        ), scale);\n\n            cur = ggml_add(ctx0, cur, inpL_delta);\n        }\n    } else {\n        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, ubatch.n_tokens);\n        ggml_set_input(inp->embd);\n\n        cur = inp->embd;\n    }\n\n    // For Granite architecture\n    if (hparams.f_embedding_scale != 0.0f) {\n        cur = ggml_scale(ctx0, cur, hparams.f_embedding_scale);\n    }\n\n    cb(cur, \"inp_embd\", -1);\n\n    res->add_input(std::move(inp));\n\n    // make sure the produced embeddings are immediately materialized in the ggml graph\n    // ref: https://github.com/ggml-org/llama.cpp/pull/18599\n    ggml_build_forward_expand(gf, cur);\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_pos() const {\n    auto inp = std::make_unique<llm_graph_input_pos>(hparams.n_pos_per_embd());\n\n    auto & cur = inp->pos;\n\n    cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t)n_tokens*hparams.n_pos_per_embd());\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_attn_scale() const {\n    auto inp = std::make_unique<llm_graph_input_attn_temp>(hparams.n_attn_temp_floor_scale, hparams.f_attn_temp_scale, hparams.f_attn_temp_offset);\n\n    auto & cur = inp->attn_scale;\n\n    // this need to be 1x1xN for broadcasting\n    cur = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_out_ids() const {\n    // note: when all tokens are output, we could skip this optimization to spare the ggml_get_rows() calls,\n    //       but this would make the graph topology depend on the number of output tokens, which can interere with\n    //       features that require constant topology such as pipline parallelism\n    //       ref: https://github.com/ggml-org/llama.cpp/pull/14275#issuecomment-2987424471\n    //if (n_outputs < n_tokens) {\n    //    return nullptr;\n    //}\n\n    auto inp = std::make_unique<llm_graph_input_out_ids>(hparams, cparams, n_outputs);\n\n    auto & cur = inp->out_ids;\n\n    cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_outputs);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_mean() const {\n    auto inp = std::make_unique<llm_graph_input_mean>(cparams);\n\n    auto & cur = inp->mean;\n\n    cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, ubatch.n_seqs_unq);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_cls() const {\n    auto inp = std::make_unique<llm_graph_input_cls>(cparams, arch);\n\n    auto & cur = inp->cls;\n\n    cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_seqs_unq);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_cross_embd() const {\n    auto inp = std::make_unique<llm_graph_input_cross_embd>(cross);\n\n    auto & cur = inp->cross_embd;\n\n    // if we have the output embeddings from the encoder, use them directly\n    // TODO: needs more work to be correct, for now just use the tensor shape\n    //if (cross->t_embd) {\n    //    cur = ggml_view_tensor(ctx0, cross->t_embd);\n\n    //    return cur;\n    //}\n\n    const auto n_embd = !cross->v_embd.empty() ? cross->n_embd : hparams.n_embd_inp();\n    const auto n_enc  = !cross->v_embd.empty() ? cross->n_enc : hparams.n_ctx_train;\n\n    cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_enc);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_pos_bucket_enc() const {\n    auto inp = std::make_unique<llm_graph_input_pos_bucket>(hparams);\n\n    auto & cur = inp->pos_bucket;\n\n    cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_tokens, n_tokens);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_inp_pos_bucket_dec() const {\n    const auto * mctx_cur = static_cast<const llama_kv_cache_context *>(mctx);\n\n    auto inp = std::make_unique<llm_graph_input_pos_bucket_kv>(hparams, mctx_cur);\n\n    const auto n_kv = mctx_cur->get_n_kv();\n\n    auto & cur = inp->pos_bucket;\n\n    cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_tokens);\n    ggml_set_input(cur);\n\n    res->add_input(std::move(inp));\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_pos_bias(ggml_tensor * pos_bucket, ggml_tensor * attn_rel_b) const {\n    ggml_tensor * pos_bucket_1d = ggml_reshape_1d(ctx0, pos_bucket, pos_bucket->ne[0] * pos_bucket->ne[1]);\n    cb(pos_bucket_1d, \"pos_bucket_1d\", -1);\n\n    ggml_tensor * pos_bias = ggml_get_rows(ctx0, attn_rel_b, pos_bucket_1d);\n\n    pos_bias = ggml_reshape_3d(ctx0, pos_bias, pos_bias->ne[0], pos_bucket->ne[0], pos_bucket->ne[1]);\n    pos_bias = ggml_permute   (ctx0, pos_bias, 2, 0, 1, 3);\n    pos_bias = ggml_cont      (ctx0, pos_bias);\n\n    cb(pos_bias, \"pos_bias\", -1);\n\n    return pos_bias;\n}\n\nggml_tensor * llm_graph_context::build_attn_mha(\n         ggml_tensor * q,\n         ggml_tensor * k,\n         ggml_tensor * v,\n         ggml_tensor * kq_b,\n         ggml_tensor * kq_mask,\n         ggml_tensor * sinks,\n         ggml_tensor * v_mla,\n               float   kq_scale,\n                 int   il) const {\n    const bool v_trans = v->nb[1] > v->nb[2];\n\n    // split the batch into streams if needed\n    const auto n_stream = k->ne[3];\n\n    q = ggml_view_4d(ctx0, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, q->nb[1], q->nb[2], q->nb[3]/n_stream, 0);\n\n    q = ggml_permute(ctx0, q, 0, 2, 1, 3);\n    k = ggml_permute(ctx0, k, 0, 2, 1, 3);\n    v = ggml_permute(ctx0, v, 0, 2, 1, 3);\n\n    ggml_tensor * cur;\n\n    if (cparams.flash_attn && kq_b == nullptr) {\n        GGML_ASSERT(kq_b == nullptr && \"Flash attention does not support KQ bias yet\");\n\n        if (v_trans) {\n            v = ggml_transpose(ctx0, v);\n        }\n\n        // this can happen when KV cache is not used (e.g. an embedding model with non-causal attn)\n        if (k->type == GGML_TYPE_F32) {\n            k = ggml_cast(ctx0, k, GGML_TYPE_F16);\n        }\n\n        if (v->type == GGML_TYPE_F32) {\n            v = ggml_cast(ctx0, v, GGML_TYPE_F16);\n        }\n\n        cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,\n                                  hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);\n        cb(cur, LLAMA_TENSOR_NAME_FATTN, il);\n\n        ggml_flash_attn_ext_add_sinks(cur, sinks);\n        ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);\n\n        if (v_mla) {\n#if 0\n            // v_mla can be applied as a matrix-vector multiplication with broadcasting across dimension 3 == n_tokens.\n            // However, the code is optimized for dimensions 0 and 1 being large, so this is ineffient.\n            cur = ggml_reshape_4d(ctx0, cur, v_mla->ne[0], 1, n_head, n_tokens);\n            cur = ggml_mul_mat(ctx0, v_mla, cur);\n#else\n            // It's preferable to do the calculation as a matrix-matrix multiplication with n_tokens in dimension 1.\n            // The permutations are noops and only change how the tensor data is interpreted.\n            cur = ggml_permute(ctx0, cur, 0, 2, 1, 3);\n            cur = ggml_mul_mat(ctx0, v_mla, cur);\n            cb(cur, \"fattn_mla\", il);\n            cur = ggml_permute(ctx0, cur, 0, 2, 1, 3);\n            cur = ggml_cont(ctx0, cur); // Needed because ggml_reshape_2d expects contiguous inputs.\n#endif\n        }\n\n        cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]);\n    } else {\n        ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);\n        cb(kq, \"kq\", il);\n\n        // note: this op tends to require high floating point range\n        //       while for some models F16 is enough, for others it is not, so we default to F32 here\n        ggml_mul_mat_set_prec(kq, GGML_PREC_F32);\n\n        if (arch == LLM_ARCH_GROK) {\n            // need to do the following:\n            // multiply by attn_output_multiplier\n            // and then :\n            // kq = 30 * tanh(kq / 30)\n            // before the softmax below\n\n            kq = ggml_tanh(ctx0, ggml_scale(ctx0, kq, hparams.f_attn_out_scale / hparams.f_attn_logit_softcapping));\n            cb(kq, \"kq_tanh\", il);\n            kq = ggml_scale(ctx0, kq, hparams.f_attn_logit_softcapping);\n            cb(kq, \"kq_scaled\", il);\n        }\n\n        if (hparams.attn_soft_cap) {\n            kq = ggml_scale(ctx0, kq, 1.0f / hparams.f_attn_logit_softcapping);\n            cb(kq, \"kq_scaled_1\", il);\n            kq = ggml_tanh (ctx0, kq);\n            cb(kq, \"kq_tanh\", il);\n            kq = ggml_scale(ctx0, kq, hparams.f_attn_logit_softcapping);\n            cb(kq, \"kq_scaled_2\", il);\n        }\n\n        if (kq_b) {\n            kq = ggml_add(ctx0, kq, kq_b);\n            cb(kq, \"kq_plus_kq_b\", il);\n        }\n\n        kq = ggml_soft_max_ext(ctx0, kq, kq_mask, kq_scale, hparams.f_max_alibi_bias);\n        ggml_soft_max_add_sinks(kq, sinks);\n        cb(kq, \"kq_soft_max\", il);\n\n        if (!v_trans) {\n            // note: avoid this branch\n            v = ggml_cont(ctx0, ggml_transpose(ctx0, v));\n            cb(v, \"v_cont\", il);\n        }\n\n        ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq);\n        cb(kqv, \"kqv\", il);\n\n        // for MLA with the absorption optimization, we need to \"decompress\" from MQA back to MHA\n        if (v_mla) {\n            kqv = ggml_mul_mat(ctx0, v_mla, kqv);\n            cb(kqv, \"kqv_mla\", il);\n        }\n\n        cur = ggml_permute(ctx0, kqv, 0, 2, 1, 3);\n\n        // recombine streams\n        cur = ggml_cont_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]);\n\n        if (!cparams.offload_kqv) {\n            // all nodes between the KV store and the attention output are run on the CPU\n            ggml_backend_sched_set_tensor_backend(sched, cur, backend_cpu);\n        }\n    }\n\n    ggml_build_forward_expand(gf, cur);\n\n    return cur;\n}\n\nllm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() const {\n    auto inp = std::make_unique<llm_graph_input_attn_no_cache>(hparams, cparams);\n\n    // note: there is no KV cache, so the number of KV values is equal to the number of tokens in the batch\n    inp->self_kq_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_tokens, n_tokens, 1, 1);\n    ggml_set_input(inp->self_kq_mask);\n\n    inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;\n\n    if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {\n        inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_tokens, n_tokens, 1, 1);\n        ggml_set_input(inp->self_kq_mask_swa);\n\n        inp->self_kq_mask_swa_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask_swa, GGML_TYPE_F16) : inp->self_kq_mask_swa;\n    } else {\n        inp->self_kq_mask_swa     = nullptr;\n        inp->self_kq_mask_swa_cnv = nullptr;\n    }\n\n    return (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));\n}\n\nggml_tensor * llm_graph_context::build_attn(\n        llm_graph_input_attn_no_cache * inp,\n        ggml_tensor * wo,\n        ggml_tensor * wo_b,\n        ggml_tensor * q_cur,\n        ggml_tensor * k_cur,\n        ggml_tensor * v_cur,\n        ggml_tensor * kq_b,\n        ggml_tensor * sinks,\n        ggml_tensor * v_mla,\n            float     kq_scale,\n            int       il) const {\n    GGML_UNUSED(n_tokens);\n\n    // these nodes are added to the graph together so that they are not reordered\n    // by doing so, the number of splits in the graph is reduced\n    ggml_build_forward_expand(gf, q_cur);\n    ggml_build_forward_expand(gf, k_cur);\n    ggml_build_forward_expand(gf, v_cur);\n\n    const bool is_swa = hparams.is_swa(il);\n\n    const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();\n\n    // [TAG_NO_CACHE_PAD]\n    // TODO: if ubatch.equal_seqs() == true, we can split the three tensors below into ubatch.n_seqs_unq streams\n    //       but it might not be worth it: https://github.com/ggml-org/llama.cpp/pull/15636\n    //assert(!ubatch.equal_seqs() || (k_cur->ne[3] == 1 && k_cur->ne[3] == ubatch.n_seqs_unq));\n\n    ggml_tensor * q = q_cur;\n    ggml_tensor * k = k_cur;\n    ggml_tensor * v = v_cur;\n\n    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);\n    cb(cur, \"kqv_out\", il);\n\n    if (wo) {\n        cur = build_lora_mm(wo, cur);\n    }\n\n    if (wo_b) {\n        //cb(cur, \"kqv_wo\", il);\n    }\n\n    if (wo_b) {\n        cur = ggml_add(ctx0, cur, wo_b);\n    }\n\n    return cur;\n}\n\nstatic std::unique_ptr<llm_graph_input_attn_kv> build_attn_inp_kv_impl(\n           ggml_context * ctx0,\n     const llama_ubatch & ubatch,\n    const llama_hparams & hparams,\n    const llama_cparams & cparams,\n    const llama_kv_cache_context * mctx_cur) {\n\n    auto inp = std::make_unique<llm_graph_input_attn_kv>(hparams, cparams, mctx_cur);\n\n    {\n        GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && \"Use llama_kv_cache_iswa for SWA\");\n\n        const auto n_kv     = mctx_cur->get_n_kv();\n        const auto n_tokens = ubatch.n_tokens;\n        const auto n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq;\n\n        inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch);\n        inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch);\n\n        inp->self_kq_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tokens/n_stream, 1, n_stream);\n        ggml_set_input(inp->self_kq_mask);\n\n        inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;\n    }\n\n    return inp;\n}\n\nllm_graph_input_attn_kv * llm_graph_context::build_attn_inp_kv() const {\n    const auto * mctx_cur = static_cast<const llama_kv_cache_context *>(mctx);\n\n    auto inp = build_attn_inp_kv_impl(ctx0, ubatch, hparams, cparams, mctx_cur);\n\n    return (llm_graph_input_attn_kv *) res->add_input(std::move(inp));\n}\n\nggml_tensor * llm_graph_context::build_attn(\n        llm_graph_input_attn_kv * inp,\n        ggml_tensor * wo,\n        ggml_tensor * wo_b,\n        ggml_tensor * q_cur,\n        ggml_tensor * k_cur,\n        ggml_tensor * v_cur,\n        ggml_tensor * kq_b,\n        ggml_tensor * sinks,\n        ggml_tensor * v_mla,\n            float     kq_scale,\n            int       il) const {\n    // these nodes are added to the graph together so that they are not reordered\n    // by doing so, the number of splits in the graph is reduced\n    // expand k later to enable rope fusion which directly writes into k-v cache\n    ggml_build_forward_expand(gf, q_cur);\n    ggml_build_forward_expand(gf, v_cur);\n    ggml_build_forward_expand(gf, k_cur);\n\n    const auto * mctx_cur = inp->mctx;\n\n    // store to KV cache\n    {\n        const auto & k_idxs = inp->get_k_idxs();\n        const auto & v_idxs = inp->get_v_idxs();\n\n        ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));\n        ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));\n    }\n\n    const auto & kq_mask = inp->get_kq_mask();\n\n    ggml_tensor * q = q_cur;\n    ggml_tensor * k = mctx_cur->get_k(ctx0, il);\n    ggml_tensor * v = mctx_cur->get_v(ctx0, il);\n\n    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);\n    cb(cur, \"kqv_out\", il);\n\n    if (wo) {\n        cur = build_lora_mm(wo, cur);\n        if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) {\n            // GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators\n            ggml_mul_mat_set_prec(cur, GGML_PREC_F32);\n        }\n    }\n\n    if (wo_b) {\n        cur = ggml_add(ctx0, cur, wo_b);\n    }\n\n    return cur;\n}\n\nggml_tensor * llm_graph_context::build_attn(\n        llm_graph_input_attn_kv_iswa * inp,\n        ggml_tensor * wo,\n        ggml_tensor * wo_b,\n        ggml_tensor * q_cur,\n        ggml_tensor * k_cur,\n        ggml_tensor * v_cur,\n        ggml_tensor * kq_b,\n        ggml_tensor * sinks,\n        ggml_tensor * v_mla,\n            float     kq_scale,\n            int       il) const {\n    // these nodes are added to the graph together so that they are not reordered\n    // by doing so, the number of splits in the graph is reduced\n    ggml_build_forward_expand(gf, q_cur);\n\n    if (k_cur) {\n        ggml_build_forward_expand(gf, k_cur);\n    }\n\n    if (v_cur) {\n        ggml_build_forward_expand(gf, v_cur);\n    }\n\n    const auto * mctx_iswa = inp->mctx;\n\n    const bool is_swa = hparams.is_swa(il);\n\n    const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();\n\n    // optionally store to KV cache\n    if (k_cur) {\n        const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();\n\n        ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));\n    }\n\n    if (v_cur) {\n        const auto & v_idxs = is_swa ? inp->get_v_idxs_swa() : inp->get_v_idxs();\n\n        ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));\n    }\n\n    const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();\n\n    ggml_tensor * q = q_cur;\n    ggml_tensor * k = mctx_cur->get_k(ctx0, il);\n    ggml_tensor * v = mctx_cur->get_v(ctx0, il);\n\n    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);\n    cb(cur, \"kqv_out\", il);\n\n    if (wo) {\n        cur = build_lora_mm(wo, cur);\n    }\n\n    if (wo_b) {\n        //cb(cur, \"kqv_wo\", il);\n    }\n\n    if (wo_b) {\n        cur = ggml_add(ctx0, cur, wo_b);\n    }\n\n    return cur;\n}\n\nllm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const {\n    auto inp = std::make_unique<llm_graph_input_attn_cross>(cross);\n\n    const int32_t n_enc = !cross->v_embd.empty() ? cross->n_enc : hparams.n_ctx_train;\n\n    inp->cross_kq_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_enc, n_tokens, 1, 1);\n    ggml_set_input(inp->cross_kq_mask);\n\n    inp->cross_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->cross_kq_mask, GGML_TYPE_F16) : inp->cross_kq_mask;\n\n    return (llm_graph_input_attn_cross *) res->add_input(std::move(inp));\n}\n\nggml_tensor * llm_graph_context::build_attn(\n        llm_graph_input_attn_cross * inp,\n        ggml_tensor * wo,\n        ggml_tensor * wo_b,\n        ggml_tensor * q_cur,\n        ggml_tensor * k_cur,\n        ggml_tensor * v_cur,\n        ggml_tensor * kq_b,\n        ggml_tensor * sinks,\n        ggml_tensor * v_mla,\n            float     kq_scale,\n            int       il) const {\n    // these nodes are added to the graph together so that they are not reordered\n    // by doing so, the number of splits in the graph is reduced\n    ggml_build_forward_expand(gf, q_cur);\n    ggml_build_forward_expand(gf, k_cur);\n    ggml_build_forward_expand(gf, v_cur);\n\n    const auto & kq_mask = inp->get_kq_mask_cross();\n\n    ggml_tensor * q = q_cur;\n    ggml_tensor * k = k_cur;\n    ggml_tensor * v = v_cur;\n\n    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);\n    cb(cur, \"kqv_out\", il);\n\n    if (wo) {\n        cur = build_lora_mm(wo, cur);\n    }\n\n    if (wo_b) {\n        //cb(cur, \"kqv_wo\", il);\n    }\n\n    if (wo_b) {\n        cur = ggml_add(ctx0, cur, wo_b);\n    }\n\n    return cur;\n}\n\n// TODO: maybe separate the inner implementation into a separate function\n//       like with the non-sliding window equivalent\n//       once sliding-window hybrid caches are a thing.\nllm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const {\n    const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);\n\n    auto inp = std::make_unique<llm_graph_input_attn_kv_iswa>(hparams, cparams, mctx_cur);\n\n    const auto n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq;\n\n    {\n        const auto n_kv = mctx_cur->get_base()->get_n_kv();\n\n        inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);\n        inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch);\n\n        inp->self_kq_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tokens/n_stream, 1, n_stream);\n        ggml_set_input(inp->self_kq_mask);\n        ggml_set_name(inp->self_kq_mask, \"self_kq_mask\");\n\n        inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;\n        ggml_set_name(inp->self_kq_mask_cnv, \"self_kq_mask_cnv\");\n    }\n\n    {\n        GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && \"Use llama_kv_cache for non-SWA\");\n\n        const auto n_kv = mctx_cur->get_swa()->get_n_kv();\n\n        inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);\n        inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch);\n\n        inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tokens/n_stream, 1, n_stream);\n        ggml_set_input(inp->self_kq_mask_swa);\n        ggml_set_name(inp->self_kq_mask_swa, \"self_kq_mask_swa\");\n\n        inp->self_kq_mask_swa_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask_swa, GGML_TYPE_F16) : inp->self_kq_mask_swa;\n        ggml_set_name(inp->self_kq_mask_swa_cnv, \"self_kq_mask_swa_cnv\");\n    }\n\n    return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));\n}\n\nggml_tensor * llm_graph_context::build_rs(\n        ggml_tensor * s,\n        ggml_tensor * state_copy_main,\n        ggml_tensor * state_copy_extra,\n            int32_t   state_size,\n            int32_t   n_seqs,\n           uint32_t   n_rs,\n           uint32_t   rs_head,\n           uint32_t   rs_size,\n            int32_t   rs_zero,\n        const llm_graph_get_rows_fn & get_state_rows) const {\n\n    ggml_tensor * states = ggml_reshape_2d(ctx0, s, state_size, rs_size);\n\n    // Clear a single state which will then be copied to the other cleared states.\n    // Note that this is a no-op when the view is zero-sized.\n    ggml_tensor * state_zero = ggml_view_1d(ctx0, states, state_size*(rs_zero >= 0), rs_zero*states->nb[1]*(rs_zero >= 0));\n    ggml_build_forward_expand(gf, ggml_scale_inplace(ctx0, state_zero, 0));\n\n    // copy states\n    // NOTE: assuming the copy destinations are ALL contained between rs_head and rs_head + n_rs\n    // {state_size, rs_size} -> {state_size, n_seqs}\n    ggml_tensor * output_states = get_state_rows(ctx0, states, state_copy_main);\n    ggml_build_forward_expand(gf, output_states);\n\n    // copy extra states which won't be changed further (between n_seqs and n_rs)\n    ggml_tensor * states_extra = ggml_get_rows(ctx0, states, state_copy_extra);\n    ggml_build_forward_expand(gf,\n        ggml_cpy(ctx0,\n            states_extra,\n            ggml_view_1d(ctx0, s, state_size*(n_rs - n_seqs), (rs_head + n_seqs)*state_size*ggml_element_size(s))));\n\n    return output_states;\n}\n\nstatic std::unique_ptr<llm_graph_input_rs> build_rs_inp_impl(\n           ggml_context * ctx0,\n     const llama_ubatch & ubatch,\n    const llama_memory_recurrent_context * mctx_cur) {\n\n    auto inp = std::make_unique<llm_graph_input_rs>(mctx_cur);\n\n    const int64_t n_rs   = mctx_cur->get_n_rs();\n    const int64_t n_seqs = ubatch.n_seqs;\n\n    inp->s_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs);\n    ggml_set_input(inp->s_copy);\n\n    inp->s_copy_main  = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0);\n    inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]);\n\n    inp->head = mctx_cur->get_head();\n    inp->rs_z = mctx_cur->get_rs_z();\n\n    return inp;\n}\n\nllm_graph_input_rs * llm_graph_context::build_rs_inp() const {\n    const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);\n\n    auto inp = build_rs_inp_impl(ctx0, ubatch, mctx_cur);\n\n    return (llm_graph_input_rs *) res->add_input(std::move(inp));\n}\n\nggml_tensor * llm_graph_context::build_rs(\n        llm_graph_input_rs * inp,\n        ggml_tensor * s,\n            int32_t   state_size,\n            int32_t   n_seqs,\n        const llm_graph_get_rows_fn & get_state_rows) const {\n    const auto * kv_state = inp->mctx;\n\n    return build_rs(s, inp->s_copy_main, inp->s_copy_extra, state_size, n_seqs,\n                    kv_state->get_n_rs(), kv_state->get_head(), kv_state->get_size(), kv_state->get_rs_z(),\n                    get_state_rows);\n}\n\nggml_tensor * llm_graph_context::build_rwkv_token_shift_load(\n    llm_graph_input_rs * inp,\n    const llama_ubatch & ubatch,\n                   int   il) const {\n    const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);\n\n    const auto token_shift_count = hparams.token_shift_count;\n\n    const int64_t n_seqs  = ubatch.n_seqs;\n\n    ggml_tensor * token_shift_all = mctx_cur->get_r_l(il);\n\n    ggml_tensor * token_shift = build_rs(\n            inp, token_shift_all,\n            hparams.n_embd_r(), n_seqs);\n\n    token_shift = ggml_reshape_3d(ctx0, token_shift, hparams.n_embd, token_shift_count, n_seqs);\n\n    return token_shift;\n}\n\nggml_tensor * llm_graph_context::build_rwkv_token_shift_store(\n         ggml_tensor * token_shift,\n  const llama_ubatch & ubatch,\n                 int   il) const {\n    const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);\n\n    const auto token_shift_count = hparams.token_shift_count;\n    const auto n_embd = hparams.n_embd;\n\n    const int64_t n_seqs = ubatch.n_seqs;\n\n    const auto kv_head = mctx_cur->get_head();\n\n    return ggml_cpy(\n        ctx0,\n        ggml_view_1d(ctx0, token_shift, n_embd * n_seqs * token_shift_count, 0),\n        ggml_view_1d(ctx0, mctx_cur->get_r_l(il), hparams.n_embd_r()*n_seqs, hparams.n_embd_r()*kv_head*ggml_element_size(mctx_cur->get_r_l(il)))\n    );\n}\n\nllm_graph_input_mem_hybrid * llm_graph_context::build_inp_mem_hybrid() const {\n    const auto * mctx_cur = static_cast<const llama_memory_hybrid_context *>(mctx);\n\n    auto inp_rs   = build_rs_inp_impl     (ctx0, ubatch, mctx_cur->get_recr());\n    auto inp_attn = build_attn_inp_kv_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn());\n\n    auto inp = std::make_unique<llm_graph_input_mem_hybrid>(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur);\n\n    return (llm_graph_input_mem_hybrid *) res->add_input(std::move(inp));\n}\n\nvoid llm_graph_context::build_dense_out(\n    ggml_tensor * dense_2,\n    ggml_tensor * dense_3) const {\n    if (!cparams.embeddings || !(dense_2 || dense_3)) {\n        return;\n    }\n    ggml_tensor * cur = res->t_embd_pooled != nullptr ? res->t_embd_pooled : res->t_embd;\n    GGML_ASSERT(cur != nullptr && \"missing t_embd_pooled/t_embd\");\n\n    if (dense_2) {\n        cur = ggml_mul_mat(ctx0, dense_2, cur);\n    }\n    if (dense_3) {\n        cur = ggml_mul_mat(ctx0, dense_3, cur);\n    }\n    cb(cur, \"result_embd_pooled\", -1);\n    res->t_embd_pooled = cur;\n    ggml_build_forward_expand(gf, cur);\n}\n\n\nvoid llm_graph_context::build_pooling(\n        ggml_tensor * cls,\n        ggml_tensor * cls_b,\n        ggml_tensor * cls_out,\n        ggml_tensor * cls_out_b) const {\n    if (!cparams.embeddings) {\n        return;\n    }\n\n    ggml_tensor * inp = res->t_embd;\n\n    //// find result_norm tensor for input\n    //for (int i = ggml_graph_n_nodes(gf) - 1; i >= 0; --i) {\n    //    inp = ggml_graph_node(gf, i);\n    //    if (strcmp(inp->name, \"result_norm\") == 0 || strcmp(inp->name, \"result_embd\") == 0) {\n    //        break;\n    //    }\n\n    //    inp = nullptr;\n    //}\n\n    GGML_ASSERT(inp != nullptr && \"missing result_norm/result_embd tensor\");\n\n    ggml_tensor * cur;\n\n    switch (pooling_type) {\n        case LLAMA_POOLING_TYPE_NONE:\n            {\n                cur = inp;\n            } break;\n        case LLAMA_POOLING_TYPE_MEAN:\n            {\n                ggml_tensor * inp_mean = build_inp_mean();\n                cur = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, inp)), inp_mean);\n            } break;\n        case LLAMA_POOLING_TYPE_CLS:\n        case LLAMA_POOLING_TYPE_LAST:\n            {\n                ggml_tensor * inp_cls = build_inp_cls();\n                cur = ggml_get_rows(ctx0, inp, inp_cls);\n            } break;\n        case LLAMA_POOLING_TYPE_RANK:\n            {\n                ggml_tensor * inp_cls = build_inp_cls();\n                cur = ggml_get_rows(ctx0, inp, inp_cls);\n\n                // classification head\n                // https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/roberta/modeling_roberta.py#L1566\n                if (cls) {\n                    cur = ggml_mul_mat(ctx0, cls, cur);\n                    if (cls_b) {\n                        cur = ggml_add(ctx0, cur, cls_b);\n                    }\n                    cur = ggml_tanh(ctx0, cur);\n                }\n\n                // some models don't have `cls_out`, for example: https://huggingface.co/jinaai/jina-reranker-v1-tiny-en\n                // https://huggingface.co/jinaai/jina-reranker-v1-tiny-en/blob/cb5347e43979c3084a890e3f99491952603ae1b7/modeling_bert.py#L884-L896\n                // Single layer classification head (direct projection)\n                // https://github.com/huggingface/transformers/blob/f4fc42216cd56ab6b68270bf80d811614d8d59e4/src/transformers/models/bert/modeling_bert.py#L1476\n                if (cls_out) {\n                    cur = ggml_mul_mat(ctx0, cls_out, cur);\n                    if (cls_out_b) {\n                        cur = ggml_add(ctx0, cur, cls_out_b);\n                    }\n                }\n\n                // softmax for qwen3 reranker\n                if (arch == LLM_ARCH_QWEN3) {\n                    cur = ggml_soft_max(ctx0, cur);\n                }\n            } break;\n        default:\n            {\n                GGML_ABORT(\"unknown pooling type\");\n            }\n    }\n\n    cb(cur, \"result_embd_pooled\", -1);\n    res->t_embd_pooled = cur;\n\n    ggml_build_forward_expand(gf, cur);\n}\n\nvoid llm_graph_context::build_sampling() const {\n    if (samplers.empty() || !res->t_logits) {\n        return;\n    }\n\n    auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);\n    res->add_input(std::move(inp_sampling));\n\n    std::map<llama_seq_id, int32_t> seq_to_logit_row;\n    int32_t logit_row_idx = 0;\n\n    for (uint32_t i = 0; i < ubatch.n_tokens; i++) {\n        if (ubatch.output[i]) {\n            llama_seq_id seq_id = ubatch.seq_id[i][0];\n            seq_to_logit_row[seq_id] = logit_row_idx;\n            logit_row_idx++;\n        }\n    }\n\n    // res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1)\n    GGML_ASSERT(res->t_logits != nullptr && \"missing t_logits tensor\");\n\n    // add a dummy row of logits\n    // this trick makes the graph static, regardless of which samplers are activated\n    // this is important in order to minimize graph reallocations\n    // TODO: use `ggml_build_forward_select()` when available (https://github.com/ggml-org/llama.cpp/pull/18550)\n    ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);\n\n    for (const auto & [seq_id, sampler] : samplers) {\n        const auto it = seq_to_logit_row.find(seq_id);\n\n        // inactive samplers always work on the first row\n        const auto row_idx = seq_to_logit_row.find(seq_id) != seq_to_logit_row.end() ? it->second : 0;\n\n        ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);\n        ggml_format_name(logits_seq, \"logits_seq_%d\", seq_id);\n\n        struct llama_sampler_data data = {\n            /*.logits      =*/ logits_seq,\n            /*.probs       =*/ nullptr,\n            /*.sampled     =*/ nullptr,\n            /*.candidates  =*/ nullptr,\n        };\n\n        assert(sampler->iface->backend_apply);\n        sampler->iface->backend_apply(sampler, ctx0, gf, &data);\n\n        if (data.sampled != nullptr) {\n            res->t_sampled[seq_id] = data.sampled;\n            ggml_build_forward_expand(gf, data.sampled);\n        }\n\n        if (data.probs != nullptr) {\n            res->t_sampled_probs[seq_id] = data.probs;\n            ggml_build_forward_expand(gf, data.probs);\n        }\n\n        if (data.logits != nullptr) {\n            res->t_sampled_logits[seq_id] = data.logits;\n            ggml_build_forward_expand(gf, data.logits);\n        }\n\n        if (data.candidates != nullptr) {\n            res->t_candidates[seq_id] = data.candidates;\n            ggml_build_forward_expand(gf, data.candidates);\n        }\n    }\n\n    // TODO: Call llama_sampler_accept_ggml after all samplers have been applied.\n    /*\n    for (const auto & [seq_id, sampler] : samplers) {\n        if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) {\n            ggml_tensor * selected_token = it->second;\n            if (selected_token != nullptr) {\n                llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token);\n            }\n        }\n    }\n    */\n}\n\nint32_t llama_relative_position_bucket(llama_pos x, llama_pos y, uint64_t n_buckets, bool bidirectional) {\n    // TODO move to hparams if a T5 variant appears that uses a different value\n    const int64_t max_distance = 128;\n\n    if (bidirectional) {\n        n_buckets >>= 1;\n    }\n\n    const int64_t max_exact = n_buckets >> 1;\n\n    int32_t relative_position = x - y;\n    int32_t relative_bucket = 0;\n\n    if (bidirectional) {\n        relative_bucket += (relative_position > 0) * n_buckets;\n        relative_position = std::abs(relative_position);\n    } else {\n        relative_position = -std::min<int32_t>(relative_position, 0);\n    }\n\n    int32_t relative_position_if_large = floorf(max_exact + logf(1.0 * relative_position / max_exact) * (n_buckets - max_exact) / log(1.0 * max_distance / max_exact));\n    relative_position_if_large = std::min<int32_t>(relative_position_if_large, n_buckets - 1);\n    relative_bucket += (relative_position < max_exact ? relative_position : relative_position_if_large);\n\n    return relative_bucket;\n}\n"}}}}