Function llm_build_ffn_sparse

llama.cpp:4652–4757

}

static struct ggml_tensor * llm_build_ffn_sparse(
        struct ggml_context * ctx,
         struct ggml_tensor * cur,
         struct ggml_tensor * up,
         struct ggml_tensor * up_b,
         struct ggml_tensor * gate,
         struct ggml_tensor * gate_b,
         struct ggml_tensor * down_t,
         struct ggml_tensor * down_b,
         struct ggml_tensor * pre_w1,
         struct ggml_tensor * pre_w2,
         struct ggml_tensor * pred_inpl,
         struct ggml_tensor * gpu_index,
         struct ggml_tensor * gpu_bucket,
         struct ggml_tensor * gate_gpu,
         struct ggml_tensor * down_gpu,
         struct ggml_tensor * up_gpu,
            llm_ffn_op_type   type_op,
          llm_ffn_gate_type   type_gate,
                     double   gpu_offload_ratio,
   const llm_build_cb_short & cb_outer) {
    bool full_gpu = gpu_offload_ratio >= 1.0;
    ggml_tensor * ffn_input = cur;

    llm_build_cb_short cb = [&cb_outer](struct ggml_tensor * tensor, const char * name) {
        cb_outer(tensor, name);
#if defined(GGML_USE_CUBLAS)
        // Determine offloading based on src[0] (weight for both mul and axpy)
        bool operates_on_gpu = tensor->src[0]->backend == GGML_BACKEND_GPU;
        if (operates_on_gpu) {
            ggml_cuda_assign_buffers_no_alloc(tensor);
        }
#endif
    };

    // prepare sparse idx
    ggml_tensor * idx = ggml_mul_mat(ctx, pre_w1, pred_inpl);
    cb(idx, "mlp_pre_hidden");
    idx = ggml_relu(ctx, idx);
    cb(idx, "mlp_pre_relu");
    idx = ggml_mul_mat(ctx, pre_w2, idx);
    // If the FFN layer is not fully offloaded, we need to transfer the sparsity index
    // back to the CPU to avoid synchronization issues.
    (full_gpu ? cb : cb_outer)(idx, "mlp_pre_out");

    auto act_fn = [&](ggml_tensor * tensor, const char * name) {
        switch (type_op) {
            case LLM_FFN_RELU:
                {
                    tensor = ggml_relu(ctx, tensor);
                    cb(tensor, name);
                } break;
            default:
                GGML_ASSERT(false && "unsupported activation function");
        }
        return tensor;
    };