Long-context: rolling KeyDiff KV-cache compaction + chunked prefill

cactus-engine/CMakeLists.txt

@@ -15,6 +15,7 @@ set(ENGINE_SOURCES
     src/sp.cpp
     src/constraints.cpp
     src/model.cpp
+    src/kv_compress.cpp
     src/model_npu.cpp
     src/engine_image.cpp
     src/index.cpp

cactus-engine/src/complete.cpp

@@ -77,23 +77,6 @@ std::vector<ToolConstraintSpec> build_tool_constraint_specs(const std::vector<To
     return specs;
 }
 
-void strip_thinking_from_cache(CactusModelHandle* handle,
-                               const std::vector<uint32_t>& generated_tokens,
-                               size_t prompt_len) {
-    const auto& cfg = handle->model->get_config();
-    uint32_t open_id = cfg.channel_open_token_id;
-    uint32_t close_id = cfg.channel_close_token_id;
-    auto ranges = find_channel_token_ranges(generated_tokens, prompt_len,
-                                            open_id, close_id);
-    if (ranges.empty()) return;
-
-    handle->model->remove_thinking_tokens(ranges);
-    for (auto it = ranges.rbegin(); it != ranges.rend(); ++it) {
-        auto start = handle->processed_tokens.begin() + it->first;
-        handle->processed_tokens.erase(start, start + it->second);
-    }
-}
-
 void setup_tool_constraints(CactusModelHandle* handle, const std::vector<ToolFunction>& tools,
                            bool force_tools, float& temperature) {
     if (!force_tools || tools.empty()) return;
@@ -918,10 +901,6 @@ int cactus_complete(
             handle->model->clear_tool_constraints();
         }
 
-        if (prompt.model_type == Config::ModelType::GEMMA4 && prompt.options.enable_thinking_if_supported && !generated_tokens.empty()) {
-            strip_thinking_from_cache(handle, generated_tokens, prompt.tokens.size());
-        }
-
         auto end_time = std::chrono::high_resolution_clock::now();
         double total_time_ms = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time).count() / 1000.0;
 
@@ -939,7 +918,7 @@ int cactus_complete(
         std::string thinking_text;
         if (prompt.model_type == Config::ModelType::GEMMA4 || prompt.options.enable_thinking_if_supported) {
             std::string stripped_content;
-            strip_thinking_block(regular_response, thinking_text, stripped_content);
+            partition_thinking_response(regular_response, thinking_text, stripped_content);
             regular_response = stripped_content;
             if (!prompt.options.enable_thinking_if_supported) {
                 thinking_text.clear();
@@ -982,7 +961,7 @@ int cactus_complete(
         std::string result = construct_response_json(primary_response, primary_function_calls, time_to_first_token,
                                                      total_time_ms, prefill_tps, decode_tps, prompt_tokens,
                                                      completion_tokens, confidence, handoff_succeeded,
-                                                     thinking_text);
+                                                     thinking_text, {}, response_text);
 
         if (result.length() >= buffer_size) {
             handle_error_response("Response buffer too small", response_buffer, buffer_size);

cactus-engine/src/engine.h

@@ -11,6 +11,7 @@
 #include <limits>
 
 #include "cactus_graph.h"
+#include "kv_compress.h"
 
 class CactusGraph;
 
@@ -172,6 +173,15 @@ struct Config {
     std::vector<std::string> layer_types;
     size_t conv_L_cache = 0;
 
+    // Rolling bounded KV compaction (default ON, 4096 -> 2048). Override at runtime with
+    // CACTUS_KV_COMPRESS_AT (trigger) / CACTUS_KV_COMPRESS_TO (target); CACTUS_KV_COMPRESS_AT=0 disables.
+    bool kv_compress = true;
+    float kv_compress_recent_frac = 0.30f;
+    uint32_t kv_compress_sink = 4;
+    int32_t kv_compress_trigger_len = 4096;
+    int32_t kv_compress_target_len = 2048;
+    bool kv_compress_preserve_special = true;
+
     uint32_t altup_num_inputs = 4;
     uint32_t laurel_rank = 64;
     static constexpr uint32_t UNSET_U32 = UINT32_MAX;
@@ -219,15 +229,16 @@ struct Config {
     uint32_t audio_fft_length = 1024;
     uint32_t audio_token_id = 0;
     bool audio_fft_overdrive = false;
-    uint32_t channel_open_token_id = 100;
-    uint32_t channel_close_token_id = 101;
 
     static bool is_gemma_family(ModelType t) {
         return t == ModelType::GEMMA || t == ModelType::GEMMA3N || t == ModelType::GEMMA4;
     }
 
     bool from_json(const std::string& json_path);
     std::string to_json() const;
+    // Disable rolling unless 0 < target < trigger (when trigger > 0).
+    void validate_kv_compress();
+    bool parse_kv_compress_override(const char* trigger_env, const char* target_env);
 };
 
 
@@ -316,6 +327,7 @@ class Tokenizer {
     virtual uint32_t get_unk_token() const = 0;
     virtual uint32_t get_bos_token() const = 0;
     virtual uint32_t get_eos_token() const = 0;
+    virtual std::unordered_set<uint32_t> special_token_ids() const { return {}; }
     virtual bool has_chat_template() const { return has_chat_template_; }
     std::string get_default_stop_sequence() const;
 
@@ -370,6 +382,11 @@ class BPETokenizer : public Tokenizer {
     uint32_t get_unk_token() const override { return unk_token_id_; }
     uint32_t get_bos_token() const override { return bos_token_id_; }
     uint32_t get_eos_token() const override { return eos_token_id_; }
+    std::unordered_set<uint32_t> special_token_ids() const override {
+        std::unordered_set<uint32_t> ids;
+        for (const auto& kv : special_tokens_) ids.insert(kv.second);
+        return ids;
+    }
 
 private:
     std::unordered_map<std::string, uint32_t> token_to_id_;
@@ -422,6 +439,11 @@ class SPTokenizer : public Tokenizer {
     uint32_t get_unk_token() const override { return unk_token_id_; }
     uint32_t get_bos_token() const override { return bos_token_id_; }
     uint32_t get_eos_token() const override { return eos_token_id_; }
+    std::unordered_set<uint32_t> special_token_ids() const override {
+        std::unordered_set<uint32_t> ids;
+        for (const auto& kv : special_tokens_) ids.insert(kv.second);
+        return ids;
+    }
 
 private:
     struct TrieNode {
@@ -618,9 +640,13 @@ class Model {
     bool load_npu_vision_encoder(const std::string& model_path);
     bool has_npu_vision_encoder() const { return npu_vision_encoder_ != nullptr; }
 
-    void remove_thinking_tokens(const std::vector<std::pair<size_t, size_t>>& ranges);
     void compact_kv_cache() {}
 
+    void compress_kv_cache_keydiff(const cactus::kvcompress::Params& params);
+    void maybe_roll_compact();
+    std::vector<size_t> compressible_layers() const;
+    void apply_kv_compress_env_override();
+
     void set_tool_constraints(const std::vector<ToolConstraintSpec>& tools);
     void clear_tool_constraints();
     void update_tool_constraints(uint32_t token_id);
@@ -685,7 +711,8 @@ class Model {
     void copy_component_outputs_to_chunk_inputs(const Component& source, Component& target, size_t token_index);
     void copy_component_outputs_to_chunk_inputs_range(const Component& source, Component& target, size_t token_offset);
     bool cache_states_compatible(const Component& source, const Component& target) const;
-    void copy_cache_states(const Component& source, Component& target, size_t logical_current = std::numeric_limits<size_t>::max());
+    void move_cache_states(Component& source, Component& target, size_t logical_current = std::numeric_limits<size_t>::max());
+    void set_cache_current_len(Component& comp, size_t len);
     void reset_component_cache_states(Component& comp);
     size_t component_chunk_tokens(const Component& comp, const std::string& input_name) const;
     size_t component_output_tokens(const Component& comp, const std::string& output_name) const;
@@ -748,6 +775,9 @@ class Model {
     std::unique_ptr<Tokenizer> tokenizer_;
     bool initialized_ = false;
     size_t cache_total_seq_len_ = 0;
+    std::vector<uint32_t> cache_token_ids_;        // token id per cache row (canonical head-0 view)
+    std::unordered_set<uint32_t> special_ids_;     // special-token ids force-kept during compaction
+    cactus::kvcompress::SpecialRowTracker special_rows_;  // per-(layer,head) special rows for compaction protect
     size_t cache_max_seq_len_ = 4096;
     size_t last_logit_position_ = 0;
     double last_prefill_cache_copy_ms_ = 0.0;