Method offload_ffn_split

llama.cpp:2969–2997

    }

    size_t offload_ffn_split(llama_model * model) {
        LLAMA_LOG_INFO("%s: applying augmentation to model - please wait ...\n", __func__);
        const int64_t t_start_aug_us = ggml_time_us();
        std::vector<uint8_t> work_buffer;

        // Set sparsity threshold via global virables
        sparse_pred_threshold = model->hparams.sparse_pred_threshold;
#if defined (GGML_USE_CUBLAS)
        ggml_cuda_set_device_constants(model->hparams.sparse_pred_threshold);
#endif

        // load gpu_idx and slice mat to gpu
        size_t offloaded_bytes = 0;
        for (llama_layer &model_layer : model -> layers) {
            // gpu_idx load
            if (model_layer.gpu_idx == NULL && model_layer.gpu_bucket == NULL) {
                ggml_tensor * gpu_idx = ggml_new_tensor_1d(aux_ctx, GGML_TYPE_I32, model_layer.mlp_pre_w2 -> ne[1]);
                ggml_set_zero(gpu_idx);
                model_layer.gpu_idx = gpu_idx;
                ggml_tensor * gpu_bucket = ggml_new_tensor_1d(aux_ctx, GGML_TYPE_I32, 0);
                model_layer.gpu_bucket = gpu_bucket;
            }
            offloaded_bytes += slice_ffn_mat_to_gpu(model_layer);
            LLAMA_LOG_INFO(".");
        }

        LLAMA_LOG_INFO(" done (%.2f ms)\n", (ggml_time_us() - t_start_aug_us) / 1000.0);
        return offloaded_bytes;
    }
};

struct buffered_tensor_allocator {