mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-04 17:10:48 -05:00
67d5978bb1
* Move MSA logic from llama-kv-cache into llama-kv-cache-msa * cont : minor * cont : ws fix --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
396 lines
12 KiB
C++
396 lines
12 KiB
C++
#include "llama-kv-cache-msa.h"
|
|
|
|
#include "llama-impl.h"
|
|
#include "llama-batch.h"
|
|
#include "llama-model.h"
|
|
|
|
#include <algorithm>
|
|
#include <cassert>
|
|
#include <cmath>
|
|
|
|
// llama_kv_cache_msa
|
|
|
|
llama_kv_cache_msa::llama_kv_cache_msa(
|
|
const llama_model & model,
|
|
ggml_type type_k,
|
|
ggml_type type_v,
|
|
bool v_trans,
|
|
bool offload,
|
|
bool unified,
|
|
uint32_t kv_size,
|
|
uint32_t n_seq_max,
|
|
uint32_t n_pad,
|
|
uint32_t n_swa,
|
|
llama_swa_type swa_type,
|
|
const layer_filter_cb & filter,
|
|
const layer_filter_cb & filter_idx,
|
|
const layer_reuse_cb & reuse) :
|
|
hparams_idx(model.hparams),
|
|
n_stream(unified ? 1 : n_seq_max), n_seq_max(n_seq_max), n_pad(n_pad),
|
|
n_swa(n_swa), swa_type(swa_type) {
|
|
|
|
LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);
|
|
|
|
kv_base = std::make_unique<llama_kv_cache>(
|
|
model, model.hparams, type_k, type_v,
|
|
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
|
|
n_swa, swa_type, nullptr, filter, reuse, nullptr);
|
|
|
|
// the MSA indexer uses a single key head per layer
|
|
std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
|
|
hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
|
|
// the rope parameters are kept identical to the main cache
|
|
|
|
LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
|
|
|
|
kv_idx = std::make_unique<llama_kv_cache>(
|
|
model, hparams_idx, type_k, type_v,
|
|
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
|
|
n_swa, swa_type, nullptr, filter_idx, reuse, nullptr);
|
|
}
|
|
|
|
void llama_kv_cache_msa::clear(bool data) {
|
|
kv_base->clear(data);
|
|
kv_idx ->clear(data);
|
|
}
|
|
|
|
bool llama_kv_cache_msa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
|
|
bool res = true;
|
|
|
|
res = res & kv_base->seq_rm(seq_id, p0, p1);
|
|
res = res & kv_idx ->seq_rm(seq_id, p0, p1);
|
|
|
|
return res;
|
|
}
|
|
|
|
void llama_kv_cache_msa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
|
|
kv_base->seq_cp(seq_id_src, seq_id_dst, p0, p1);
|
|
kv_idx ->seq_cp(seq_id_src, seq_id_dst, p0, p1);
|
|
}
|
|
|
|
void llama_kv_cache_msa::seq_keep(llama_seq_id seq_id) {
|
|
kv_base->seq_keep(seq_id);
|
|
kv_idx ->seq_keep(seq_id);
|
|
}
|
|
|
|
void llama_kv_cache_msa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
|
|
kv_base->seq_add(seq_id, p0, p1, shift);
|
|
kv_idx ->seq_add(seq_id, p0, p1, shift);
|
|
}
|
|
|
|
void llama_kv_cache_msa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
|
|
kv_base->seq_div(seq_id, p0, p1, d);
|
|
kv_idx ->seq_div(seq_id, p0, p1, d);
|
|
}
|
|
|
|
llama_pos llama_kv_cache_msa::seq_pos_min(llama_seq_id seq_id) const {
|
|
return kv_base->seq_pos_min(seq_id);
|
|
}
|
|
|
|
llama_pos llama_kv_cache_msa::seq_pos_max(llama_seq_id seq_id) const {
|
|
return kv_base->seq_pos_max(seq_id);
|
|
}
|
|
|
|
std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_msa::memory_breakdown() const {
|
|
std::map<ggml_backend_buffer_type_t, size_t> mb = kv_base->memory_breakdown();
|
|
for (const auto & buft_size : kv_idx->memory_breakdown()) {
|
|
mb[buft_size.first] += buft_size.second;
|
|
}
|
|
return mb;
|
|
}
|
|
|
|
llama_memory_context_ptr llama_kv_cache_msa::init_batch(
|
|
llama_batch_allocr & balloc,
|
|
uint32_t n_ubatch,
|
|
bool embd_all) {
|
|
GGML_UNUSED(embd_all);
|
|
|
|
do {
|
|
balloc.split_reset();
|
|
|
|
std::vector<llama_ubatch> ubatches;
|
|
while (true) {
|
|
auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
|
|
|
|
if (ubatch.n_tokens == 0) {
|
|
break;
|
|
}
|
|
|
|
ubatches.push_back(std::move(ubatch));
|
|
}
|
|
|
|
if (balloc.get_n_used() < balloc.get_n_tokens()) {
|
|
// failed to find a suitable split
|
|
break;
|
|
}
|
|
|
|
auto sinfos_base = kv_base->prepare(ubatches);
|
|
if (sinfos_base.empty()) {
|
|
break;
|
|
}
|
|
|
|
auto sinfos_idx = kv_idx->prepare(ubatches);
|
|
if (sinfos_idx.empty()) {
|
|
break;
|
|
}
|
|
|
|
assert(sinfos_base.size() == sinfos_idx.size());
|
|
|
|
return std::make_unique<llama_kv_cache_msa_context>(
|
|
this, std::move(sinfos_base), std::move(sinfos_idx), std::move(ubatches));
|
|
} while (false);
|
|
|
|
return std::make_unique<llama_kv_cache_msa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
|
|
}
|
|
|
|
llama_memory_context_ptr llama_kv_cache_msa::init_full() {
|
|
return std::make_unique<llama_kv_cache_msa_context>(this);
|
|
}
|
|
|
|
llama_memory_context_ptr llama_kv_cache_msa::init_update(llama_context * lctx, bool optimize) {
|
|
return std::make_unique<llama_kv_cache_msa_context>(this, lctx, optimize);
|
|
}
|
|
|
|
bool llama_kv_cache_msa::get_can_shift() const {
|
|
return kv_base->get_can_shift() &&
|
|
kv_idx ->get_can_shift() &&
|
|
kv_base->get_size() == kv_idx->get_size();
|
|
}
|
|
|
|
void llama_kv_cache_msa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
|
|
kv_base->state_write(io, seq_id, flags);
|
|
kv_idx ->state_write(io, seq_id, flags);
|
|
}
|
|
|
|
void llama_kv_cache_msa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
|
|
kv_base->state_read(io, seq_id, flags);
|
|
kv_idx ->state_read(io, seq_id, flags);
|
|
}
|
|
|
|
llama_kv_cache * llama_kv_cache_msa::get_base() const {
|
|
return kv_base.get();
|
|
}
|
|
|
|
llama_kv_cache * llama_kv_cache_msa::get_idx() const {
|
|
return kv_idx.get();
|
|
}
|
|
|
|
// llama_kv_cache_msa_context
|
|
|
|
llama_kv_cache_msa_context::llama_kv_cache_msa_context(llama_memory_status status) :
|
|
kv(nullptr), status(status) {}
|
|
|
|
llama_kv_cache_msa_context::llama_kv_cache_msa_context(
|
|
llama_kv_cache_msa * kv) :
|
|
kv(kv),
|
|
ctx_base(kv->get_base()->init_full()),
|
|
ctx_idx (kv->get_idx ()->init_full()),
|
|
status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
|
|
}
|
|
|
|
llama_kv_cache_msa_context::llama_kv_cache_msa_context(
|
|
llama_kv_cache_msa * kv,
|
|
llama_context * lctx,
|
|
bool optimize) :
|
|
kv(kv),
|
|
ctx_base(kv->get_base()->init_update(lctx, optimize)),
|
|
ctx_idx (kv->get_idx ()->init_update(lctx, optimize)),
|
|
status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
|
|
}
|
|
|
|
llama_kv_cache_msa_context::llama_kv_cache_msa_context(
|
|
llama_kv_cache_msa * kv,
|
|
slot_info_vec_t sinfos_base,
|
|
slot_info_vec_t sinfos_idx,
|
|
std::vector<llama_ubatch> ubatches) :
|
|
kv(kv),
|
|
ubatches(std::move(ubatches)),
|
|
// here we copy the ubatches. not sure if this is ideal
|
|
ctx_base(new llama_kv_cache_context(kv->get_base(), std::move(sinfos_base), this->ubatches)),
|
|
ctx_idx (new llama_kv_cache_context(kv->get_idx (), std::move(sinfos_idx), this->ubatches)),
|
|
status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
|
|
}
|
|
|
|
llama_kv_cache_msa_context::~llama_kv_cache_msa_context() = default;
|
|
|
|
bool llama_kv_cache_msa_context::next() {
|
|
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
|
|
|
|
ctx_base->next();
|
|
ctx_idx ->next();
|
|
|
|
if (++i_next >= ubatches.size()) {
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool llama_kv_cache_msa_context::apply() {
|
|
assert(!llama_memory_status_is_fail(status));
|
|
|
|
bool res = true;
|
|
|
|
res = res & ctx_base->apply();
|
|
res = res & ctx_idx ->apply();
|
|
|
|
return res;
|
|
}
|
|
|
|
llama_memory_status llama_kv_cache_msa_context::get_status() const {
|
|
return status;
|
|
}
|
|
|
|
const llama_ubatch & llama_kv_cache_msa_context::get_ubatch() const {
|
|
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
|
|
|
|
return ubatches[i_next];
|
|
}
|
|
|
|
const llama_kv_cache_context * llama_kv_cache_msa_context::get_base() const {
|
|
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
|
|
|
|
return static_cast<const llama_kv_cache_context *>(ctx_base.get());
|
|
}
|
|
|
|
const llama_kv_cache_context * llama_kv_cache_msa_context::get_idx() const {
|
|
assert(status == LLAMA_MEMORY_STATUS_SUCCESS);
|
|
|
|
return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
|
|
}
|
|
|
|
uint32_t llama_kv_cache_msa_context::get_n_pos() const {
|
|
// pad the value so that the graph remains constant across batches and can be reused
|
|
const uint32_t n_pad_cur = std::max(kv->get_n_pad(), 256u);
|
|
|
|
llama_pos pos_max = -1;
|
|
|
|
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) kv->get_n_seq_max(); ++seq_id) {
|
|
pos_max = std::max(pos_max, kv->seq_pos_max(seq_id));
|
|
}
|
|
|
|
return std::max(n_pad_cur, GGML_PAD((uint32_t) (pos_max + 1), n_pad_cur));
|
|
}
|
|
|
|
void llama_kv_cache_msa_context::set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const {
|
|
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
|
|
GGML_ASSERT(dst->type == GGML_TYPE_I32);
|
|
GGML_ASSERT(div > 0);
|
|
|
|
const int64_t n_tokens = ubatch->n_tokens;
|
|
const int64_t n_kv = dst->ne[0];
|
|
const int64_t n_stream_ub = dst->ne[1];
|
|
|
|
GGML_ASSERT(n_tokens % n_stream_ub == 0);
|
|
const int64_t n_tps = n_tokens/n_stream_ub;
|
|
|
|
int32_t * data = (int32_t *) dst->data;
|
|
|
|
for (int64_t s = 0; s < n_stream_ub; ++s) {
|
|
const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];
|
|
|
|
const auto & cells = kv->get_base()->get_cells(seq_id);
|
|
|
|
for (int64_t j = 0; j < n_kv; ++j) {
|
|
// the value for empty or other-sequence cells is irrelevant as consumers mask them
|
|
data[s*n_kv + j] =
|
|
cells.is_empty(j) || !cells.seq_has(j, seq_id)
|
|
? 0
|
|
: (int32_t) (cells.pos_get(j)/div);
|
|
}
|
|
}
|
|
}
|
|
|
|
void llama_kv_cache_msa_context::set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const {
|
|
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
|
|
GGML_ASSERT(dst->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_F32);
|
|
|
|
const int64_t n_tokens = ubatch->n_tokens;
|
|
const int64_t n_pos = dst->ne[0];
|
|
const int64_t n_stream_ub = dst->ne[1];
|
|
|
|
GGML_ASSERT(n_tokens % n_stream_ub == 0);
|
|
const int64_t n_tps = n_tokens/n_stream_ub;
|
|
|
|
for (int64_t s = 0; s < n_stream_ub; ++s) {
|
|
const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];
|
|
|
|
const auto & cells = kv->get_base()->get_cells(seq_id);
|
|
|
|
std::vector<int32_t> map(n_pos, 0);
|
|
|
|
for (uint32_t j = 0; j < cells.size(); ++j) {
|
|
if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
|
|
continue;
|
|
}
|
|
|
|
const llama_pos p0 = cells.pos_get(j);
|
|
|
|
if (p0 < 0 || p0 >= n_pos) {
|
|
continue;
|
|
}
|
|
|
|
map[p0] = (int32_t) j;
|
|
}
|
|
|
|
if (dst->type == GGML_TYPE_I32) {
|
|
int32_t * data = (int32_t *) dst->data + s*n_pos;
|
|
std::copy(map.begin(), map.end(), data);
|
|
} else {
|
|
float * data = (float *) dst->data + s*n_pos;
|
|
for (int64_t p = 0; p < n_pos; ++p) {
|
|
data[p] = (float) map[p];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void llama_kv_cache_msa_context::set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const {
|
|
GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
|
|
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
|
|
|
const int64_t n_tokens = ubatch->n_tokens;
|
|
const int64_t n_pos = dst->ne[0];
|
|
|
|
GGML_ASSERT(dst->ne[1] == n_tokens);
|
|
|
|
const uint32_t n_swa = kv->get_n_swa();
|
|
const llama_swa_type swa_type = kv->get_swa_type();
|
|
|
|
float * data = (float *) dst->data;
|
|
|
|
std::fill(data, data + n_pos*n_tokens, -INFINITY);
|
|
|
|
for (int64_t i = 0; i < n_tokens; ++i) {
|
|
const llama_seq_id seq_id = ubatch->seq_id[i][0];
|
|
|
|
const auto & cells = kv->get_base()->get_cells(seq_id);
|
|
|
|
const llama_pos p1 = ubatch->pos[i];
|
|
|
|
for (uint32_t j = 0; j < cells.size(); ++j) {
|
|
if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
|
|
continue;
|
|
}
|
|
|
|
const llama_pos p0 = cells.pos_get(j);
|
|
|
|
if (p0 < 0 || p0 >= n_pos) {
|
|
continue;
|
|
}
|
|
|
|
// causal mask
|
|
if (p0 > p1) {
|
|
continue;
|
|
}
|
|
|
|
// apply SWA if any
|
|
if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
|
|
continue;
|
|
}
|
|
|
|
data[i*n_pos + p0] = 0.0f;
|
|
}
|
|
}
|
|
}
|