Felipe97/llama-cpp-compiled
01.1k
1#include "llama-kv-cache-dsv4.h"2 3#include "ggml-backend.h"4#include "llama-impl.h"5#include "llama-batch.h"6#include "llama-io.h"7#include "llama-model.h"8 9#include <algorithm>10#include <cassert>11#include <climits>12#include <cstdlib>13#include <cstring>14#include <map>15#include <sstream>16#include <stdexcept>17 18static constexpr uint32_t DSV4_CSA_RATIO = 4;19static constexpr uint32_t DSV4_HCA_RATIO = 128;20 21static constexpr uint32_t DSV4_STATE_MAGIC = 0x34565344; // DSV422static constexpr uint32_t DSV4_STATE_VERSION = 1;23static constexpr uint32_t DSV4_STATE_MODE_FULL = 0;24static constexpr uint32_t DSV4_STATE_MODE_PARTIAL = 1;25static constexpr uint32_t DSV4_K_CACHE_STATE_VER = 2;26static constexpr uint32_t DSV4_COMP_STATE_VER = 1;27 28static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {29 return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);30}31 32static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) {33 GGML_ASSERT(ggml_is_contiguous(tensor));34 GGML_ASSERT(tensor->ne[3] == 1);35 GGML_ASSERT(stream < (uint32_t) tensor->ne[2]);36 37 const size_t stream_size = tensor->nb[2];38 ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size);39}40 41static uint32_t dsv4_state_n_used_k_rows(llama_pos pos_max, uint32_t ratio, uint32_t kv_size) {42 if (pos_max < 0) {43 return 0;44 }45 46 const uint64_t n_rows = ((uint64_t) pos_max + 1)/ratio;47 48 return (uint32_t) std::min<uint64_t>(kv_size, n_rows);49}50 51static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {52 if (n_stream <= 1) {53 return 0;54 }55 if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {56 throw std::runtime_error("DSV4 sequence id out of stream range");57 }58 59 return (int64_t) seq_id*size;60}61 62static bool dsv4_ubatch_has_coupled(const llama_ubatch & ubatch) {63 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {64 if (ubatch.n_seq_id[i] > 1) {65 return true;66 }67 }68 69 return false;70}71 72static bool dsv4_token_has_seq(const llama_ubatch & ubatch, uint32_t i, llama_seq_id seq_id) {73 for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {74 if (ubatch.seq_id[i][s] == seq_id) {75 return true;76 }77 }78 79 return false;80}81 82static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) {83 if (!dsv4_ubatch_has_coupled(ubatch)) {84 return ubatch;85 }86 if (ubatch.embd) {87 throw std::runtime_error("DSV4 coupled embedding ubatches are not supported");88 }89 90 std::vector<uint32_t> counts(ubatch.n_seqs_unq, 0);91 uint32_t n_tokens = 0;92 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {93 const llama_seq_id seq_id = ubatch.seq_id_unq[s];94 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {95 if (dsv4_token_has_seq(ubatch, i, seq_id)) {96 ++counts[s];97 ++n_tokens;98 }99 }100 }101 102 if (n_tokens == 0) {103 return ubatch;104 }105 106 const uint32_t n_seq_tokens = counts[0];107 for (uint32_t s = 1; s < counts.size(); ++s) {108 if (counts[s] != n_seq_tokens) {109 throw std::runtime_error("DSV4 coupled raw writes require equal sequence lengths");110 }111 }112 113 auto data = std::make_shared<llama_ubatch::data_t>();114 data->pos.resize((size_t) n_tokens*ubatch.n_pos);115 data->n_seq_id.reserve(n_tokens);116 data->seq_id.reserve(n_tokens);117 data->seq_id_data.reserve(n_tokens);118 data->seq_id_unq.assign(ubatch.seq_id_unq, ubatch.seq_id_unq + ubatch.n_seqs_unq);119 data->seq_idx.assign(LLAMA_MAX_SEQ, -1);120 data->output.assign(n_tokens, 0);121 if (ubatch.token) {122 data->token.reserve(n_tokens);123 }124 125 for (uint32_t s = 0; s < data->seq_id_unq.size(); ++s) {126 data->seq_idx[data->seq_id_unq[s]] = s;127 }128 129 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {130 const llama_seq_id seq_id = ubatch.seq_id_unq[s];131 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {132 if (!dsv4_token_has_seq(ubatch, i, seq_id)) {133 continue;134 }135 136 const uint32_t dst = data->n_seq_id.size();137 if (ubatch.token) {138 data->token.push_back(ubatch.token[i]);139 }140 for (uint32_t p = 0; p < ubatch.n_pos; ++p) {141 data->pos[(size_t) p*n_tokens + dst] = ubatch.pos[(size_t) p*ubatch.n_tokens + i];142 }143 data->n_seq_id.push_back(1);144 data->seq_id_data.push_back(seq_id);145 }146 }147 148 for (uint32_t i = 0; i < n_tokens; ++i) {149 data->seq_id.push_back(&data->seq_id_data[i]);150 }151 152 llama_ubatch res {153 /*.b_equal_seqs =*/ true,154 /*.n_tokens =*/ n_tokens,155 /*.n_seq_tokens =*/ n_seq_tokens,156 /*.n_seqs =*/ ubatch.n_seqs_unq,157 /*.n_seqs_unq =*/ ubatch.n_seqs_unq,158 /*.n_pos =*/ ubatch.n_pos,159 /*.token =*/ data->token.empty() ? nullptr : data->token.data(),160 /*.embd =*/ nullptr,161 /*.pos =*/ data->pos.data(),162 /*.n_seq_id =*/ data->n_seq_id.data(),163 /*.seq_id =*/ data->seq_id.data(),164 /*.seq_id_unq =*/ data->seq_id_unq.data(),165 /*.seq_idx =*/ data->seq_idx.data(),166 /*.output =*/ data->output.data(),167 /*.data =*/ data,168 };169 170 return res;171}172 173static std::vector<llama_ubatch> dsv4_build_raw_write_ubatches(const std::vector<llama_ubatch> & ubatches) {174 std::vector<llama_ubatch> res;175 res.reserve(ubatches.size());176 for (const llama_ubatch & ubatch : ubatches) {177 res.push_back(dsv4_build_raw_write_ubatch(ubatch));178 }179 return res;180}181 182static bool dsv4_batch_has_coupled(const llama_batch & batch) {183 if (!batch.n_seq_id) {184 return false;185 }186 187 for (int32_t i = 0; i < batch.n_tokens; ++i) {188 if (batch.n_seq_id[i] > 1) {189 return true;190 }191 }192 193 return false;194}195 196static int64_t dsv4_comp_graph_n_stream(const llama_ubatch & ubatch, uint32_t n_stream) {197 // Coupled sequence sets must stay in one graph stream because their198 // compressed state is shared. Independent per-seq state can fan out.199 if (n_stream <= 1 || ubatch.n_seqs_unq <= 1 || dsv4_ubatch_has_coupled(ubatch)) {200 return 1;201 }202 203 return ubatch.n_seqs_unq;204}205 206static void dsv4_state_src_stream_range(207 uint32_t n_stream,208 llama_seq_id seq_id,209 uint32_t & s0,210 uint32_t & ns) {211 if (seq_id >= 0 && n_stream > 1) {212 if ((uint32_t) seq_id >= n_stream) {213 throw std::runtime_error("DSV4 state sequence id out of stream range");214 }215 216 s0 = (uint32_t) seq_id;217 ns = 1;218 return;219 }220 221 s0 = 0;222 ns = seq_id >= 0 ? 1 : n_stream;223}224 225static void dsv4_state_dst_stream_range(226 uint32_t n_stream,227 llama_seq_id seq_id,228 uint32_t ns,229 uint32_t & s0) {230 if (seq_id >= 0) {231 if (ns != 1) {232 throw std::runtime_error("DSV4 sequence state stream count mismatch");233 }234 if (n_stream > 1 && (uint32_t) seq_id >= n_stream) {235 throw std::runtime_error("DSV4 state sequence id out of stream range");236 }237 238 s0 = n_stream > 1 ? (uint32_t) seq_id : 0;239 return;240 }241 242 if (ns != n_stream) {243 throw std::runtime_error("DSV4 full state stream count mismatch");244 }245 246 s0 = 0;247}248 249static void dsv4_state_write_tensor_streams(250 llama_io_write_i & io,251 ggml_tensor * tensor,252 uint32_t tensor_rows,253 uint32_t n_rows,254 uint32_t s0,255 uint32_t ns,256 const std::vector<uint32_t> * stream_ids = nullptr) {257 const int32_t type_i = (int32_t) tensor->type;258 const uint64_t ne0 = tensor->ne[0];259 const uint64_t rows = n_rows;260 const uint64_t row_size = ggml_row_size(tensor->type, tensor->ne[0]);261 262 if (n_rows > tensor_rows) {263 throw std::runtime_error("DSV4 state tensor row count exceeds storage");264 }265 266 io.write(&type_i, sizeof(type_i));267 io.write(&ne0, sizeof(ne0));268 io.write(&rows, sizeof(rows));269 io.write(&row_size, sizeof(row_size));270 271 const size_t stream_stride = (size_t) tensor_rows*row_size;272 const size_t size = (size_t) n_rows*row_size;273 if (size == 0) {274 return;275 }276 277 if (stream_ids && stream_ids->size() != ns) {278 throw std::runtime_error("DSV4 state tensor stream map size mismatch");279 }280 281 for (uint32_t s = 0; s < ns; ++s) {282 const uint32_t stream = stream_ids ? (*stream_ids)[s] : s0 + s;283 if ((int64_t) stream >= tensor->ne[2]) {284 throw std::runtime_error("DSV4 state tensor stream out of range");285 }286 const size_t offset = (size_t) stream*stream_stride;287 io.write_tensor(tensor, offset, size);288 }289}290 291static void dsv4_state_read_tensor_streams(292 llama_io_read_i & io,293 ggml_tensor * tensor,294 uint32_t tensor_rows,295 uint32_t n_rows,296 uint32_t s0,297 uint32_t ns) {298 int32_t type_i_ref;299 uint64_t ne0_ref;300 uint64_t rows_ref;301 uint64_t row_size_ref;302 303 io.read(&type_i_ref, sizeof(type_i_ref));304 io.read(&ne0_ref, sizeof(ne0_ref));305 io.read(&rows_ref, sizeof(rows_ref));306 io.read(&row_size_ref, sizeof(row_size_ref));307 308 const int32_t type_i = (int32_t) tensor->type;309 const uint64_t ne0 = tensor->ne[0];310 const uint64_t rows = n_rows;311 const uint64_t row_size = ggml_row_size(tensor->type, tensor->ne[0]);312 313 if (type_i != type_i_ref || ne0 != ne0_ref || rows != rows_ref || row_size != row_size_ref) {314 throw std::runtime_error("DSV4 state tensor metadata mismatch");315 }316 if (n_rows > tensor_rows) {317 throw std::runtime_error("DSV4 state tensor row count exceeds storage");318 }319 320 const size_t stream_stride = (size_t) tensor_rows*row_size;321 const size_t size = (size_t) n_rows*row_size;322 if (size == 0) {323 return;324 }325 326 for (uint32_t s = 0; s < ns; ++s) {327 const size_t offset = (size_t) (s0 + s)*stream_stride;328 io.read_tensor(tensor, offset, size);329 }330}331 332static void dsv4_state_write_k_cache(333 llama_io_write_i & io,334 const llama_kv_cache * kv,335 llama_seq_id seq_id,336 llama_state_seq_flags flags,337 uint32_t n_rows) {338 GGML_UNUSED(flags);339 340 uint32_t s0;341 uint32_t ns;342 dsv4_state_src_stream_range(kv->get_n_stream(), seq_id, s0, ns);343 344 const uint32_t version = DSV4_K_CACHE_STATE_VER;345 const uint32_t kv_size = kv->get_size();346 const auto layer_ids = kv->get_layer_ids();347 const uint32_t n_layer = layer_ids.size();348 349 if (n_rows > kv_size) {350 throw std::runtime_error("DSV4 K-cache state row count exceeds cache size");351 }352 353 io.write(&version, sizeof(version));354 io.write(&n_rows, sizeof(n_rows));355 io.write(&ns, sizeof(ns));356 io.write(&n_layer, sizeof(n_layer));357 358 for (uint32_t il : layer_ids) {359 io.write(&il, sizeof(il));360 dsv4_state_write_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows, s0, ns);361 }362}363 364static void dsv4_state_read_k_cache(365 llama_io_read_i & io,366 llama_kv_cache * kv,367 llama_seq_id seq_id,368 llama_state_seq_flags flags) {369 GGML_UNUSED(flags);370 371 uint32_t version;372 uint32_t n_rows_ref;373 uint32_t ns;374 uint32_t n_layer_ref;375 376 io.read(&version, sizeof(version));377 io.read(&n_rows_ref, sizeof(n_rows_ref));378 io.read(&ns, sizeof(ns));379 io.read(&n_layer_ref, sizeof(n_layer_ref));380 381 if (version != 1 && version != DSV4_K_CACHE_STATE_VER) {382 throw std::runtime_error("DSV4 K-cache state version mismatch");383 }384 385 const uint32_t kv_size = kv->get_size();386 if (version == 1 && n_rows_ref != kv_size) {387 LLAMA_LOG_INFO("kv size ref %d kv %d\n", n_rows_ref, kv_size);388 throw std::runtime_error("DSV4 K-cache state size mismatch");389 }390 if (n_rows_ref > kv_size) {391 LLAMA_LOG_INFO("kv rows ref %d kv %d\n", n_rows_ref, kv_size);392 throw std::runtime_error("DSV4 K-cache state size mismatch");393 }394 395 uint32_t s0;396 dsv4_state_dst_stream_range(kv->get_n_stream(), seq_id, ns, s0);397 398 const auto layer_ids = kv->get_layer_ids();399 if (n_layer_ref != layer_ids.size()) {400 throw std::runtime_error("DSV4 K-cache layer count mismatch");401 }402 403 for (uint32_t il : layer_ids) {404 uint32_t il_ref;405 io.read(&il_ref, sizeof(il_ref));406 if (il_ref != il) {407 throw std::runtime_error("DSV4 K-cache layer id mismatch");408 }409 410 dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv_size, n_rows_ref, s0, ns);411 }412}413 414static std::string dsv4_plan_positions(const std::vector<int32_t> & values) {415 std::ostringstream ss;416 ss << "[";417 for (size_t i = 0; i < values.size(); ++i) {418 if (i > 0) {419 ss << ", ";420 }421 ss << values[i];422 }423 ss << "]";424 return ss.str();425}426 427static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(428 const llama_ubatch & ubatch,429 uint32_t ratio,430 bool overlap,431 uint32_t state_size,432 uint32_t kv_size,433 uint32_t n_stream,434 uint32_t n_rs_seq,435 const std::vector<uint32_t> & rs_idx) {436 llama_kv_cache_dsv4_context::comp_plan plan;437 plan.n_visible.resize(ubatch.n_tokens);438 plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);439 440 // n_stream is the persistent cache/state layout; plan.n_stream is the441 // graph view for this ubatch and can be a subset of those streams.442 if (n_stream <= 1 && ubatch.n_seqs_unq > 1) {443 throw std::runtime_error("DSV4 single compressed stream cannot serve multiple sequences");444 }445 446 const int64_t state_rows = (int64_t) state_size*n_stream;447 448 struct persist_row {449 int32_t dst;450 int32_t src;451 llama_pos pos;452 };453 454 std::vector<persist_row> persist_rows;455 456 // For the overlap compressor, build_overlap_compressed_kv_from_state() consumes457 // state_read_idxs as two contiguous halves: the first ratio*n_blocks entries are458 // the "previous-window" gather indices for every block, followed by the459 // "current-window" indices for every block. Collect them separately here and460 // append cur after prev once the loop has visited all completed blocks461 std::vector<int32_t> overlap_prev_reads;462 std::vector<int32_t> overlap_cur_reads;463 464 std::map<std::pair<llama_seq_id, llama_pos>, int64_t> curr_token_idx_map;465 std::map<llama_seq_id, uint32_t> state_write_counts;466 467 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {468 for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {469 curr_token_idx_map[std::make_pair(ubatch.seq_id[i][s], ubatch.pos[i])] = i;470 }471 }472 473 const auto state_source_idx = [&](llama_seq_id seq_id, llama_pos pos) -> int32_t {474 if (pos < 0) {475 // The overlap compressor needs a zero/-inf source for the first476 // block's previous half. The graph appends that row after the477 // current-ubatch scratch rows.478 return (int32_t) (state_rows + ubatch.n_tokens);479 }480 481 const auto key = std::make_pair(seq_id, pos);482 if (curr_token_idx_map.find(key) != curr_token_idx_map.end()) {483 return (int32_t) (state_rows + curr_token_idx_map.at(key));484 }485 486 const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);487 return (int32_t) (stream_off + pos%state_size);488 };489 490 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {491 const llama_pos pos = ubatch.pos[i];492 493 if (pos < 0) {494 continue;495 }496 497 plan.state_pos.push_back((int32_t) (pos%ratio));498 499 const int64_t n_visible = (int64_t) (pos + 1)/ratio;500 plan.n_visible[i] = (int32_t) n_visible;501 plan.n_kv = std::max(plan.n_kv, n_visible);502 503 for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {504 const llama_seq_id seq_id = ubatch.seq_id[i][s];505 const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);506 const int32_t state_idx = (int32_t) (stream_off + pos%state_size);507 508 const auto it = std::find_if(persist_rows.begin(), persist_rows.end(),509 [state_idx](const persist_row & row) {510 return row.dst == state_idx;511 });512 if (it == persist_rows.end()) {513 persist_rows.push_back({ state_idx, (int32_t) i, pos });514 } else if (pos > it->pos) {515 it->src = (int32_t) i;516 it->pos = pos;517 }518 519 if ((pos + 1) % ratio != 0) {520 continue;521 }522 523 const llama_pos source_start = pos + 1 - ratio;524 const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);525 526 plan.state_write_idxs.push_back(cache_off + pos/ratio);527 plan.state_write_pos.push_back((int32_t) source_start);528 ++state_write_counts[seq_id];529 530 if (overlap) {531 const llama_pos prev_start = source_start - ratio;532 533 for (uint32_t j = 0; j < ratio; ++j) {534 overlap_prev_reads.push_back(state_source_idx(seq_id, prev_start + j));535 }536 for (uint32_t j = 0; j < ratio; ++j) {537 overlap_cur_reads.push_back(state_source_idx(seq_id, source_start + j));538 }539 } else {540 for (uint32_t j = 0; j < ratio; ++j) {541 plan.state_read_idxs.push_back(state_source_idx(seq_id, source_start + j));542 }543 }544 }545 }546 547 if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {548 assert(kv_size > 0);549 550 // Pad each stream to the reserve plan's block count.551 const auto append_dummy_block = [&](llama_seq_id seq_id, uint32_t i) {552 const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);553 const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);554 555 plan.state_write_idxs.push_back(cache_off + kv_size - 1);556 plan.state_write_pos .push_back(0);557 558 if (overlap) {559 for (uint32_t j = 0; j < ratio; ++j) {560 overlap_prev_reads.push_back(source_idx);561 overlap_cur_reads .push_back(source_idx);562 }563 } else {564 for (uint32_t j = 0; j < ratio; ++j) {565 plan.state_read_idxs.push_back(source_idx);566 }567 }568 };569 570 if (dsv4_ubatch_has_coupled(ubatch)) {571 if (plan.state_write_idxs.empty()) {572 uint32_t i = 0;573 while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {574 ++i;575 }576 assert(i < ubatch.n_tokens);577 append_dummy_block(ubatch.seq_id[i][0], i);578 }579 } else {580 const uint32_t n_blocks = (std::max<uint32_t>(1, ubatch.n_seq_tokens) + ratio - 1)/ratio;581 582 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {583 const llama_seq_id seq_id = ubatch.seq_id_unq[s];584 const uint32_t n_writes = state_write_counts[seq_id];585 if (n_writes >= n_blocks) {586 continue;587 }588 if (n_writes + 1 != n_blocks) {589 throw std::runtime_error("DSV4 CSA sequence positions are not contiguous");590 }591 592 uint32_t i = 0;593 while (i < ubatch.n_tokens && (ubatch.pos[i] < 0 || !dsv4_token_has_seq(ubatch, i, seq_id))) {594 ++i;595 }596 assert(i < ubatch.n_tokens);597 append_dummy_block(seq_id, i);598 }599 }600 }601 602 if (ratio == DSV4_HCA_RATIO && !plan.state_pos.empty() && plan.state_write_idxs.empty()) {603 assert(kv_size > 0);604 // the last slot must not be live, or the dummy write would corrupt it;605 // a full stream implies a completed block, which implies real writes606 assert(plan.n_kv < (int64_t) kv_size);607 608 // Keep the compress/write ops in the graph when no HCA block completes609 // in this ubatch. The dummy block writes to the last cache slot and is610 // masked out.611 uint32_t i = 0;612 while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {613 ++i;614 }615 assert(i < ubatch.n_tokens);616 617 const llama_seq_id seq_id = ubatch.seq_id[i][0];618 const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);619 const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);620 621 plan.state_write_idxs.push_back(cache_off + kv_size - 1);622 plan.state_write_pos .push_back(0);623 624 for (uint32_t j = 0; j < ratio; ++j) {625 plan.state_read_idxs.push_back(source_idx);626 }627 }628 629 if (overlap) {630 // [ all blocks' prev-window indices | all blocks' cur-window indices ]631 plan.state_read_idxs.reserve(overlap_prev_reads.size() + overlap_cur_reads.size());632 plan.state_read_idxs.insert(plan.state_read_idxs.end(),633 overlap_prev_reads.begin(), overlap_prev_reads.end());634 plan.state_read_idxs.insert(plan.state_read_idxs.end(),635 overlap_cur_reads.begin(), overlap_cur_reads.end());636 }637 638 // Keep the mask (and with it the compressed-attention branch) present even639 // before the first block is visible, so the graph topology never changes.640 // Padded slots are masked out; comp cache buffers are zero-initialized.641 plan.n_kv = std::max<int64_t>(GGML_PAD(plan.n_kv, 256u), 256);642 643 std::sort(persist_rows.begin(), persist_rows.end(),644 [](const persist_row & a, const persist_row & b) {645 return a.dst < b.dst;646 });647 648 for (const persist_row & row : persist_rows) {649 plan.state_persist_src_idxs.push_back(row.src);650 plan.state_persist_dst_idxs.push_back(row.dst);651 }652 653 if (n_rs_seq > 0) {654 // Emit restore/snapshot entries for all layout streams so that the655 // graph tensor sizes do not depend on the ubatch's sequence count.656 // Streams not present in the ubatch get no-op entries.657 for (uint32_t stream = 0; stream < n_stream; ++stream) {658 llama_seq_id seq_id = -1;659 if (n_stream == 1) {660 // a unified stream serves any single sequence661 seq_id = ubatch.n_seqs_unq > 0 ? ubatch.seq_id_unq[0] : -1;662 } else {663 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {664 if (ubatch.seq_id_unq[s] == (llama_seq_id) stream) {665 seq_id = ubatch.seq_id_unq[s];666 break;667 }668 }669 }670 671 const int64_t stream_off = (int64_t) stream*state_size;672 const uint32_t rollback = seq_id >= 0 && (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;673 // Keep the restore graph fixed-width when no rollback is pending.674 const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0;675 for (uint32_t r = 0; r < state_size; ++r) {676 plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r));677 plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r));678 }679 680 std::vector<uint32_t> token_idxs;681 token_idxs.reserve(ubatch.n_tokens);682 if (seq_id >= 0) {683 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {684 if (dsv4_token_has_seq(ubatch, i, seq_id)) {685 token_idxs.push_back(i);686 }687 }688 }689 690 const uint32_t n_seq_tokens = (uint32_t) token_idxs.size();691 const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq);692 for (uint32_t d = 1; d <= n_rs_seq; ++d) {693 const int64_t dst_plane = (int64_t) d*state_rows;694 const uint32_t prefix = d <= n_seq_tokens ? n_seq_tokens - d : 0;695 696 for (uint32_t r = 0; r < state_size; ++r) {697 int32_t src = (int32_t) (stream_off + r);698 699 for (uint32_t j = 0; j < prefix; ++j) {700 const uint32_t i_tok = token_idxs[j];701 if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {702 src = (int32_t) (scratch_off + i_tok);703 }704 }705 706 if (n_seq_tokens == 0) {707 // no-op: copy the snapshot plane onto itself708 src = (int32_t) (dst_plane + stream_off + r);709 }710 711 plan.state_snapshot_src_idxs.push_back(src);712 plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r));713 }714 }715 }716 }717 718 static const bool debug = []() {719 const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG");720 return env && atoi(env) > 0;721 }();722 723 if (debug) {724 LLAMA_LOG_DEBUG("%s: ratio=%u, n_tokens=%u, n_seqs_unq=%u, state_persist_dst=%s, state_write_pos=%s\n",725 __func__, ratio, ubatch.n_tokens, ubatch.n_seqs_unq,726 dsv4_plan_positions(plan.state_persist_dst_idxs).c_str(),727 dsv4_plan_positions(plan.state_write_pos).c_str());728 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {729 const llama_seq_id seq_id = ubatch.seq_id_unq[s];730 const uint32_t rollback = seq_id >= 0 && (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;731 LLAMA_LOG_DEBUG("%s: seq %d pos [%d, %d] rollback=%u\n", __func__, seq_id,732 ubatch.pos[0], ubatch.pos[ubatch.n_tokens - 1], rollback);733 }734 }735 736 return plan;737}738 739static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans(740 const std::vector<llama_ubatch> & ubatches,741 uint32_t ratio,742 bool overlap,743 uint32_t state_size,744 uint32_t kv_size,745 uint32_t n_stream,746 uint32_t n_rs_seq,747 const std::vector<uint32_t> & rs_idx) {748 std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;749 plans.reserve(ubatches.size());750 751 // the first ubatch touching a seq consumes its rollback restore752 std::vector<uint32_t> rs(rs_idx);753 for (const llama_ubatch & ubatch : ubatches) {754 plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs));755 756 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {757 const llama_seq_id seq_id = ubatch.seq_id_unq[s];758 if (seq_id >= 0 && (size_t) seq_id < rs.size()) {759 rs[seq_id] = 0;760 }761 }762 }763 764 return plans;765}766 767static llama_kv_cache::slot_info_vec_t dsv4_build_comp_sinfos(768 const std::vector<llama_ubatch> & ubatches,769 uint32_t n_stream) {770 llama_kv_cache::slot_info_vec_t sinfos;771 sinfos.reserve(ubatches.size());772 773 for (const llama_ubatch & ubatch : ubatches) {774 if (n_stream <= 1 && ubatch.n_seqs_unq > 1) {775 throw std::runtime_error("DSV4 single compressed stream cannot serve multiple sequences");776 }777 778 const uint32_t ns = (uint32_t) dsv4_comp_graph_n_stream(ubatch, n_stream);779 llama_kv_cache::slot_info sinfo;780 sinfo.s0 = n_stream > 1 ? LLAMA_MAX_SEQ : 0;781 sinfo.s1 = 0;782 sinfo.resize(ns);783 784 for (uint32_t s = 0; s < ns; ++s) {785 const llama_seq_id seq_id = n_stream > 1 ? ubatch.seq_id_unq[s] : 0;786 const uint32_t strm = (uint32_t) dsv4_stream_offset(n_stream, seq_id, 1);787 788 sinfo.s0 = std::min(sinfo.s0, strm);789 sinfo.s1 = std::max(sinfo.s1, strm);790 sinfo.strm[s] = strm;791 sinfo.idxs[s].resize(1, 0);792 }793 794 if (n_stream > 1 && sinfo.s1 - sinfo.s0 + 1 != ns) {795 throw std::runtime_error("DSV4 compressed streams are not contiguous in ubatch");796 }797 798 sinfos.push_back(std::move(sinfo));799 }800 801 return sinfos;802}803 804static llama_kv_cache::slot_info_vec_t dsv4_build_raw_read_sinfos(805 const llama_kv_cache::slot_info_vec_t & sinfos_write,806 const std::vector<llama_ubatch> & ubatches) {807 llama_kv_cache::slot_info_vec_t sinfos;808 sinfos.reserve(ubatches.size());809 810 for (size_t i = 0; i < ubatches.size(); ++i) {811 const llama_ubatch & ubatch = ubatches[i];812 const auto & sinfo_write = sinfos_write[i];813 814 if (!dsv4_ubatch_has_coupled(ubatch)) {815 sinfos.push_back(sinfo_write);816 continue;817 }818 819 const llama_seq_id seq_id = ubatch.seq_id[0][0];820 uint32_t i_stream = 0;821 for (; i_stream < sinfo_write.n_stream(); ++i_stream) {822 if (sinfo_write.strm[i_stream] == seq_id) {823 break;824 }825 }826 if (i_stream == sinfo_write.n_stream()) {827 throw std::runtime_error("DSV4 raw write stream not found for coupled read");828 }829 830 llama_kv_cache::slot_info sinfo;831 sinfo.s0 = sinfo_write.strm[i_stream];832 sinfo.s1 = sinfo_write.strm[i_stream];833 sinfo.resize(1);834 sinfo.strm[0] = sinfo_write.strm[i_stream];835 sinfo.idxs[0] = sinfo_write.idxs[i_stream];836 sinfos.push_back(std::move(sinfo));837 }838 839 return sinfos;840}841 842static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(843 const llama_ubatch & ubatch,844 uint32_t ratio,845 bool overlap,846 uint32_t state_size,847 uint32_t kv_size,848 uint32_t n_stream,849 uint32_t n_rs_seq) {850 llama_kv_cache_dsv4_context::comp_plan plan;851 plan.n_visible.resize(ubatch.n_tokens);852 plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);853 plan.n_kv = kv_size;854 855 if (ubatch.n_tokens == 0) {856 return plan;857 }858 859 // worst case over every seq split: sum of per-seq ceil(tokens/ratio) is at860 // most floor(n_tokens/ratio) + n_seqs861 const uint32_t n_seqs = std::max<uint32_t>(1, ubatch.n_seqs);862 const size_t n_blocks = (size_t) ubatch.n_tokens/ratio + n_seqs;863 864 const uint64_t state_rows = (uint64_t) state_size*n_stream;865 const size_t n_persist = (size_t) std::min<uint64_t>(ubatch.n_tokens, state_rows);866 const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*n_stream : 0;867 const size_t n_snapshot = (size_t) n_rs_seq*state_size*n_stream;868 869 plan.state_pos .resize(ubatch.n_tokens);870 plan.state_persist_src_idxs.resize(n_persist);871 plan.state_persist_dst_idxs.resize(n_persist);872 plan.state_restore_src_idxs.resize(n_restore);873 plan.state_restore_dst_idxs.resize(n_restore);874 plan.state_snapshot_src_idxs.resize(n_snapshot);875 plan.state_snapshot_dst_idxs.resize(n_snapshot);876 plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks);877 plan.state_write_idxs.resize(n_blocks);878 plan.state_write_pos .resize(n_blocks);879 880 return plan;881}882 883static void dsv4_make_k_only(llama_hparams & hparams) {884 // llama_kv_cache uses hparams.is_mla() to allocate K-only storage.885 hparams.n_embd_head_k_mla_impl = hparams.n_embd_head_k();886 hparams.n_embd_head_v_mla_impl = hparams.n_embd_head_k();887}888 889//890// llama_dsv4_comp_state891//892 893llama_dsv4_comp_state::llama_dsv4_comp_state(894 const llama_model & model,895 bool offload,896 bool unified,897 uint32_t n_seq_max,898 uint32_t ratio,899 uint32_t state_size,900 uint32_t n_embd_state,901 uint32_t n_rs_seq,902 const char * name,903 const llama_memory_i::layer_filter_cb & filter) :904 ratio(ratio),905 state_size(state_size),906 n_embd_state(n_embd_state),907 n_stream(unified ? 1 : n_seq_max),908 n_rs_seq(n_rs_seq) {909 const llama_hparams & hparams = model.hparams;910 911 struct ggml_backend_buft_comparator {912 bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {913 return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;914 }915 };916 917 std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> ctx_map;918 919 auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {920 auto it = ctx_map.find(buft);921 if (it == ctx_map.end()) {922 ggml_init_params params = {923 /*.mem_size =*/ size_t(2u*(1 + n_stream)*hparams.n_layer()*ggml_tensor_overhead()),924 /*.mem_buffer =*/ NULL,925 /*.no_alloc =*/ true,926 };927 928 ggml_context * ctx = ggml_init(params);929 if (!ctx) {930 return nullptr;931 }932 933 ctx_map.emplace(buft, ctx);934 935 return ctx;936 }937 938 return it->second.get();939 };940 941 for (uint32_t il = 0; il < hparams.n_layer(); ++il) {942 if (filter && !filter(il)) {943 continue;944 }945 946 const char * dev_name = "CPU";947 948 ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();949 950 if (offload) {951 auto * dev = model.dev_layer(il);952 buft = ggml_backend_dev_buffer_type(dev);953 954 dev_name = ggml_backend_dev_name(dev);955 }956 957 LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);958 959 ggml_context * ctx = ctx_for_buft(buft);960 if (!ctx) {961 throw std::runtime_error("failed to create ggml context for DSV4 compressor state");962 }963 964 const uint32_t n_planes = n_stream*(1 + n_rs_seq);965 ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);966 ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);967 968 ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il);969 ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);970 971 std::vector<ggml_tensor *> kv_stream;972 std::vector<ggml_tensor *> score_stream;973 974 for (uint32_t s = 0; s < n_stream; ++s) {975 kv_stream.push_back(ggml_view_2d(ctx, kv, n_embd_state, state_size, kv->nb[1], s*kv->nb[2]));976 score_stream.push_back(ggml_view_2d(ctx, score, n_embd_state, state_size, score->nb[1], s*score->nb[2]));977 }978 979 map_layer_ids[il] = layers.size();980 981 layers.push_back({ il, kv, score, std::move(kv_stream), std::move(score_stream) });982 }983 984 for (auto & [buft, ctx] : ctx_map) {985 ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);986 if (!buf) {987 throw std::runtime_error("failed to allocate buffer for DSV4 compressor state");988 }989 990 ggml_backend_buffer_clear(buf, 0);991 992 LLAMA_LOG_INFO("%s: %10s DSV4 %s state buffer size = %8.2f MiB\n",993 __func__, ggml_backend_buffer_name(buf), name, ggml_backend_buffer_get_size(buf)/1024.0/1024.0);994 995 ctxs_bufs.emplace_back(std::move(ctx), buf);996 }997 998 LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n",999 __func__, name, ratio, state_size, n_embd_state, n_stream, n_rs_seq, layers.size(), total_size()/1024.0/1024.0);1000}1001 1002void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {1003 if (!data) {1004 return;1005 }1006 1007 if (seq_id >= 0) {1008 GGML_ASSERT((uint32_t) seq_id < n_stream);1009 1010 for (const auto & layer : layers) {1011 for (uint32_t d = 0; d <= n_rs_seq; ++d) {1012 const uint32_t stream = d*n_stream + (uint32_t) seq_id;1013 dsv4_clear_tensor_stream(layer.kv, stream);1014 dsv4_clear_tensor_stream(layer.score, stream);1015 }1016 }1017 return;1018 }1019 1020 for (auto & [_, buf] : ctxs_bufs) {1021 ggml_backend_buffer_clear(buf.get(), 0);1022 }1023}1024 1025void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) {1026 GGML_ASSERT(seq_id_src >= 0 && (uint32_t) seq_id_src < n_stream);1027 GGML_ASSERT(seq_id_dst >= 0 && (uint32_t) seq_id_dst < n_stream);1028 1029 if (seq_id_src == seq_id_dst) {1030 return;1031 }1032 1033 clear(seq_id_dst, true);1034 1035 sc_info.ssrc.push_back((uint32_t) seq_id_src);1036 sc_info.sdst.push_back((uint32_t) seq_id_dst);1037}1038 1039void llama_dsv4_comp_state::apply_copies(const stream_copy_info & sc_info) const {1040 for (size_t i = 0; i < sc_info.ssrc.size(); ++i) {1041 const uint32_t ssrc = sc_info.ssrc[i];1042 const uint32_t sdst = sc_info.sdst[i];1043 1044 for (const auto & layer : layers) {1045 ggml_backend_tensor_copy(layer.kv_stream[ssrc], layer.kv_stream[sdst]);1046 ggml_backend_tensor_copy(layer.score_stream[ssrc], layer.score_stream[sdst]);1047 }1048 }1049}1050 1051uint32_t llama_dsv4_comp_state::get_ratio() const {1052 return ratio;1053}1054 1055uint32_t llama_dsv4_comp_state::get_state_size() const {1056 return state_size;1057}1058 1059uint32_t llama_dsv4_comp_state::get_n_stream() const {1060 return n_stream;1061}1062 1063uint32_t llama_dsv4_comp_state::get_n_rs_seq() const {1064 return n_rs_seq;1065}1066 1067uint32_t llama_dsv4_comp_state::get_n_rows() const {1068 return state_size*n_stream;1069}1070 1071std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {1072 std::map<ggml_backend_buffer_type_t, size_t> ret;1073 for (const auto & [_, buf] : ctxs_bufs) {1074 ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());1075 ret[buft] += ggml_backend_buffer_get_size(buf.get());1076 }1077 return ret;1078}1079 1080void llama_dsv4_comp_state::state_write(1081 llama_io_write_i & io,1082 llama_seq_id seq_id,1083 llama_state_seq_flags flags,1084 const std::vector<uint32_t> & rs_idx) const {1085 GGML_UNUSED(flags);1086 1087 uint32_t s0;1088 uint32_t ns;1089 dsv4_state_src_stream_range(n_stream, seq_id, s0, ns);1090 1091 std::vector<uint32_t> stream_ids(ns);1092 for (uint32_t s = 0; s < ns; ++s) {1093 const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s;1094 if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) {1095 throw std::runtime_error("DSV4 recurrent state rollback index out of range");1096 }1097 stream_ids[s] = rs_idx[seq]*n_stream + s0 + s;1098 }1099 1100 const uint32_t version = DSV4_COMP_STATE_VER;1101 const uint32_t n_layer = layers.size();1102 1103 io.write(&version, sizeof(version));1104 io.write(&ratio, sizeof(ratio));1105 io.write(&state_size, sizeof(state_size));1106 io.write(&n_embd_state, sizeof(n_embd_state));1107 io.write(&ns, sizeof(ns));1108 io.write(&n_layer, sizeof(n_layer));1109 1110 for (const auto & layer : layers) {1111 io.write(&layer.il, sizeof(layer.il));1112 1113 dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns, &stream_ids);1114 dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns, &stream_ids);1115 }1116}1117 1118void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {1119 GGML_UNUSED(flags);1120 1121 uint32_t version;1122 uint32_t ratio_ref;1123 uint32_t state_size_ref;1124 uint32_t n_embd_state_ref;1125 uint32_t ns;1126 uint32_t n_layer_ref;1127 1128 io.read(&version, sizeof(version));1129 io.read(&ratio_ref, sizeof(ratio_ref));1130 io.read(&state_size_ref, sizeof(state_size_ref));1131 io.read(&n_embd_state_ref, sizeof(n_embd_state_ref));1132 io.read(&ns, sizeof(ns));1133 io.read(&n_layer_ref, sizeof(n_layer_ref));1134 1135 if (version != DSV4_COMP_STATE_VER) {1136 throw std::runtime_error("DSV4 compressor state version mismatch");1137 }1138 if (ratio_ref != ratio || state_size_ref != state_size || n_embd_state_ref != n_embd_state) {1139 throw std::runtime_error("DSV4 compressor state metadata mismatch");1140 }1141 if (n_layer_ref != layers.size()) {1142 throw std::runtime_error("DSV4 compressor state layer count mismatch");1143 }1144 1145 uint32_t s0;1146 dsv4_state_dst_stream_range(n_stream, seq_id, ns, s0);1147 1148 for (const auto & layer : layers) {1149 uint32_t il_ref;1150 io.read(&il_ref, sizeof(il_ref));1151 if (il_ref != layer.il) {1152 throw std::runtime_error("DSV4 compressor state layer id mismatch");1153 }1154 1155 dsv4_state_read_tensor_streams(io, layer.kv, state_size, state_size, s0, ns);1156 dsv4_state_read_tensor_streams(io, layer.score, state_size, state_size, s0, ns);1157 }1158}1159 1160ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const {1161 const int32_t ids = map_layer_ids.at(il);1162 ggml_tensor * state = layers[ids].kv;1163 1164 return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);1165}1166 1167ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const {1168 const int32_t ids = map_layer_ids.at(il);1169 ggml_tensor * state = layers[ids].score;1170 1171 return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);1172}1173 1174ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {1175 ggml_tensor * state = get_kv_all(ctx, il);1176 const size_t row_size = ggml_row_size(state->type, state->ne[0]);1177 1178 return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);1179}1180 1181ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const {1182 ggml_tensor * state = get_score_all(ctx, il);1183 const size_t row_size = ggml_row_size(state->type, state->ne[0]);1184 1185 return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);1186}1187 1188ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {1189 return ggml_set_rows(ctx, get_kv_all(ctx, il), cur, idxs);1190}1191 1192ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {1193 return ggml_set_rows(ctx, get_score_all(ctx, il), cur, idxs);1194}1195 1196size_t llama_dsv4_comp_state::total_size() const {1197 size_t size = 0;1198 1199 for (const auto & [_, buf] : ctxs_bufs) {1200 size += ggml_backend_buffer_get_size(buf.get());