Function llama_kv_cache_init

llama.cpp:1614–1677

kv cache helpers

// kv cache helpers
//
static bool llama_kv_cache_init(
        const struct llama_hparams & hparams,
             struct llama_kv_cache & cache,
                         ggml_type   wtype,
                          uint32_t   n_ctx,
                               int   n_gpu_layers) {
    const uint32_t n_embd  = hparams.n_embd_gqa();
    const uint32_t n_layer = hparams.n_layer;

    const int64_t n_mem      = n_layer*n_ctx;
    const int64_t n_elements = n_embd*n_mem;

    cache.has_shift = false;

    cache.head = 0;
    cache.size = n_ctx;

    cache.cells.clear();
    cache.cells.resize(n_ctx);

    cache.buf.resize(2u*n_elements*ggml_type_size(wtype) + 2u*ggml_tensor_overhead());
    memset(cache.buf.data, 0, cache.buf.size);

    struct ggml_init_params params;
    params.mem_size   = cache.buf.size;
    params.mem_buffer = cache.buf.data;
    params.no_alloc   = false;

    cache.ctx = ggml_init(params);

    if (!cache.ctx) {
        LLAMA_LOG_ERROR("%s: failed to allocate memory for kv cache\n", __func__);
        return false;
    }

    cache.k = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);
    cache.v = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);
    ggml_set_name(cache.k, "cache_k");
    ggml_set_name(cache.v, "cache_v");

    (void) n_gpu_layers;

#ifdef GGML_USE_CUBLAS
    if (ggml_cublas_loaded()) {
        size_t vram_kv_cache = 0;

        if (n_gpu_layers > (int)n_layer + 1) {
            ggml_cuda_assign_buffers_no_scratch(cache.v);
            LLAMA_LOG_INFO("%s: offloading v cache to GPU\n", __func__);
            vram_kv_cache += ggml_nbytes(cache.v);
        }
        if (n_gpu_layers > (int)n_layer + 2) {
            ggml_cuda_assign_buffers_no_scratch(cache.k);
            LLAMA_LOG_INFO("%s: offloading k cache to GPU\n", __func__);
            vram_kv_cache += ggml_nbytes(cache.k);
        }
        if (vram_kv_cache > 0) {
            LLAMA_LOG_INFO("%s: VRAM kv self = %.2f MB\n", __func__, vram_kv_cache / 1024.0 / 1024.0);