Method process_ubatch

smallthinker/src/llama-context.cpp:654–702

}

llm_graph_result_ptr llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_state_i * mstate, ggml_status & ret) {
    if (mstate && !mstate->apply()) {
        LLAMA_LOG_ERROR("%s: failed to apply memory state\n", __func__);
        ret = GGML_STATUS_FAILED;
        return nullptr;
    }

    auto * gf = graph_init();
    if (!gf) {
        LLAMA_LOG_ERROR("%s: failed to initialize graph\n", __func__);
        ret = GGML_STATUS_FAILED;
        return nullptr;
    }

    auto res = graph_build(ctx_compute.get(), gf, ubatch, gtype, mstate);
    if (!res) {
        LLAMA_LOG_ERROR("%s: failed to build graph\n", __func__);
        ret = GGML_STATUS_FAILED;
        return nullptr;
    }

    // LLAMA_LOG_INFO("graph build time: %.3f ms (%d nodes, %d leafs)\n", (ggml_time_us() - t_start_us)/1000.0, gf->n_nodes, gf->n_leafs);

    if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) {
        LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__);
        ret = GGML_STATUS_ALLOC_FAILED;
        return nullptr;
    }

    res->set_inputs(&ubatch);

    const int n_threads = ubatch.n_tokens > 1 ? cparams.n_threads_batch : cparams.n_threads;
    if (powerinfer_has_global_expert_cache()) {
        powerinfer_init_moe_pipeline(n_threads, model.hparams.n_layer, model.hparams.n_embd, model.hparams.n_ff_exp,
                                     ubatch.n_tokens, model.hparams.n_expert, model.hparams.n_expert_used, true);
    }
    az::global_spin_barrier.init(n_threads);

    const auto status = graph_compute(gf, ubatch.n_tokens > 1);
    if (status != GGML_STATUS_SUCCESS) {
        LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status);
        ret = status;
        return nullptr;
    }

    ret = GGML_STATUS_SUCCESS;

    return res;
}

int llama_context::encode(llama_batch & inp_batch) {
    if (inp_batch.n_tokens == 0) {