return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);
}
+static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) {
+ GGML_ASSERT(ggml_is_contiguous(tensor));
+ GGML_ASSERT(tensor->ne[3] == 1);
+ GGML_ASSERT(stream < (uint32_t) tensor->ne[2]);
+
+ const size_t stream_size = tensor->nb[2];
+ ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size);
+}
+
static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {
if (n_stream <= 1) {
return 0;
__func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
}
-void llama_dsv4_comp_state::clear(bool data) {
+void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
if (!data) {
return;
}
+ if (seq_id >= 0) {
+ GGML_ASSERT((uint32_t) seq_id < n_stream);
+ for (const auto & layer : layers) {
+ dsv4_clear_tensor_stream(layer.kv, (uint32_t) seq_id);
+ dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id);
+ }
+ return;
+ }
+
for (auto & [_, buf] : ctxs_bufs) {
ggml_backend_buffer_clear(buf.get(), 0);
}
// graph does not necessarily overwrite; uninitialized buffer contents would
// otherwise leak in (instance-specific garbage) and corrupt recall. Zero all
// compressed buffers up front so reads of un-written rows are deterministic.
- clear_compressed(true);
+ clear_compressed(-1, true);
}
llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
void llama_kv_cache_dsv4::clear(bool data) {
kv_raw->clear(data);
- clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows
+ clear_compressed(-1, true); // DSV4 compressed buffers must never expose stale/uninit rows
}
bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
if (res) {
- clear_compressed(true);
+ clear_compressed(seq_id, true);
}
return res;
void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1);
- clear_compressed(true);
}
void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
+ GGML_ASSERT(seq_id >= 0 && (uint32_t) seq_id < n_seq_max);
+
kv_raw->seq_keep(seq_id);
- clear_compressed(true);
+
+ for (llama_seq_id id = 0; id < (llama_seq_id) n_seq_max; ++id) {
+ if (id == seq_id) {
+ continue;
+ }
+
+ kv_raw->seq_rm(id, -1, -1);
+ clear_compressed(id, true);
+ }
}
void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
kv_raw->seq_add(seq_id, p0, p1, shift);
- clear_compressed(true);
}
void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
kv_raw->seq_div(seq_id, p0, p1, d);
- clear_compressed(true);
}
llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const {
return lid_state.get();
}
-void llama_kv_cache_dsv4::clear_compressed(bool data) {
- kv_csa->clear(data);
- kv_hca->clear(data);
- kv_lid->clear(data);
- csa_state->clear(data);
- hca_state->clear(data);
- lid_state->clear(data);
+void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
+ if (seq_id < 0) {
+ kv_csa->clear(data);
+ kv_hca->clear(data);
+ kv_lid->clear(data);
+ } else {
+ GGML_ASSERT((uint32_t) seq_id < n_seq_max);
+
+ const auto clear_seq = [seq_id, data](llama_kv_cache * kv) {
+ kv->seq_rm(seq_id, -1, -1);
+
+ if (data) {
+ for (uint32_t il : kv->get_layer_ids()) {
+ dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id);
+ }
+ }
+ };
+
+ clear_seq(kv_csa.get());
+ clear_seq(kv_hca.get());
+ clear_seq(kv_lid.get());
+ }
+
+ csa_state->clear(seq_id, data);
+ hca_state->clear(seq_id, data);
+ lid_state->clear(seq_id, data);
}
//
const char * name,
const llama_memory_i::layer_filter_cb & filter);
- void clear(bool data);
+ void clear(llama_seq_id seq_id, bool data);
uint32_t get_ratio() const;
uint32_t get_state_size() const;
// DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches.
// The compressed caches are storage only; DSV4-specific visibility and block
// planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4.
+// FIXME: currently the cache only supports non-unified mode even if unified flag is passed
+// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819
class llama_kv_cache_dsv4 : public llama_memory_i {
public:
std::unique_ptr<llama_dsv4_comp_state> hca_state;
std::unique_ptr<llama_dsv4_comp_state> lid_state;
- void clear_compressed(bool data);
+ void clear_compressed(llama_seq_id seq_id, bool data);
};
// DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept
}
-// Test 2: state load
+// Test 2: sequence removal isolation
+// - decode the same prefix into two sequences
+// - remove sequence 0
+// - verify that sequence 1 remains unchanged
+static bool test_seq_rm_isolated(
+ struct llama_model * model,
+ const struct common_params & params,
+ const llama_tokens & tokens) {
+ auto params_ctx = common_context_params_to_llama(params);
+ params_ctx.n_ctx = 256;
+ params_ctx.n_seq_max = 2;
+ params_ctx.kv_unified = true;
+
+ auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
+ if (!ctx) {
+ LOG_ERR("%s: failed to create context\n", __func__);
+ return false;
+ }
+
+ LOG("\n=== Test 2: sequence removal isolation ===\n");
+
+ const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128;
+ for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
+ llama_batch_ptr batch(n_tokens, 0, 1);
+ for (size_t i = 0; i < n_tokens; ++i) {
+ common_batch_add(batch.get(), tokens[i], i, { seq_id }, false);
+ }
+
+ if (llama_decode(ctx.get(), batch.get())) {
+ LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id);
+ return false;
+ }
+ }
+
+ const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) {
+ const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id);
+ if (state_size == 0) {
+ LOG_ERR("%s: sequence state is empty\n", __func__);
+ return false;
+ }
+
+ state.resize(state_size);
+ const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id);
+ if (ncopy != state.size()) {
+ LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
+ __func__, ncopy, state.size());
+ return false;
+ }
+
+ return true;
+ };
+
+ std::vector<uint8_t> state_before;
+ if (!get_seq_state(1, state_before)) {
+ return false;
+ }
+
+ if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
+ LOG_ERR("%s: failed to remove sequence 0\n", __func__);
+ return false;
+ }
+
+ std::vector<uint8_t> state_after;
+ if (!get_seq_state(1, state_after)) {
+ return false;
+ }
+
+ if (state_before != state_after) {
+ LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__);
+ return false;
+ }
+
+ LOG("PASS\n");
+ return true;
+}
+
+
+// Test 3: state load
// - create a new context
// - load state from file
// - replay the last prompt token
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
- LOG("\n=== Test 2: state load ===\n");
+ LOG("\n=== Test 3: state load ===\n");
// Load state from file
llama_tokens unused_sts(tokens.size());
}
-// Test 3: seq copy (host)
+// Test 4: seq copy (host)
// - create a multi-seq context
// - load state from file
// - replay the last prompt token
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
- LOG("\n=== Test 3: seq copy (host) ===\n");
+ LOG("\n=== Test 4: seq copy (host) ===\n");
// Load state from file
llama_tokens unused_sts(tokens.size());
}
-// Test 4: seq copy (device)
+// Test 5: seq copy (device)
// - create a multi-seq context
// - load state from file
// - replay the last prompt token
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
- LOG("\n=== Test 4: seq copy (device) ===\n");
+ LOG("\n=== Test 5: seq copy (device) ===\n");
// Load state from file
llama_tokens unused_sts(tokens.size());
return 1;
}
- // Test 2: state load
+ // Test 2: sequence removal isolation
+ if (!test_seq_rm_isolated(model, params, tokens)) {
+ return 1;
+ }
+
+ // Test 3: state load
if (!test_state_load(model, params, tokens, result_baseline)) {
return 1;
}
- // Test 3: seq copy (host)
+ // Test 4: seq copy (host)
if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
return 1;
}
- // Test 4: seq copy (device)
+ // Test 5: seq copy (device)
if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
return 1;
}
overview / pkg / ggml / sources / llama.cpp / commitdiff