diff --git a/cactus-engine/CMakeLists.txt b/cactus-engine/CMakeLists.txt index 1755844cf..24c7c8ca4 100644 --- a/cactus-engine/CMakeLists.txt +++ b/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 diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index a6c20e65b..c3a7a7b0a 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -77,23 +77,6 @@ std::vector build_tool_constraint_specs(const std::vector& 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& 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(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); diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index bfac47540..e63041b25 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -11,6 +11,7 @@ #include #include "cactus_graph.h" +#include "kv_compress.h" class CactusGraph; @@ -172,6 +173,15 @@ struct Config { std::vector 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,8 +229,6 @@ 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; @@ -228,6 +236,9 @@ struct Config { 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 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 special_token_ids() const override { + std::unordered_set ids; + for (const auto& kv : special_tokens_) ids.insert(kv.second); + return ids; + } private: std::unordered_map 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 special_token_ids() const override { + std::unordered_set 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>& ranges); void compact_kv_cache() {} + void compress_kv_cache_keydiff(const cactus::kvcompress::Params& params); + void maybe_roll_compact(); + std::vector compressible_layers() const; + void apply_kv_compress_env_override(); + void set_tool_constraints(const std::vector& 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::max()); + void move_cache_states(Component& source, Component& target, size_t logical_current = std::numeric_limits::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_; bool initialized_ = false; size_t cache_total_seq_len_ = 0; + std::vector cache_token_ids_; // token id per cache row (canonical head-0 view) + std::unordered_set 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; diff --git a/cactus-engine/src/kv_compress.cpp b/cactus-engine/src/kv_compress.cpp new file mode 100644 index 000000000..eff9c4eb6 --- /dev/null +++ b/cactus-engine/src/kv_compress.cpp @@ -0,0 +1,482 @@ +#include "kv_compress.h" + +#include +#include +#include +#include +#include + +#if defined(__ARM_NEON) +#include +#define CACTUS_KV_NEON 1 +#else +#define CACTUS_KV_NEON 0 +#endif + +namespace cactus { +namespace kvcompress { + +namespace { +// Toggled off by tests to exercise/compare the scalar fallback on a NEON build. +bool g_use_simd = true; +} + +void kv_set_simd(bool on) { g_use_simd = on; } + +namespace { + +// Round half-to-even like numpy. recent_frac widens from float to a slightly-off double, so snap +// near-halves first to avoid flipping an exact .5 boundary the Python reference keeps. +long py_round(double x) { + double twice = x * 2.0; + double twice_rounded = std::nearbyint(twice); + if (std::fabs(twice - twice_rounded) < 1e-6) x = twice_rounded * 0.5; + double r = std::nearbyint(x); // FE_TONEAREST default == round-half-to-even + return static_cast(r); +} + +} // namespace + +void keydiff_score(const float* keys, size_t n, size_t head_dim, float* out) { + // double accumulation matches the float64 reference. + std::vector mu(head_dim, 0.0); + for (size_t i = 0; i < n; ++i) { + const float* k = keys + i * head_dim; + for (size_t d = 0; d < head_dim; ++d) mu[d] += static_cast(k[d]); + } + double mu_norm = 0.0; + for (size_t d = 0; d < head_dim; ++d) { + mu[d] /= static_cast(n); + mu_norm += mu[d] * mu[d]; + } + mu_norm = std::sqrt(mu_norm) + 1e-8; + for (size_t d = 0; d < head_dim; ++d) mu[d] /= mu_norm; + + for (size_t i = 0; i < n; ++i) { + const float* k = keys + i * head_dim; + double knorm = 0.0, dot = 0.0; + for (size_t d = 0; d < head_dim; ++d) { + double v = static_cast(k[d]); + knorm += v * v; + dot += v * mu[d]; + } + knorm = std::sqrt(knorm) + 1e-8; + out[i] = static_cast(-(dot / knorm)); + } +} + +std::vector keepset_for_head(const float* scores, size_t n, const Params& p) { + long budget = std::max(1, static_cast(p.abs_budget)); + long B = std::min(budget, static_cast(n)); + long sink = std::min(std::max(static_cast(p.sink), 0), static_cast(n)); + long n_recent = std::min(py_round(static_cast(p.recent_frac) * static_cast(B)), + static_cast(n)); + + auto in_range = [&](int idx) { return idx >= 0 && idx < static_cast(n); }; + std::set reserved; + for (long i = 0; i < sink; ++i) reserved.insert(i); + for (int idx : p.protect) if (in_range(idx)) reserved.insert(static_cast(idx)); + for (long i = static_cast(n) - n_recent; i < static_cast(n); ++i) + if (i >= 0) reserved.insert(i); + + // Reserved may exceed B: keep sink, protected, then recent, until B. + if (static_cast(reserved.size()) > B) { + std::vector ordered; + for (long i = 0; i < sink; ++i) ordered.push_back(i); + for (int idx : p.protect) if (in_range(idx)) ordered.push_back(static_cast(idx)); + for (long i = static_cast(n) - 1; i >= sink; --i) ordered.push_back(i); + reserved.clear(); + for (long i : ordered) { + if (i >= 0 && i < static_cast(n)) reserved.insert(i); + if (static_cast(reserved.size()) == B) break; + } + } + + std::set keep(reserved); + long remaining = B - static_cast(keep.size()); + if (remaining > 0) { + std::vector order(n); + std::iota(order.begin(), order.end(), 0L); + std::stable_sort(order.begin(), order.end(), [&](long a, long b) { + return scores[a] > scores[b]; + }); + for (long i : order) { + if (keep.count(i)) continue; + keep.insert(i); + if (--remaining == 0) break; + } + } + + std::vector result; + result.reserve(keep.size()); + for (long i : keep) result.push_back(static_cast(i)); + return result; +} + +namespace { + +std::vector rope_inv_freq(size_t head_dim, double rope_theta) { + std::vector inv(head_dim / 2); + for (size_t i = 0; i < inv.size(); ++i) + inv[i] = std::pow(rope_theta, -(2.0 * static_cast(i)) / static_cast(head_dim)); + return inv; +} + +RopeRotation rope_rot(const std::vector& inv_freq, double delta_pos) { + RopeRotation r; + r.cos.resize(inv_freq.size()); + r.sin.resize(inv_freq.size()); + for (size_t i = 0; i < inv_freq.size(); ++i) { + double a = delta_pos * inv_freq[i]; + r.cos[i] = std::cos(a); + r.sin[i] = std::sin(a); + } + return r; +} + +void rope_apply(float* row, const RopeRotation& r) { + size_t half = r.cos.size(); + for (size_t i = 0; i < half; ++i) { + double x1 = row[i], x2 = row[i + half]; + row[i] = static_cast(x1 * r.cos[i] - x2 * r.sin[i]); + row[i + half] = static_cast(x2 * r.cos[i] + x1 * r.sin[i]); + } +} + +// Apply the inverse rotation (negate the table angle): re-rope by +t given the un-rope table[t]. +void rope_apply_conj(float* row, const RopeRotation& r) { + size_t half = r.cos.size(); + for (size_t i = 0; i < half; ++i) { + double x1 = row[i], x2 = row[i + half]; + row[i] = static_cast(x1 * r.cos[i] + x2 * r.sin[i]); + row[i + half] = static_cast(x2 * r.cos[i] - x1 * r.sin[i]); + } +} + +size_t max_kept_index(const std::vector>& kept) { + size_t m = 0; + for (const auto& kh : kept) for (int idx : kh) m = std::max(m, static_cast(idx)); + return m; +} + +} // namespace + +void rope_rotate_row(float* row, size_t head_dim, double rope_theta, double delta_pos) { + rope_apply(row, rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos)); +} + +std::vector unrope_table(size_t n, size_t head_dim, double rope_theta) { + auto inv = rope_inv_freq(head_dim, rope_theta); + std::vector table(n); + for (size_t t = 0; t < n; ++t) table[t] = rope_rot(inv, -static_cast(t)); + return table; +} + +namespace { + +// Matches the engine's KV quant convention (quantize_group_fp16_to_int8 in +// cactus-kernels/src/quants.cpp): scale floored at 1e-10, roundf, clamp to [-128, 127]. +void requant_row(const float* row, int8_t* dst, float* dsc, size_t head_dim, size_t group_size) { + size_t groups = (head_dim + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + size_t lo = g * group_size, hi = std::min(head_dim, lo + group_size); + float amax = 0.0f; + for (size_t d = lo; d < hi; ++d) amax = std::max(amax, std::fabs(row[d])); + float scale = amax / 127.0f; + if (scale < 1e-10f) scale = 1e-10f; + dsc[g] = scale; + float inv = 1.0f / scale; + for (size_t d = lo; d < hi; ++d) { + int32_t q = static_cast(std::roundf(row[d] * inv)); + q = std::max(-128, std::min(127, q)); + dst[d] = static_cast(q); + } + } +} + +void dequant_row(const int8_t* src, const float* scale, size_t head_dim, size_t group_size, + float* out) { +#if CACTUS_KV_NEON + if (g_use_simd) { + size_t groups = (head_dim + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + size_t lo = g * group_size, hi = std::min(head_dim, lo + group_size); + float32x4_t scv = vdupq_n_f32(scale[g]); + size_t d = lo; + for (; d + 4 <= hi; d += 4) { + int32_t four; + std::memcpy(&four, src + d, 4); + int16x8_t i16 = vmovl_s8(vreinterpret_s8_s32(vdup_n_s32(four))); + vst1q_f32(out + d, vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(i16))), scv)); + } + for (; d < hi; ++d) out[d] = static_cast(src[d]) * scale[g]; + } + return; + } +#endif + for (size_t d = 0; d < head_dim; ++d) out[d] = static_cast(src[d]) * scale[d / group_size]; +} + +void rotate_int8_row_rot(int8_t* int8, float* scale, size_t head_dim, size_t group_size, + const RopeRotation& rot) { + std::vector row(head_dim); + dequant_row(int8, scale, head_dim, group_size, row.data()); + rope_apply(row.data(), rot); + requant_row(row.data(), int8, scale, head_dim, group_size); +} + +// fill_post(h, post) gathers head h's post-RoPE rows; scoring is shared across fp16/int8. +template +std::vector> keepsets_per_head(size_t n, size_t kv_heads, size_t head_dim, + const std::vector& unrope, + const Params& p, FillPost fill_post, + const std::vector>& protect_per_head) { + const bool per_head_protect = protect_per_head.size() == kv_heads; + std::vector post(n * head_dim), pre(n * head_dim), scores(n); + std::vector> out; + out.reserve(kv_heads); + Params ph = p; + for (size_t h = 0; h < kv_heads; ++h) { + fill_post(h, post.data()); + for (size_t t = 0; t < n; ++t) { + const float* src = post.data() + t * head_dim; + float* dst = pre.data() + t * head_dim; + for (size_t d = 0; d < head_dim; ++d) dst[d] = src[d]; + rope_apply(dst, unrope[t]); + } + keydiff_score(pre.data(), n, head_dim, scores.data()); + if (per_head_protect) ph.protect = protect_per_head[h]; + out.push_back(keepset_for_head(scores.data(), n, per_head_protect ? ph : p)); + } + return out; +} + +} // namespace + +void compact_fp16(uint16_t* key_rows_u, uint16_t* val_rows_u, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, + const std::vector& unrope) { + __fp16* key_rows = reinterpret_cast<__fp16*>(key_rows_u); + __fp16* val_rows = reinterpret_cast<__fp16*>(val_rows_u); + std::vector krow(head_dim), vrow(head_dim); + for (size_t h = 0; h < kv_heads; ++h) { + const std::vector& kept = kept_per_head[h]; + for (size_t rank = 0; rank < kept.size(); ++rank) { + int abs_pos = kept[rank]; + const __fp16* ksrc = key_rows + (static_cast(abs_pos) * kv_heads + h) * head_dim; + const __fp16* vsrc = val_rows + (static_cast(abs_pos) * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { krow[d] = ksrc[d]; vrow[d] = vsrc[d]; } + // renumber abs_pos -> rank: un-rope by -abs_pos then re-rope by +rank (no per-row trig) + rope_apply(krow.data(), unrope[static_cast(abs_pos)]); + rope_apply_conj(krow.data(), unrope[rank]); + __fp16* kdst = key_rows + (rank * kv_heads + h) * head_dim; + __fp16* vdst = val_rows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + kdst[d] = static_cast<__fp16>(krow[d]); + vdst[d] = static_cast<__fp16>(vrow[d]); + } + } + } +} + +void compact_fp16(uint16_t* key_rows, uint16_t* val_rows, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, double rope_theta) { + compact_fp16(key_rows, val_rows, kv_heads, head_dim, kept_per_head, + unrope_table(max_kept_index(kept_per_head) + 1, head_dim, rope_theta)); +} + +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, + const std::vector& unrope, bool renumber) { + size_t groups = (head_dim + group_size - 1) / group_size; + size_t int8_stride = kv_heads * head_dim; + size_t scale_stride = kv_heads * groups; + std::vector row(head_dim); + for (size_t h = 0; h < kv_heads; ++h) { + const std::vector& kept = kept_per_head[h]; + for (size_t rank = 0; rank < kept.size(); ++rank) { + int abs_pos = kept[rank]; + size_t src_t = static_cast(abs_pos); + const int8_t* src = int8_rows + src_t * int8_stride + h * head_dim; + const float* ssc = scale_rows + src_t * scale_stride + h * groups; + dequant_row(src, ssc, head_dim, group_size, row.data()); + if (renumber) { // un-rope by -abs_pos then re-rope by +rank (no per-row trig) + rope_apply(row.data(), unrope[static_cast(abs_pos)]); + rope_apply_conj(row.data(), unrope[rank]); + } + int8_t* dst = int8_rows + rank * int8_stride + h * head_dim; + float* dsc = scale_rows + rank * scale_stride + h * groups; + requant_row(row.data(), dst, dsc, head_dim, group_size); + } + } +} + +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, double rope_theta, + bool renumber) { + std::vector unrope; + if (renumber) unrope = unrope_table(max_kept_index(kept_per_head) + 1, head_dim, rope_theta); + compact_int8(int8_rows, scale_rows, kv_heads, head_dim, group_size, kept_per_head, unrope, renumber); +} + +void rotate_int8_row(int8_t* int8, float* scale, size_t head_dim, size_t group_size, + double rope_theta, double delta_pos) { + rotate_int8_row_rot(int8, scale, head_dim, group_size, + rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos)); +} + +void rerope_recent_fp16(uint16_t* key_rows_u, size_t kv_heads, size_t head_dim, + size_t lo, size_t hi, double rope_theta, double delta_pos) { + if (delta_pos == 0.0 || hi <= lo) return; + __fp16* key_rows = reinterpret_cast<__fp16*>(key_rows_u); + RopeRotation rot = rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos); + std::vector row(head_dim); + for (size_t t = lo; t < hi; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + __fp16* r = key_rows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) row[d] = r[d]; + rope_apply(row.data(), rot); + for (size_t d = 0; d < head_dim; ++d) r[d] = static_cast<__fp16>(row[d]); + } +} + +void rerope_recent_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, size_t head_dim, + size_t group_size, size_t lo, size_t hi, double rope_theta, double delta_pos) { + if (delta_pos == 0.0 || hi <= lo) return; + size_t groups = (head_dim + group_size - 1) / group_size; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + RopeRotation rot = rope_rot(rope_inv_freq(head_dim, rope_theta), delta_pos); + for (size_t t = lo; t < hi; ++t) + for (size_t h = 0; h < kv_heads; ++h) + rotate_int8_row_rot(int8_rows + t * int8_stride + h * head_dim, + scale_rows + t * scale_stride + h * groups, + head_dim, group_size, rot); +} + +std::vector> keepsets_from_fp16(const uint16_t* key_rows_u, size_t n, + size_t kv_heads, size_t head_dim, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head) { + const __fp16* key_rows = reinterpret_cast(key_rows_u); + return keepsets_per_head(n, kv_heads, head_dim, unrope, p, [&](size_t h, float* post) { + for (size_t t = 0; t < n; ++t) { + const __fp16* src = key_rows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) post[t * head_dim + d] = src[d]; + } + }, protect_per_head); +} + +std::vector> keepsets_from_fp16(const uint16_t* key_rows, size_t n, + size_t kv_heads, size_t head_dim, + double rope_theta, const Params& p) { + return keepsets_from_fp16(key_rows, n, kv_heads, head_dim, + unrope_table(n, head_dim, rope_theta), p); +} + +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head) { + size_t groups = (head_dim + group_size - 1) / group_size; + size_t int8_stride = kv_heads * head_dim; + size_t scale_stride = kv_heads * groups; + return keepsets_per_head(n, kv_heads, head_dim, unrope, p, [&](size_t h, float* post) { + for (size_t t = 0; t < n; ++t) { + const int8_t* src = int8_rows + t * int8_stride + h * head_dim; + const float* ssc = scale_rows + t * scale_stride + h * groups; + dequant_row(src, ssc, head_dim, group_size, post + t * head_dim); + } + }, protect_per_head); +} + +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, double rope_theta, + const Params& p) { + return keepsets_from_int8(int8_rows, scale_rows, n, kv_heads, head_dim, group_size, + unrope_table(n, head_dim, rope_theta), p); +} + +std::vector remap_rows_through_kept(const std::vector& rows, const std::vector& kept) { + std::vector out; + size_t i = 0; + for (size_t rank = 0; rank < kept.size(); ++rank) { + while (i < rows.size() && rows[i] < kept[rank]) ++i; + if (i < rows.size() && rows[i] == kept[rank]) out.push_back(static_cast(rank)); + } + return out; +} + +void SpecialRowTracker::add_appended(size_t layer, size_t kv_heads, + const std::vector& appended_rows) { + if (layer_rows_.size() <= layer) layer_rows_.resize(layer + 1); + auto& heads = layer_rows_[layer]; + if (heads.empty()) heads.resize(kv_heads); + for (auto& rows : heads) + rows.insert(rows.end(), appended_rows.begin(), appended_rows.end()); // appended >= existing, stays sorted +} + +const std::vector>& SpecialRowTracker::protect(size_t layer) const { + static const std::vector> kEmpty; + return layer < layer_rows_.size() ? layer_rows_[layer] : kEmpty; +} + +size_t SpecialRowTracker::max_reserved(size_t layer, size_t sink, + const std::vector& appended) const { + const auto& heads = protect(layer); + if (heads.empty()) { + size_t extra = 0; + for (int r : appended) if (static_cast(r) >= sink) ++extra; + return sink + extra; + } + size_t worst = 0; + for (const auto& rows : heads) { + size_t count = sink; + for (int r : rows) if (static_cast(r) >= sink) ++count; + for (int r : appended) if (static_cast(r) >= sink) ++count; + worst = std::max(worst, count); + } + return worst; +} + +void SpecialRowTracker::remap(size_t layer, const std::vector>& kept_per_head) { + if (layer >= layer_rows_.size()) return; + auto& heads = layer_rows_[layer]; + for (size_t h = 0; h < heads.size() && h < kept_per_head.size(); ++h) + heads[h] = remap_rows_through_kept(heads[h], kept_per_head[h]); +} + +bool is_sliding_layer(const std::vector& layer_types, size_t li) { + return li < layer_types.size() && layer_types[li].find("sliding") != std::string::npos; +} + +std::vector physical_compressible_layers(const std::vector& layer_types, + size_t num_layers, size_t num_kv_shared) { + auto is_full = [&](size_t i) { return !is_sliding_layer(layer_types, i); }; + std::vector full; + for (size_t i = 0; i < num_layers; ++i) if (is_full(i)) full.push_back(i); + + if (num_kv_shared == 0 || num_kv_shared >= num_layers) return full; + size_t first_shared = num_layers - num_kv_shared; + + std::set sources; + std::vector> last_of_type; + for (size_t i = 0; i < first_shared && i < layer_types.size(); ++i) { + bool found = false; + for (auto& e : last_of_type) if (e.first == layer_types[i]) { e.second = i; found = true; break; } + if (!found) last_of_type.emplace_back(layer_types[i], i); + } + for (auto& e : last_of_type) sources.insert(e.second); + + std::vector out; + for (size_t i : full) if (i < first_shared && !sources.count(i)) out.push_back(i); + return out; +} + +} // namespace kvcompress +} // namespace cactus diff --git a/cactus-engine/src/kv_compress.h b/cactus-engine/src/kv_compress.h new file mode 100644 index 000000000..c33a5e659 --- /dev/null +++ b/cactus-engine/src/kv_compress.h @@ -0,0 +1,127 @@ +#ifndef CACTUS_KV_COMPRESS_H +#define CACTUS_KV_COMPRESS_H + +#include +#include +#include +#include +#include + +namespace cactus { +namespace kvcompress { + +struct Params { + float recent_frac = 0.30f; + size_t sink = 4; + int abs_budget = 0; // per (layer, kv-head) keep budget, clamped to [1, n] + std::vector protect; // positions always kept (special tokens) +}; + +// Mirrors cactus-graph CacheMetadata. +struct CacheHeader { + uint64_t current_seq_len; + uint64_t max_seq_len; + uint64_t num_kv_heads; + uint64_t head_dim; + uint64_t sink_size; + uint64_t reserved[3]; +}; +static_assert(sizeof(CacheHeader) == 64, "CacheHeader must be 64 bytes"); + +void kv_set_simd(bool on); + +// score s_i = -cos(k_i, mean(k)); keys pre-RoPE. +void keydiff_score(const float* keys, size_t n, size_t head_dim, float* out); + +// Keep-set: sink + recent + top-score middle; ties by ascending index. +std::vector keepset_for_head(const float* scores, size_t n, const Params& p); + +// Keys stored POST-RoPE; rotating by delta_pos re-RoPEs to orig+delta_pos. +void rope_rotate_row(float* row, size_t head_dim, double rope_theta, double delta_pos); + +void rotate_int8_row(int8_t* int8, float* scale, size_t head_dim, size_t group_size, + double rope_theta, double delta_pos); + +struct RopeRotation { std::vector cos, sin; }; + +// Un-rope rotations for [0, n): row t rotates by -t. +std::vector unrope_table(size_t n, size_t head_dim, double rope_theta); + +// Gather survivors in rank order; re-rope K to 0..B-1, gather V unchanged. +void compact_fp16(uint16_t* key_rows, uint16_t* val_rows, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, + const std::vector& unrope); +void compact_fp16(uint16_t* key_rows, uint16_t* val_rows, size_t kv_heads, size_t head_dim, + const std::vector>& kept_per_head, double rope_theta); + +// renumber=true (K) re-ropes each row to its new rank; false (V) gathers as-is. +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, + const std::vector& unrope, bool renumber); +void compact_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, + size_t head_dim, size_t group_size, + const std::vector>& kept_per_head, double rope_theta, + bool renumber); + +// Rotate recent K rows [lo, hi) by delta_pos; sink and V stay fixed. +void rerope_recent_fp16(uint16_t* key_rows, size_t kv_heads, size_t head_dim, + size_t lo, size_t hi, double rope_theta, double delta_pos); +void rerope_recent_int8(int8_t* int8_rows, float* scale_rows, size_t kv_heads, size_t head_dim, + size_t group_size, size_t lo, size_t hi, double rope_theta, double delta_pos); + +// Per-head keep-sets from POST-RoPE keys. protect_per_head overrides Params::protect; empty -> shared. +std::vector> keepsets_from_fp16(const uint16_t* key_rows, size_t n, + size_t kv_heads, size_t head_dim, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head = {}); +std::vector> keepsets_from_fp16(const uint16_t* key_rows, size_t n, + size_t kv_heads, size_t head_dim, + double rope_theta, const Params& p); + +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, + const std::vector& unrope, + const Params& p, + const std::vector>& protect_per_head = {}); +std::vector> keepsets_from_int8(const int8_t* int8_rows, const float* scale_rows, + size_t n, size_t kv_heads, size_t head_dim, + size_t group_size, double rope_theta, + const Params& p); + +std::vector remap_rows_through_kept(const std::vector& rows, const std::vector& kept); + +// Per-(compressible layer, head) special rows; lets compaction re-protect each head's own specials. +class SpecialRowTracker { +public: + void clear() { tracked_len_ = 0; layer_rows_.clear(); valid_ = true; } + size_t tracked_len() const { return tracked_len_; } + void set_tracked_len(size_t len) { tracked_len_ = len; } + bool valid() const { return valid_; } + void invalidate() { valid_ = false; layer_rows_.clear(); } + + // appended_rows must be >= tracked_len (the still head-aligned region). + void add_appended(size_t layer, size_t kv_heads, const std::vector& appended_rows); + const std::vector>& protect(size_t layer) const; + size_t max_reserved(size_t layer, size_t sink, const std::vector& appended) const; + void remap(size_t layer, const std::vector>& kept_per_head); + +private: + size_t tracked_len_ = 0; + bool valid_ = true; + std::vector>> layer_rows_; // [layer][head] -> sorted rows +}; + +// Sliding (local) vs full-attention; selects local vs global theta. +bool is_sliding_layer(const std::vector& layer_types, size_t li); + +// Compressible layers: empty -> all (Qwen); num_kv_shared drops shared consumers+sources (Gemma -> {4,9}). +std::vector physical_compressible_layers(const std::vector& layer_types, + size_t num_layers, size_t num_kv_shared); + +} // namespace kvcompress +} // namespace cactus + +#endif // CACTUS_KV_COMPRESS_H diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index e2fb0fd49..7017cfe08 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -362,6 +362,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, CACTUS_LOG_ERROR("model", "Failed to load config.txt from: " << bundle_dir); return false; } + apply_kv_compress_env_override(); if (!load_manifest()) { CACTUS_LOG_ERROR("model", "Failed to load bundle manifest from: " << bundle_dir); return false; @@ -820,7 +821,7 @@ bool Model::cache_states_compatible(const Component& source, const Component& ta return true; } -void Model::copy_cache_states(const Component& source, Component& target, size_t logical_current) { +void Model::move_cache_states(Component& source, Component& target, size_t logical_current) { if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) { throw std::runtime_error("prefill and step cache states are not compatible"); } @@ -830,150 +831,35 @@ void Model::copy_cache_states(const Component& source, Component& target, size_t if (src.layer_key != dst.layer_key) { throw std::runtime_error("prefill and step cache layer mismatch: " + src.layer_key + " != " + dst.layer_key); } + int moved_src = -1, moved_dst = -1; for (auto [src_node, dst_node] : {std::pair{src.key_node_id, dst.key_node_id}, std::pair{src.value_node_id, dst.value_node_id}}) { - const auto& src_desc = source.graph->get_output_buffer(static_cast(src_node)); - const auto& dst_desc = target.graph->get_output_buffer(static_cast(dst_node)); - if (src_desc.precision != dst_desc.precision) { - std::ostringstream oss; - oss << "prefill and step cache precision mismatch at layer " << src.layer_key - << ": " << static_cast(src_desc.precision) - << " vs " << static_cast(dst_desc.precision); - throw std::runtime_error(oss.str()); - } - void* src_ptr = source.graph->get_output(static_cast(src_node)); - void* dst_ptr = target.graph->get_output(static_cast(dst_node)); - - const OpType src_op = source.graph->get_node_op_type(static_cast(src_node)); - const OpType dst_op = target.graph->get_node_op_type(static_cast(dst_node)); - if (src_op != dst_op) { - throw std::runtime_error( - "cache state op_type mismatch between prefill and step at layer " - + src.layer_key); - } - if (src_op == OpType::RECURRENT_CACHE_STATE) { - if (src_desc.byte_size != dst_desc.byte_size) { - throw std::runtime_error( - "recurrent cache buffer shape mismatch between prefill and step at layer " - + src.layer_key); - } - std::memcpy(dst_ptr, src_ptr, src_desc.byte_size); - continue; - } - if (src_op == OpType::CONV_CACHE_STATE) { - if (src_desc.byte_size != dst_desc.byte_size) { - throw std::runtime_error( - "conv cache buffer shape mismatch between prefill and step at layer " - + src.layer_key); + if (src_node < 0 || dst_node < 0) continue; + // Conv/recurrent caches share one node for key and value; move it only once. + if (src_node == moved_src && dst_node == moved_dst) continue; + moved_src = src_node; + moved_dst = dst_node; + target.graph->steal_cache_buffer(static_cast(dst_node), *source.graph, static_cast(src_node)); + if (target.graph->get_node_op_type(static_cast(dst_node)) == OpType::KV_CACHE_STATE && + logical_current != std::numeric_limits::max()) { + auto* meta = static_cast(target.graph->get_output(static_cast(dst_node))); + if (meta && logical_current < meta[0]) { + meta[4] = std::min(meta[4], logical_current); + meta[0] = logical_current; } - std::memcpy(dst_ptr, src_ptr, src_desc.byte_size); - continue; - } - if (src_desc.byte_size == dst_desc.byte_size && logical_current == std::numeric_limits::max()) { - std::memcpy(dst_ptr, src_ptr, src_desc.byte_size); - continue; - } - if (src_desc.precision == Precision::INT8) { - auto* src_meta = static_cast(src_ptr); - auto* dst_meta = static_cast(dst_ptr); - const size_t src_current = static_cast(src_meta[0]); - const size_t effective_current = std::min(src_current, logical_current); - const size_t src_max = static_cast(src_meta[1]); - const size_t kv_heads = static_cast(src_meta[2]); - const size_t head_dim = static_cast(src_meta[3]); - const size_t sink = static_cast(src_meta[4]); - if (src_max == 0 || kv_heads == 0 || head_dim == 0) { - throw std::runtime_error("prefill cache metadata is not initialized for layer " + src.layer_key); - } - const size_t groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; - const size_t int8_stride = kv_heads * head_dim; - const size_t scale_stride = kv_heads * groups; - const size_t row_bytes = int8_stride + scale_stride * sizeof(float); - const size_t dst_max = (dst_desc.byte_size - 64) / row_bytes; - if (dst_max == 0) { - throw std::runtime_error("step cache capacity is zero for layer " + src.layer_key); - } - const size_t dst_current = std::min(effective_current, dst_max); - dst_meta[0] = dst_current; - dst_meta[1] = dst_max; - dst_meta[2] = kv_heads; - dst_meta[3] = head_dim; - dst_meta[4] = std::min(sink, dst_current); - std::memset(static_cast(dst_ptr) + 64, 0, dst_desc.byte_size - 64); - - const auto* src_i8 = static_cast(src_ptr) + 64; - const auto* src_scales = reinterpret_cast( - static_cast(src_ptr) + 64 + src_max * int8_stride); - auto* dst_i8 = static_cast(dst_ptr) + 64; - auto* dst_scales = reinterpret_cast( - static_cast(dst_ptr) + 64 + dst_max * int8_stride); - auto copy_rows = [&](size_t dst_row, size_t src_row, size_t rows) { - if (rows == 0) return; - std::memcpy( - dst_i8 + dst_row * int8_stride, - src_i8 + src_row * int8_stride, - rows * int8_stride); - std::memcpy( - dst_scales + dst_row * scale_stride, - src_scales + src_row * scale_stride, - rows * scale_stride * sizeof(float)); - }; - if (effective_current <= dst_max) { - copy_rows(0, 0, effective_current); - } else { - const size_t copied_sink = std::min(sink, dst_max); - const size_t tail_rows = dst_max - copied_sink; - copy_rows(0, 0, copied_sink); - if (tail_rows > 0) copy_rows(copied_sink, effective_current - tail_rows, tail_rows); - } - continue; - } - if (PrecisionTraits::is_cq(src_desc.precision) || src_desc.byte_size < 64 || dst_desc.byte_size < 64) { - std::ostringstream oss; - oss << "prefill and step cache buffer mismatch at layer " << src.layer_key - << ": " << src_desc.byte_size << " bytes vs " << dst_desc.byte_size << " bytes"; - throw std::runtime_error(oss.str()); } + } + } +} - auto* src_meta = static_cast(src_ptr); - auto* dst_meta = static_cast(dst_ptr); - const size_t src_current = static_cast(src_meta[0]); - const size_t effective_current = std::min(src_current, logical_current); - const size_t src_max = static_cast(src_meta[1]); - const size_t kv_heads = static_cast(src_meta[2]); - const size_t head_dim = static_cast(src_meta[3]); - const size_t sink = static_cast(src_meta[4]); - if (src_max == 0 || kv_heads == 0 || head_dim == 0) { - throw std::runtime_error("prefill cache metadata is not initialized for layer " + src.layer_key); - } - const size_t row_bytes = kv_heads * head_dim * PrecisionTraits::size_of(src_desc.precision); - const size_t dst_max = (dst_desc.byte_size - 64) / row_bytes; - if (dst_max == 0) { - throw std::runtime_error("step cache capacity is zero for layer " + src.layer_key); - } - const size_t dst_current = std::min(effective_current, dst_max); - dst_meta[0] = dst_current; - dst_meta[1] = dst_max; - dst_meta[2] = kv_heads; - dst_meta[3] = head_dim; - dst_meta[4] = std::min(sink, dst_current); - std::memset(static_cast(dst_ptr) + 64, 0, dst_desc.byte_size - 64); - const auto* src_rows = static_cast(src_ptr) + 64; - auto* dst_rows = static_cast(dst_ptr) + 64; - if (effective_current <= dst_max) { - std::memcpy(dst_rows, src_rows, effective_current * row_bytes); - } else { - const size_t copied_sink = std::min(sink, dst_max); - const size_t tail_rows = dst_max - copied_sink; - if (copied_sink > 0) { - std::memcpy(dst_rows, src_rows, copied_sink * row_bytes); - } - if (tail_rows > 0) { - std::memcpy( - dst_rows + copied_sink * row_bytes, - src_rows + (effective_current - tail_rows) * row_bytes, - tail_rows * row_bytes); - } - } +void Model::set_cache_current_len(Component& comp, size_t len) { + for (const auto& state : comp.cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; + auto* meta = static_cast(comp.graph->get_output(static_cast(node_id))); + if (!meta || len >= meta[0]) continue; + meta[4] = std::min(meta[4], len); + meta[0] = len; } } } @@ -1043,26 +929,31 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector 0 ? std::min(chunk_size, component_tokens) : component_tokens; if (effective_chunk != component_tokens) effective_chunk = component_tokens; size_t whole_chunks_end = (tokens.size() / effective_chunk) * effective_chunk; - const bool has_recurrent_state = [&]() { + auto any_cache_node = [&](auto predicate) { if (!decoder_prefill_->graph) return false; for (const auto& state : decoder_prefill_->cache_states) { for (int node_id : {state.key_node_id, state.value_node_id}) { if (node_id < 0) continue; - if (decoder_prefill_->graph->get_node_op_type(static_cast(node_id)) - == OpType::RECURRENT_CACHE_STATE) { - return true; - } + if (predicate(static_cast(node_id))) return true; } } return false; - }(); + }; + const bool has_recurrent_state = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::RECURRENT_CACHE_STATE; + }); if (has_recurrent_state && whole_chunks_end > effective_chunk) { whole_chunks_end = effective_chunk; } + const bool has_sliding_window_cache = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::KV_CACHE_STATE + && decoder_prefill_->graph->get_node_window_size(id) > 0; + }); const size_t tail_tokens = tokens.size() - whole_chunks_end; const size_t padding_cutoff = std::max(1, effective_chunk / 16); const bool pad_tail = family_ != "lfm2_vl" && !has_recurrent_state + && !has_sliding_window_cache && tail_tokens >= padding_cutoff; const size_t executable_tokens = whole_chunks_end + (pad_tail ? effective_chunk : 0); if (executable_tokens == 0) { @@ -1123,7 +1014,7 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vectorinput_buffers[i].begin(), decoder_->input_buffers[i].end(), 0); } auto copy_start = std::chrono::high_resolution_clock::now(); - copy_cache_states(*decoder_prefill_, *decoder_, start_position + result.logical_tokens); + move_cache_states(*decoder_prefill_, *decoder_, start_position + result.logical_tokens); auto copy_end = std::chrono::high_resolution_clock::now(); last_prefill_cache_copy_ms_ = std::chrono::duration_cast(copy_end - copy_start).count() / 1000.0; } @@ -1492,6 +1383,7 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, if (tokens.empty() || !decoder_ || cache_total_seq_len_ != 0) { return false; } + cache_token_ids_ = tokens; if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { context_tokens_.insert(context_tokens_.end(), tokens.begin(), tokens.end()); run_full_context_text(); @@ -1516,13 +1408,15 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, if (decoder_prefill_) { chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), true); if (chunked.logical_tokens == tokens.size() && chunked.padding_tokens > 0 && tokens.size() > 0) { - copy_cache_states(*decoder_prefill_, *decoder_, tokens.size() - 1); + // Cache already moved into the step component; drop the padded last row and re-run it for logits. + set_cache_current_len(*decoder_, tokens.size() - 1); cache_total_seq_len_ = tokens.size() - 1; run_step(tokens.back(), cache_total_seq_len_, true); ++cache_total_seq_len_; out_token = argmax_last_logits(); record_sampled_token(out_token); last_prefill_scalar_tail_tokens_ = 1; + maybe_roll_compact(); return true; } cache_total_seq_len_ += chunked.logical_tokens; @@ -1538,6 +1432,7 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, out_token = argmax_last_logits(); } record_sampled_token(out_token); + maybe_roll_compact(); return true; } @@ -1549,6 +1444,7 @@ void Model::prefill(const std::vector& tokens, size_t /*chunk_size*/, context_tokens_.insert(context_tokens_.end(), tokens.begin(), tokens.end()); if (!context_tokens_.empty()) run_full_context_text(); cache_total_seq_len_ = context_tokens_.size(); + cache_token_ids_ = context_tokens_; return; } ChunkedPrefillResult chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), prepare_decode); @@ -1557,7 +1453,13 @@ void Model::prefill(const std::vector& tokens, size_t /*chunk_size*/, run_step(tokens[i], cache_total_seq_len_, /*read_logits=*/false); ++cache_total_seq_len_; } + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); last_prefill_scalar_tail_tokens_ = tokens.size() - chunked.logical_tokens; + + if (prepare_decode) { + // After the prompt reaches full length here -- never mid-chunk -- bound it to target_len. + maybe_roll_compact(); + } } void Model::prefill_with_images(const std::vector& tokens, @@ -2135,7 +2037,7 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, decoder_prefill_chunk_->graph->execute(); } if (whole_chunks_end > 0 && decoder_ != nullptr) { - copy_cache_states(*decoder_prefill_chunk_, *decoder_); + move_cache_states(*decoder_prefill_chunk_, *decoder_); decoder_prefill_chunk_->graph->release_runtime_buffers(); unload_component_graph(*decoder_prefill_chunk_); } @@ -2281,6 +2183,7 @@ void Model::prefill_with_media(const std::vector& tokens, run_step(t, pos, false); } cache_total_seq_len_ += tokens.size(); + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); (void)profile_file; } @@ -2292,6 +2195,7 @@ uint32_t Model::decode(const std::vector& tokens, float /*temperature* context_tokens_.insert(context_tokens_.end(), tokens.begin(), tokens.end()); run_full_context_text(); cache_total_seq_len_ = context_tokens_.size(); + cache_token_ids_ = context_tokens_; uint32_t result = argmax_last_logits(out_entropy); record_sampled_token(result); return result; @@ -2301,6 +2205,8 @@ uint32_t Model::decode(const std::vector& tokens, float /*temperature* } run_step(tokens.back(), cache_total_seq_len_ + tokens.size() - 1, /*read_logits=*/true); cache_total_seq_len_ += tokens.size(); + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); + maybe_roll_compact(); uint32_t result = argmax_last_logits(out_entropy); record_sampled_token(result); return result; @@ -2802,6 +2708,8 @@ void Model::reset_cache() { cache_total_seq_len_ = 0; last_logit_position_ = 0; context_tokens_.clear(); + cache_token_ids_.clear(); + special_rows_.clear(); token_history_.clear(); media_features_.clear(); media_feature_shapes_.clear(); @@ -2815,91 +2723,202 @@ void Model::reset_cache() { void Model::set_cache_window(size_t /*window_size*/, size_t /*sink_size*/) {} -void Model::remove_thinking_tokens(const std::vector>& ranges) { - size_t total_removed = 0; - for (const auto& r : ranges) total_removed += r.second; - if (cache_total_seq_len_ >= total_removed) - cache_total_seq_len_ -= total_removed; - else - cache_total_seq_len_ = 0; - - struct CacheHeader { - uint64_t current_seq_len; - uint64_t max_seq_len; - uint64_t num_kv_heads; - uint64_t head_dim; - uint64_t sink_size; - uint64_t reserved[3]; - }; - constexpr size_t kHeaderBytes = 64; - static_assert(sizeof(CacheHeader) == kHeaderBytes, "CacheHeader layout mismatch"); - - auto sorted_ranges = ranges; - std::sort(sorted_ranges.begin(), sorted_ranges.end(), - [](const auto& a, const auto& b) { return a.first < b.first; }); - - for (auto& kv : components_) { - Component& comp = kv.second; - if (!comp.graph) continue; - for (const auto& cs : comp.cache_states) { - for (int node_id : {cs.key_node_id, cs.value_node_id}) { - if (node_id < 0) continue; - const auto& desc = comp.graph->get_output_buffer(static_cast(node_id)); - if (desc.byte_size <= kHeaderBytes || !desc.get_data()) continue; - void* raw = comp.graph->get_output(static_cast(node_id)); - if (!raw) continue; - auto* hdr = static_cast(raw); - size_t cur = hdr->current_seq_len; - if (cur == 0) continue; - size_t kv_heads = hdr->num_kv_heads; - size_t hdim = hdr->head_dim; - if (kv_heads == 0 || hdim == 0) continue; - size_t token_elems = kv_heads * hdim; - size_t num_groups = (hdim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; - size_t token_scales = kv_heads * num_groups; - size_t max_seq = hdr->max_seq_len; - - size_t new_len = cur; - if (desc.precision == Precision::FP16) { - auto* base = reinterpret_cast<__fp16*>(static_cast(raw) + kHeaderBytes); - for (auto it = sorted_ranges.rbegin(); it != sorted_ranges.rend(); ++it) { - size_t start = it->first; - if (start >= new_len) continue; - size_t count = std::min(it->second, new_len - start); - size_t tail_start = start + count; - size_t tail_count = new_len - tail_start; - if (tail_count > 0) { - std::memmove(base + start * token_elems, - base + tail_start * token_elems, - tail_count * token_elems * sizeof(__fp16)); - } - new_len -= count; - } - } else { - auto* int8_base = reinterpret_cast(static_cast(raw) + kHeaderBytes); - auto* scale_base = reinterpret_cast(static_cast(raw) + kHeaderBytes + - max_seq * kv_heads * hdim); - for (auto it = sorted_ranges.rbegin(); it != sorted_ranges.rend(); ++it) { - size_t start = it->first; - if (start >= new_len) continue; - size_t count = std::min(it->second, new_len - start); - size_t tail_start = start + count; - size_t tail_count = new_len - tail_start; - if (tail_count > 0) { - std::memmove(int8_base + start * token_elems, - int8_base + tail_start * token_elems, - tail_count * token_elems); - std::memmove(scale_base + start * token_scales, - scale_base + tail_start * token_scales, - tail_count * token_scales * sizeof(float)); - } - new_len -= count; - } - } - hdr->current_seq_len = new_len; +void Model::apply_kv_compress_env_override() { + config_.parse_kv_compress_override(std::getenv("CACTUS_KV_COMPRESS_AT"), + std::getenv("CACTUS_KV_COMPRESS_TO")); +} + +std::vector Model::compressible_layers() const { + size_t shared = (config_.num_kv_shared_layers == Config::UNSET_U32) + ? 0 : config_.num_kv_shared_layers; + return cactus::kvcompress::physical_compressible_layers( + config_.layer_types, config_.num_layers, shared); +} + +// Keep params.abs_budget survivors per (layer, kv-head) -- sink + recent + most distinctive middle +// tokens by KeyDiff -- compact them, and renumber RoPE positions to a contiguous window 0..B-1. +void Model::compress_kv_cache_keydiff(const cactus::kvcompress::Params& params) { + if (!decoder_) return; + using cactus::kvcompress::CacheHeader; + constexpr size_t kHeaderBytes = sizeof(CacheHeader); + + std::vector layers = compressible_layers(); + if (layers.empty()) return; + std::set compressible(layers.begin(), layers.end()); + const double rope_theta = static_cast(config_.rope_theta); + const double rope_local_theta = (config_.rope_local_base_freq == Config::UNSET_F32) + ? rope_theta : static_cast(config_.rope_local_base_freq); + const size_t old_total = cache_total_seq_len_; + + cactus::kvcompress::Params params_local = params; + bool preserve = config_.kv_compress_preserve_special; + if (const char* e = std::getenv("CACTUS_KV_PRESERVE_SPECIAL")) preserve = (std::atoi(e) != 0); + const bool map_valid = cache_token_ids_.size() == old_total && media_features_.empty(); + const bool per_head_protect = preserve && map_valid && special_rows_.valid(); + // cache_token_ids_ past tracked_len is still head-aligned, so specials there apply to every head. + std::vector appended_special; + if (per_head_protect) { + if (special_ids_.empty() && tokenizer_) special_ids_ = tokenizer_->special_token_ids(); + for (size_t r = special_rows_.tracked_len(); r < old_total && r < cache_token_ids_.size(); ++r) + if (special_ids_.count(cache_token_ids_[r])) appended_special.push_back(static_cast(r)); + } + + Component& comp = *decoder_; + if (!comp.graph) return; + // Skip the whole pass (before mutating) if any layer's V dim differs from K (MLA), or a head + // can't fit sink + all its specials in the budget. + const size_t protect_budget = params.abs_budget > 0 ? static_cast(params.abs_budget) : 0; + for (size_t li = 0; li < comp.cache_states.size(); ++li) { + if (!compressible.count(li)) continue; + const auto& cs = comp.cache_states[li]; + if (cs.key_node_id < 0 || cs.value_node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(cs.key_node_id)) != OpType::KV_CACHE_STATE) continue; + void* kraw = comp.graph->get_output(static_cast(cs.key_node_id)); + void* vraw = comp.graph->get_output(static_cast(cs.value_node_id)); + if (!kraw || !vraw) continue; + if (static_cast(vraw)->head_dim != static_cast(kraw)->head_dim) return; + if (per_head_protect && protect_budget > 0 && + special_rows_.max_reserved(li, params.sink, appended_special) > protect_budget) return; + } + // B is identical across the compacted layers; capture it from the keep-set, not a mutated header. + size_t new_seq_len = 0; + bool have_new_seq_len = false; + std::vector canonical_keep; + bool canonical_captured = false; + // The un-rope table is identical across compressible layers (all global, same theta/dims), so + // build it once on the first such layer and reuse it for every layer's keep-set scoring. + std::vector unrope; + // Reclaim capacity a long prefill grew into; next_pow2(trigger) still fits the decode oscillation. + size_t shrink_cap = 1; + while (shrink_cap < static_cast(config_.kv_compress_trigger_len)) shrink_cap <<= 1; + for (size_t li = 0; li < comp.cache_states.size(); ++li) { + if (!compressible.count(li)) continue; + const auto& cs = comp.cache_states[li]; + if (cs.key_node_id < 0 || cs.value_node_id < 0) continue; + // Skip non-KV caches: hybrid models (LFM2) interleave header-less conv/recurrent states + // here, which KeyDiff would corrupt and read out of bounds. + if (comp.graph->get_node_op_type(static_cast(cs.key_node_id)) != OpType::KV_CACHE_STATE) continue; + + const auto& kdesc = comp.graph->get_output_buffer(static_cast(cs.key_node_id)); + const auto& vdesc = comp.graph->get_output_buffer(static_cast(cs.value_node_id)); + if (kdesc.byte_size <= kHeaderBytes || vdesc.byte_size <= kHeaderBytes) continue; + void* kraw = comp.graph->get_output(static_cast(cs.key_node_id)); + void* vraw = comp.graph->get_output(static_cast(cs.value_node_id)); + if (!kraw || !vraw) continue; + + auto* khdr = static_cast(kraw); + auto* vhdr = static_cast(vraw); + size_t n = khdr->current_seq_len; + size_t kv_heads = khdr->num_kv_heads; + size_t head_dim = khdr->head_dim; + if (kv_heads == 0 || head_dim == 0) continue; + if (unrope.empty()) unrope = cactus::kvcompress::unrope_table(n, head_dim, rope_theta); + + static const std::vector> kNoProtect; + if (per_head_protect) special_rows_.add_appended(li, kv_heads, appended_special); + const std::vector>& pph = per_head_protect ? special_rows_.protect(li) : kNoProtect; + + if (kdesc.precision == Precision::FP16) { + auto* kbase = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + auto* vbase = reinterpret_cast(static_cast(vraw) + kHeaderBytes); + auto kept = cactus::kvcompress::keepsets_from_fp16( + kbase, n, kv_heads, head_dim, unrope, params_local, pph); + if (!canonical_captured) { canonical_keep = kept.empty() ? std::vector{} : kept[0]; canonical_captured = true; } + cactus::kvcompress::compact_fp16(kbase, vbase, kv_heads, head_dim, kept, unrope); + if (per_head_protect) special_rows_.remap(li, kept); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; + vhdr->current_seq_len = B; + new_seq_len = B; + have_new_seq_len = true; + } else if (kdesc.precision == Precision::INT8) { + size_t max_seq = khdr->max_seq_len; + auto* k_i8 = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + auto* k_sc = reinterpret_cast(static_cast(kraw) + kHeaderBytes + + max_seq * kv_heads * head_dim); + auto* v_i8 = reinterpret_cast(static_cast(vraw) + kHeaderBytes); + auto* v_sc = reinterpret_cast(static_cast(vraw) + kHeaderBytes + + max_seq * kv_heads * head_dim); + auto kept = cactus::kvcompress::keepsets_from_int8( + k_i8, k_sc, n, kv_heads, head_dim, KV_QUANT_GROUP_SIZE, unrope, params_local, pph); + if (!canonical_captured) { canonical_keep = kept.empty() ? std::vector{} : kept[0]; canonical_captured = true; } + cactus::kvcompress::compact_int8(k_i8, k_sc, kv_heads, head_dim, KV_QUANT_GROUP_SIZE, + kept, unrope, /*renumber=*/true); + cactus::kvcompress::compact_int8(v_i8, v_sc, kv_heads, head_dim, KV_QUANT_GROUP_SIZE, + kept, unrope, /*renumber=*/false); + if (per_head_protect) special_rows_.remap(li, kept); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; + vhdr->current_seq_len = B; + new_seq_len = B; + have_new_seq_len = true; + } + comp.graph->shrink_cache_buffer(static_cast(cs.key_node_id), shrink_cap); + comp.graph->shrink_cache_buffer(static_cast(cs.value_node_id), shrink_cap); + } + + // Shift non-compacted layers' recent K rows by -Δ into the compressed frame (RoPE is relative, so a + // uniform shift preserves q·k offsets); each uses its own theta -- local for sliding, global for shared sources. + const size_t Delta = (have_new_seq_len && old_total >= new_seq_len) ? old_total - new_seq_len : 0; + if (Delta > 0) { + const double dpos = -static_cast(Delta); + for (size_t li = 0; li < comp.cache_states.size(); ++li) { + if (compressible.count(li)) continue; + const auto& cs = comp.cache_states[li]; + if (cs.key_node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(cs.key_node_id)) != OpType::KV_CACHE_STATE) continue; + const auto& kdesc = comp.graph->get_output_buffer(static_cast(cs.key_node_id)); + if (kdesc.byte_size <= kHeaderBytes) continue; + void* kraw = comp.graph->get_output(static_cast(cs.key_node_id)); + if (!kraw) continue; + auto* khdr = static_cast(kraw); + size_t kv_heads = khdr->num_kv_heads, head_dim = khdr->head_dim; + if (kv_heads == 0 || head_dim == 0) continue; + size_t hi = khdr->current_seq_len; + size_t lo = std::min(khdr->sink_size, hi); + const double layer_theta = cactus::kvcompress::is_sliding_layer(config_.layer_types, li) + ? rope_local_theta : rope_theta; + if (kdesc.precision == Precision::FP16) { + auto* kbase = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + cactus::kvcompress::rerope_recent_fp16(kbase, kv_heads, head_dim, lo, hi, + layer_theta, dpos); + } else if (kdesc.precision == Precision::INT8) { + size_t max_seq = khdr->max_seq_len; + auto* k_i8 = reinterpret_cast(static_cast(kraw) + kHeaderBytes); + auto* k_sc = reinterpret_cast(static_cast(kraw) + kHeaderBytes + + max_seq * kv_heads * head_dim); + cactus::kvcompress::rerope_recent_int8(k_i8, k_sc, kv_heads, head_dim, + KV_QUANT_GROUP_SIZE, lo, hi, layer_theta, dpos); } } } + + // All layers now share the same compressed frame, so the next decode query uses position_ids = B. + if (have_new_seq_len) { + cache_total_seq_len_ = new_seq_len; + if (per_head_protect) special_rows_.set_tracked_len(new_seq_len); + else special_rows_.invalidate(); + if (map_valid && canonical_captured) { + std::vector compacted; + compacted.reserve(canonical_keep.size()); + for (int idx : canonical_keep) + if (idx >= 0 && idx < static_cast(cache_token_ids_.size())) compacted.push_back(cache_token_ids_[idx]); + cache_token_ids_ = std::move(compacted); + } else { + cache_token_ids_.clear(); + } + } +} + +void Model::maybe_roll_compact() { + if (!config_.kv_compress || config_.kv_compress_trigger_len <= 0) return; + if (cache_total_seq_len_ < static_cast(config_.kv_compress_trigger_len)) return; + + cactus::kvcompress::Params p; + p.recent_frac = config_.kv_compress_recent_frac; + p.sink = config_.kv_compress_sink; + p.abs_budget = config_.kv_compress_target_len; + compress_kv_cache_keydiff(p); } std::vector Model::get_embeddings(const std::vector& tokens, bool pooled, @@ -3088,6 +3107,12 @@ bool Config::from_json(const std::string& config_path) { else model_variant = ModelVariant::DEFAULT; } else if (key == "conv_L_cache") conv_L_cache = static_cast(std::stoul(value)); + else if (key == "kv_compress") kv_compress = (value == "true" || value == "1"); + else if (key == "kv_compress_recent_frac") kv_compress_recent_frac = std::stof(value); + else if (key == "kv_compress_sink") kv_compress_sink = static_cast(std::stoul(value)); + else if (key == "kv_compress_trigger_len") kv_compress_trigger_len = static_cast(std::stol(value)); + else if (key == "kv_compress_target_len") kv_compress_target_len = static_cast(std::stol(value)); + else if (key == "kv_compress_preserve_special") kv_compress_preserve_special = (value == "true" || value == "1"); else if (key == "layer_types") { layer_types.clear(); std::string sanitized; @@ -3200,8 +3225,6 @@ bool Config::from_json(const std::string& config_path) { audio_fft_length = audio_fft_overdrive ? 1024u : 512u; } else if (key == "audio_token_id") audio_token_id = static_cast(std::stoul(value)); - else if (key == "channel_open_token_id") channel_open_token_id = static_cast(std::stoul(value)); - else if (key == "channel_close_token_id") channel_close_token_id = static_cast(std::stoul(value)); else if (key == "activation_sparsity_ppf") { activation_sparsity_ppf.clear(); std::stringstream ss(value); @@ -3256,9 +3279,42 @@ bool Config::from_json(const std::string& config_path) { } } + validate_kv_compress(); + return true; +} + +bool Config::parse_kv_compress_override(const char* trigger_env, const char* target_env) { + const bool has_trigger = trigger_env && *trigger_env; + const bool has_target = target_env && *target_env; + if (!has_trigger && !has_target) return false; + if (has_trigger) kv_compress_trigger_len = static_cast(std::stol(trigger_env)); + if (has_target) kv_compress_target_len = static_cast(std::stol(target_env)); + if (kv_compress_trigger_len <= 0) { + kv_compress = false; + kv_compress_trigger_len = 0; + kv_compress_target_len = 0; + CACTUS_LOG_INFO("kv_compress", "rolling compaction disabled (CACTUS_KV_COMPRESS_AT <= 0)"); + return true; + } + kv_compress = true; + validate_kv_compress(); + CACTUS_LOG_INFO("kv_compress", "rolling override: trigger_len=" << kv_compress_trigger_len + << " target_len=" << kv_compress_target_len); return true; } +void Config::validate_kv_compress() { + if (kv_compress_trigger_len <= 0) return; + if (kv_compress_target_len <= 0 || kv_compress_target_len >= kv_compress_trigger_len) { + CACTUS_LOG_WARN("kv_compress", "invalid rolling config (target_len=" << kv_compress_target_len + << ", trigger_len=" << kv_compress_trigger_len + << "): require 0 < target_len < trigger_len; disabling rolling compaction"); + kv_compress = false; + kv_compress_trigger_len = 0; + kv_compress_target_len = 0; + } +} + std::string Config::to_json() const { return "{}"; } diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index 8f5930d8e..14f1336d0 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -426,6 +426,10 @@ std::string Tokenizer::format_qwen_style(const std::vector& message if (is_lfm2) result += "\n"; } } + // Replicating generation prompt so the kv cache can and is re-used across turns. + if (role == "assistant" && chat_template_.find("") != std::string::npos) { + result += msg.content.find("") != std::string::npos ? "\n" : "\n\n\n\n"; + } result += msg.content; if (role == "assistant" && !msg.tool_calls.empty()) { for (const auto& tc : msg.tool_calls) { @@ -471,27 +475,6 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa result += "\n"; } - auto strip_channel = [](const std::string& text) -> std::string { - const std::string open_tag = "<|channel>"; - const std::string close_tag = ""; - std::string out; - size_t pos = 0; - while (pos < text.size()) { - size_t open_pos = text.find(open_tag, pos); - if (open_pos == std::string::npos) { - out += text.substr(pos); - break; - } - out += text.substr(pos, open_pos - pos); - size_t close_pos = text.find(close_tag, open_pos + open_tag.size()); - if (close_pos == std::string::npos) { - break; - } - pos = close_pos + close_tag.size(); - } - return out; - }; - auto compute_soft_tokens = [&](const std::string& image_path) -> size_t { int w = 0, h = 0, c = 0; unsigned char* data = cactus_image_load(image_path.c_str(), &w, &h, &c, 3); @@ -516,7 +499,7 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa std::string role = (msg.role == "assistant") ? "model" : msg.role; result += "<|turn>" + role + "\n"; if (role == "model") { - result += strip_channel(msg.content); + result += msg.content; if (!msg.tool_calls.empty()) { for (const auto& tc : msg.tool_calls) { result += format_tool_call_for_prompt(tc.name, tc.arguments, true); diff --git a/cactus-engine/src/utils.h b/cactus-engine/src/utils.h index 9f66f4895..ca64eb549 100644 --- a/cactus-engine/src/utils.h +++ b/cactus-engine/src/utils.h @@ -1534,30 +1534,6 @@ inline void parse_function_calls_from_response(const std::string& response_text, } } -inline std::vector> find_channel_token_ranges( - const std::vector& tokens, size_t offset, - uint32_t channel_open_id, uint32_t channel_close_id) { - std::vector> ranges; - size_t pos = 0; - while (pos < tokens.size()) { - if (tokens[pos] != channel_open_id) { - pos++; - continue; - } - - size_t block_start = pos; - pos++; - while (pos < tokens.size() && tokens[pos] != channel_close_id) { - pos++; - } - if (pos < tokens.size()) { - pos++; - } - ranges.push_back({offset + block_start, pos - block_start}); - } - return ranges; -} - inline void strip_tag_blocks(std::string& text, std::string& extracted, const std::string& open_tag, const std::string& close_tag) { std::string result; @@ -1592,7 +1568,7 @@ inline void strip_tag_blocks(std::string& text, std::string& extracted, text = result; } -inline void strip_thinking_block(const std::string& input, std::string& thinking, std::string& content) { +inline void partition_thinking_response(const std::string& input, std::string& thinking, std::string& content) { thinking.clear(); content = input; @@ -1634,13 +1610,17 @@ inline std::string construct_response_json(const std::string& regular_response, float confidence = 0.0f, bool cloud_handoff = false, const std::string& thinking = "", - const std::vector& segments = {}) { + const std::vector& segments = {}, + const std::string& context_response = "") { std::ostringstream json; json << "{"; json << "\"success\":true,"; json << "\"error\":null,"; json << "\"cloud_handoff\":" << (cloud_handoff ? "true" : "false") << ","; json << "\"response\":\"" << escape_json_string(regular_response) << "\","; + if (!context_response.empty()) { + json << "\"context_response\":\"" << escape_json_string(context_response) << "\","; + } if (!thinking.empty()) { json << "\"thinking\":\"" << escape_json_string(thinking) << "\","; } diff --git a/cactus-engine/tests/run.cpp b/cactus-engine/tests/run.cpp index 7fcdcb0fb..8f63a05d0 100644 --- a/cactus-engine/tests/run.cpp +++ b/cactus-engine/tests/run.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -301,7 +302,11 @@ bool file_exists(const std::string& path) { std::vector parse_token_ids(const std::string& text) { std::vector tokens; - std::stringstream ss(text); + std::string normalized = text; + for (char& ch : normalized) { + if (ch == ' ' || ch == '\t' || ch == '\r' || ch == '\n') ch = ','; + } + std::stringstream ss(normalized); std::string part; while (std::getline(ss, part, ',')) { size_t start = part.find_first_not_of(" \t\r\n"); @@ -318,6 +323,16 @@ std::vector parse_token_ids(const std::string& text) { return tokens; } +bool read_text_file(const std::string& path, std::string& text) { + std::ifstream in(path); + if (!in) return false; + std::ostringstream buffer; + buffer << in.rdbuf(); + if (!in.good() && !in.eof()) return false; + text = buffer.str(); + return true; +} + bool write_text_file(const std::string& path, const std::string& text) { std::ofstream out(path); if (!out) return false; @@ -395,7 +410,7 @@ std::string build_messages(const std::string& system_prompt, void print_usage(const char* argv0) { std::cerr << "Usage: " << argv0 << " [--system ] [--image ] [--audio ]" - << " [--prompt ] [--input-ids ] [--max-new-tokens ]" + << " [--prompt ] [--input-ids ] [--input-ids-file ] [--max-new-tokens ]" << " [--result-json ] [--thinking] [--no-cloud-handoff]" << " [--confidence-threshold ] [--cloud-timeout-ms ]\n"; } @@ -417,6 +432,7 @@ int main(int argc, char** argv) { std::string current_audio; std::string initial_prompt; std::string input_ids; + std::string input_ids_file; std::string result_json; int max_new_tokens = kMaxTokens; bool thinking = false; @@ -436,6 +452,8 @@ int main(int argc, char** argv) { initial_prompt = argv[++i]; } else if (arg == "--input-ids" && i + 1 < argc) { input_ids = argv[++i]; + } else if (arg == "--input-ids-file" && i + 1 < argc) { + input_ids_file = expand_tilde(argv[++i]); } else if (arg == "--max-new-tokens" && i + 1 < argc) { max_new_tokens = std::max(0, std::atoi(argv[++i])); } else if (arg == "--result-json" && i + 1 < argc) { @@ -459,6 +477,20 @@ int main(int argc, char** argv) { std::cerr << "Audio file not found: " << current_audio << "\n"; return 1; } + if (!input_ids.empty() && !input_ids_file.empty()) { + std::cerr << "Use either --input-ids or --input-ids-file, not both\n"; + return 1; + } + if (!input_ids_file.empty()) { + if (!file_exists(input_ids_file)) { + std::cerr << "Input ids file not found: " << input_ids_file << "\n"; + return 1; + } + if (!read_text_file(input_ids_file, input_ids)) { + std::cerr << "Failed to read input ids file: " << input_ids_file << "\n"; + return 1; + } + } std::cout << "Loading model from " << model_path << "...\n"; cactus_model_t model = cactus_init(model_path.c_str(), nullptr, false); @@ -517,7 +549,9 @@ int main(int argc, char** argv) { std::string input; if (auto_send) { auto_send = false; - input = initial_prompt.empty() ? "Describe the attached input." : initial_prompt; + if (!initial_prompt.empty()) input = initial_prompt; + else if (!current_image.empty()) input = "Describe this image."; + else input = "Transcribe or respond to this audio."; std::cout << "You: " << input << "\n"; } else { std::cout << "You: " << std::flush; @@ -646,7 +680,8 @@ int main(int argc, char** argv) { continue; } - std::string assistant = json_string_value(response_json, "response"); + std::string assistant = json_string_value(response_json, "context_response"); + if (assistant.empty()) assistant = json_string_value(response_json, "response"); history.push_back({"assistant", assistant}); current_image.clear(); current_audio.clear(); diff --git a/cactus-engine/tests/test_kv_compress.cpp b/cactus-engine/tests/test_kv_compress.cpp new file mode 100644 index 000000000..20f46afa1 --- /dev/null +++ b/cactus-engine/tests/test_kv_compress.cpp @@ -0,0 +1,1007 @@ +#include "test_utils.h" +#include "../src/kv_compress.h" +#include "../src/engine.h" +#include "cactus_graph.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace TestUtils; +using namespace cactus::kvcompress; + +namespace { + +using Header = cactus::kvcompress::CacheHeader; +constexpr size_t kHeaderBytes = sizeof(Header); +constexpr size_t kGroupSize = 32; // KV_QUANT_GROUP_SIZE + +float f16_to_f32(uint16_t u) { __fp16 v; std::memcpy(&v, &u, 2); return static_cast(v); } +uint16_t f32_to_f16(float f) { __fp16 v = static_cast<__fp16>(f); uint16_t u; std::memcpy(&u, &v, 2); return u; } + +using EngineTestUtils::rope_reference; + +// `pre_rope[h][t][d]` is the planted pre-RoPE key, stored post-RoPE at position t. +std::vector make_fp16_cache(size_t n, size_t max_seq, size_t kv_heads, size_t head_dim, + double theta, std::vector>>& pre_rope, + std::vector>>& values) { + std::vector buf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + auto* h = reinterpret_cast(buf.data()); + h->current_seq_len = n; h->max_seq_len = max_seq; h->num_kv_heads = kv_heads; + h->head_dim = head_dim; h->sink_size = 4; + auto* rows = reinterpret_cast(buf.data() + kHeaderBytes); + pre_rope.assign(kv_heads, std::vector>(n)); + values.assign(kv_heads, std::vector>(n)); + for (size_t hh = 0; hh < kv_heads; ++hh) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim), v(head_dim); + for (size_t d = 0; d < head_dim; ++d) { + k[d] = std::sin(0.1f * (hh + 1) * (t + 1) + 0.3f * d) + 0.05f * d; + v[d] = std::cos(0.07f * (hh + 1) * (t + 2) + 0.2f * d); + } + pre_rope[hh][t] = k; values[hh][t] = v; + std::vector kr = rope_reference(k, (double)t, theta); + for (size_t d = 0; d < head_dim; ++d) + rows[(t * kv_heads + hh) * head_dim + d] = f32_to_f16(kr[d]); + } + return buf; +} + +std::vector make_fp16_value_cache(size_t n, size_t max_seq, size_t kv_heads, size_t head_dim, + const std::vector>>& values) { + std::vector buf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + auto* h = reinterpret_cast(buf.data()); + h->current_seq_len = n; h->max_seq_len = max_seq; h->num_kv_heads = kv_heads; + h->head_dim = head_dim; h->sink_size = 4; + auto* rows = reinterpret_cast(buf.data() + kHeaderBytes); + for (size_t hh = 0; hh < kv_heads; ++hh) + for (size_t t = 0; t < n; ++t) + for (size_t d = 0; d < head_dim; ++d) + rows[(t * kv_heads + hh) * head_dim + d] = f32_to_f16(values[hh][t][d]); + return buf; +} + +// Independent oracle quantizer, kept separate from production requant so tests don't self-validate. +void quantize_row(const std::vector& row, int8_t* dst, float* sc, + size_t head_dim, size_t group_size) { + size_t groups = (head_dim + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + size_t lo = g * group_size, hi = std::min(head_dim, lo + group_size); + float amax = 0.f; for (size_t d = lo; d < hi; ++d) amax = std::max(amax, std::fabs(row[d])); + float scale = amax / 127.f; if (scale < 1e-10f) scale = 1e-10f; sc[g] = scale; + float inv = 1.f / scale; + for (size_t d = lo; d < hi; ++d) { int32_t q = (int32_t)std::roundf(row[d] * inv); q = std::max(-128, std::min(127, q)); dst[d] = (int8_t)q; } + } +} + +} // namespace + +bool test_compact_fp16_cache() { + const size_t n = 60, max_seq = 128, kv_heads = 3, head_dim = 16; + const double theta = 1000000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + std::vector> kept = { + {0, 5, 10, 20, 30, 59}, {1, 4, 12, 25, 40, 58}, {0, 2, 15, 33, 50, 55}}; + size_t B = kept[0].size(); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + + bool ok = true; + for (size_t h = 0; h < kv_heads; ++h) { + for (size_t rank = 0; rank < B; ++rank) { + int abs_pos = kept[h][rank]; + std::vector expK = rope_reference(pre[h][abs_pos], (double)rank, theta); + const uint16_t* kdst = krows + (rank * kv_heads + h) * head_dim; + const uint16_t* vdst = vrows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + if (std::abs(f16_to_f32(kdst[d]) - expK[d]) > 5e-2f) ok = false; + if (std::abs(f16_to_f32(vdst[d]) - vals[h][abs_pos][d]) > 5e-2f) ok = false; + } + } + } + return ok && B == 6; +} + +bool test_dense_check_full_budget() { + // abs_budget == n keeps every index in order, so renumber delta is 0 and the buffer is unchanged. + const size_t n = 40, max_seq = 64, kv_heads = 2, head_dim = 16; + const double theta = 1000000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + std::vector kbuf0 = kbuf, vbuf0 = vbuf; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + Params p; p.recent_frac = 0.3f; p.sink = 4; p.abs_budget = (int)n; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, theta, p); + for (auto& k : kept) if (k.size() != n) return false; + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + + size_t live = n * kv_heads * head_dim * sizeof(uint16_t); + if (std::memcmp(kbuf.data() + kHeaderBytes, kbuf0.data() + kHeaderBytes, live) != 0) return false; + if (std::memcmp(vbuf.data() + kHeaderBytes, vbuf0.data() + kHeaderBytes, live) != 0) return false; + return true; +} + +bool test_rope_renumber_contiguous() { + // rope(stored_post_rope_at_abs, rank - abs) == rope(pre, rank). + const size_t head_dim = 16; const double theta = 1000000.0; + std::vector pre(head_dim); + for (size_t d = 0; d < head_dim; ++d) pre[d] = std::sin(0.4 * d) + 0.1 * d; + int abs_pos = 137; size_t rank = 5; + std::vector at_abs = rope_reference(pre, (double)abs_pos, theta); + std::vector at_rank = rope_reference(pre, (double)rank, theta); + std::vector delta = at_abs; + rope_rotate_row(delta.data(), head_dim, theta, (double)rank - (double)abs_pos); + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(delta[d] - at_rank[d]) > 1e-3f) return false; + return true; +} + +bool test_fp16_storage_round_trip() { + // Each survivor must read back as rope(pre, rank) within fp16 tolerance. + const size_t n = 64, max_seq = 128, kv_heads = 2, head_dim = 16; + const double theta = 1000000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + Params p; p.recent_frac = 0.3f; p.sink = 4; p.abs_budget = 16; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, theta, p); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + + bool ok = true; + for (size_t h = 0; h < kv_heads; ++h) { + for (size_t rank = 0; rank < kept[h].size(); ++rank) { + int abs_pos = kept[h][rank]; + std::vector expK = rope_reference(pre[h][abs_pos], (double)rank, theta); + const uint16_t* kdst = krows + (rank * kv_heads + h) * head_dim; + const uint16_t* vdst = vrows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + if (std::abs(f16_to_f32(kdst[d]) - expK[d]) > 5e-2f) ok = false; + if (std::abs(f16_to_f32(vdst[d]) - vals[h][abs_pos][d]) > 5e-2f) ok = false; + } + } + } + return ok; +} + +bool test_gemma_layer_selection() { + std::vector qwen(28, "full_attention"); + auto q = physical_compressible_layers(qwen, 28, 0); + if (q.size() != 28) return false; + for (size_t i = 0; i < 28; ++i) if (q[i] != i) return false; + + std::vector g; + for (int i = 0; i < 35; ++i) g.push_back(((i + 1) % 5 == 0) ? "global" : "sliding"); + auto gg = physical_compressible_layers(g, 35, 20); + std::vector expect = {4, 9}; + if (gg != expect) return false; + + // Pass-2 theta follows is_sliding_layer, NOT compressibility: layer 14 is a global KV-shared + // *source* excluded from compaction (not in {4,9}) yet must still re-rope with the global theta. + if (is_sliding_layer(g, 14)) return false; + if (is_sliding_layer(g, 4)) return false; + if (!is_sliding_layer(g, 3)) return false; + for (size_t i = 0; i < 28; ++i) if (is_sliding_layer(qwen, i)) return false; + return true; +} + +bool test_compact_int8_cache() { + const size_t n = 50, max_seq = 96, kv_heads = 2, head_dim = 32; + const double theta = 1000000.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + + std::vector>> pre(kv_heads, std::vector>(n)); + std::vector>> vals(kv_heads, std::vector>(n)); + std::vector k_i8(max_seq * int8_stride, 0), v_i8(max_seq * int8_stride, 0); + std::vector k_sc(max_seq * scale_stride, 0.f), v_sc(max_seq * scale_stride, 0.f); + + auto quant_row = [&](const std::vector& row, int8_t* dst, float* sc) { + quantize_row(row, dst, sc, head_dim, kGroupSize); + }; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim), v(head_dim); + for (size_t d = 0; d < head_dim; ++d) { k[d] = std::sin(0.1 * (h + 1) * (t + 1) + 0.2 * d); v[d] = std::cos(0.05 * (t + 1) + 0.1 * d); } + pre[h][t] = k; vals[h][t] = v; + std::vector kr = rope_reference(k, (double)t, theta); + quant_row(kr, k_i8.data() + t * int8_stride + h * head_dim, k_sc.data() + t * scale_stride + h * groups); + quant_row(v, v_i8.data() + t * int8_stride + h * head_dim, v_sc.data() + t * scale_stride + h * groups); + } + + std::vector> kept = {{0, 5, 11, 22, 33, 49}, {1, 4, 13, 26, 40, 48}}; + size_t B = kept[0].size(); + compact_int8(k_i8.data(), k_sc.data(), kv_heads, head_dim, kGroupSize, kept, theta, true); + compact_int8(v_i8.data(), v_sc.data(), kv_heads, head_dim, kGroupSize, kept, theta, false); + + bool ok = true; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t rank = 0; rank < B; ++rank) { + int abs_pos = kept[h][rank]; + std::vector expK = rope_reference(pre[h][abs_pos], (double)rank, theta); + const int8_t* kd = k_i8.data() + rank * int8_stride + h * head_dim; + const float* ks = k_sc.data() + rank * scale_stride + h * groups; + const int8_t* vd = v_i8.data() + rank * int8_stride + h * head_dim; + const float* vs = v_sc.data() + rank * scale_stride + h * groups; + for (size_t d = 0; d < head_dim; ++d) { + float dq = (float)kd[d] * ks[d / kGroupSize]; + if (std::abs(dq - expK[d]) > 0.05f) ok = false; + float dqv = (float)vd[d] * vs[d / kGroupSize]; + if (std::abs(dqv - vals[h][abs_pos][d]) > 0.05f) ok = false; + } + } + return ok; +} + +// Mirrors Model::maybe_roll_compact on an in-place FP16 K/V pair; returns the new bounded length B. +size_t roll_compact_once(uint16_t* krows, uint16_t* vrows, Header* khdr, Header* vhdr, + size_t kv_heads, size_t head_dim, double theta, int target_len) { + size_t n = khdr->current_seq_len; + Params p; p.recent_frac = 0.30f; p.sink = 4; + p.abs_budget = target_len; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, theta, p); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; + vhdr->current_seq_len = B; + return B; +} + +bool test_rolling_bounded_compaction() { + // Grow to trigger_len -> compact to target_len -> repeat, asserting the cache stays bounded and + // a planted distinctive mid-token survives every compaction. + const int trigger_len = 4096, target_len = 2048; + const size_t kv_heads = 2, head_dim = 16, max_seq = trigger_len + 8; + const double theta = 1000000.0; + + std::vector kbuf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + std::vector vbuf(kbuf.size(), 0); + auto* khdr = reinterpret_cast(kbuf.data()); + auto* vhdr = reinterpret_cast(vbuf.data()); + *khdr = Header{0, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + *vhdr = *khdr; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + // A centroid-opposite direction that KeyDiff must keep. + std::vector distinctive(head_dim); + for (size_t d = 0; d < head_dim; ++d) distinctive[d] = (d % 2 == 0 ? -1.0f : 1.0f) * 4.0f; + + // The token at `plant_abs` is the distinctive one; the rest cluster on a common centroid. + auto append = [&](size_t from, size_t to, long plant_abs) { + for (size_t t = from; t < to; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector k(head_dim); + if ((long)t == plant_abs) { + k = distinctive; + } else { + for (size_t d = 0; d < head_dim; ++d) + k[d] = std::sin(0.05 * (h + 1) + 0.3 * d) + 0.01f * std::sin(0.01 * t); + } + std::vector kr = rope_reference(k, (double)t, theta); + std::vector vr(head_dim); + for (size_t d = 0; d < head_dim; ++d) vr[d] = std::cos(0.07 * (t + 1) + 0.2 * d); + for (size_t d = 0; d < head_dim; ++d) { + krows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(kr[d]); + vrows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(vr[d]); + } + } + }; + + auto distinctive_survives = [&](size_t B) -> bool { + for (size_t h = 0; h < kv_heads; ++h) + for (size_t rank = 0; rank < B; ++rank) { + std::vector pre(head_dim); + const uint16_t* src = krows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) pre[d] = f16_to_f32(src[d]); + rope_rotate_row(pre.data(), head_dim, theta, -(double)rank); + double dot = 0, na = 0, nb = 0; + for (size_t d = 0; d < head_dim; ++d) { + dot += pre[d] * distinctive[d]; na += pre[d] * pre[d]; nb += distinctive[d] * distinctive[d]; + } + if (dot / (std::sqrt(na) * std::sqrt(nb) + 1e-9) > 0.99) return true; + } + return false; + }; + + bool ok = true; + const int cycles = 3; + for (int c = 0; c < cycles; ++c) { + size_t start = khdr->current_seq_len; + long plant_abs = (c == 0) ? (trigger_len / 2) : -1; + append(start, (size_t)trigger_len, plant_abs); + khdr->current_seq_len = trigger_len; + vhdr->current_seq_len = trigger_len; + + if (c == 0 && !distinctive_survives(trigger_len)) ok = false; + + size_t B = roll_compact_once(krows, vrows, khdr, vhdr, kv_heads, head_dim, theta, target_len); + if (B != (size_t)target_len) ok = false; + if (khdr->current_seq_len > (size_t)trigger_len) ok = false; + if (khdr->current_seq_len != (size_t)target_len) ok = false; + if (!distinctive_survives(B)) ok = false; + } + if (khdr->current_seq_len != (size_t)target_len) ok = false; + return ok; +} + +bool test_config_parse_rolling_fields() { + cactus::engine::Config def; + if (!def.kv_compress) return false; + if (def.kv_compress_trigger_len != 4096 || def.kv_compress_target_len != 2048) return false; + + char tmpl[] = "/tmp/cactus_kvcfg_XXXXXX"; + int fd = mkstemp(tmpl); + if (fd < 0) return false; + ::close(fd); + { + std::ofstream f(tmpl); + // model_type=qwen avoids the Gemma4-only required-field validation. + f << "model_type=qwen\n" + << "kv_compress=true\n" + << "kv_compress_trigger_len=4096\n" + << "kv_compress_target_len=2048\n"; + } + cactus::engine::Config cfg; + bool parsed = cfg.from_json(tmpl); + std::remove(tmpl); + if (!parsed) return false; + return cfg.kv_compress && cfg.kv_compress_trigger_len == 4096 && cfg.kv_compress_target_len == 2048; +} + +bool test_trigger_zero_gates_rolling() { + // The maybe_roll_compact gate fires iff kv_compress && trigger_len > 0 && seq_len >= trigger_len. + const size_t n = 5000, kv_heads = 2, head_dim = 16, max_seq = 5008; + const double theta = 1000000.0; + + auto fresh_cache = [&]() { + std::vector buf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + *reinterpret_cast(buf.data()) = Header{n, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + auto* rows = reinterpret_cast(buf.data() + kHeaderBytes); + for (size_t i = 0; i < n * kv_heads * head_dim; ++i) rows[i] = f32_to_f16(0.01f * (i % 97)); + return buf; + }; + auto roll_if_gated = [&](cactus::engine::Config& cfg, std::vector& kbuf, std::vector& vbuf) { + auto* khdr = reinterpret_cast(kbuf.data()); + bool fire = cfg.kv_compress && cfg.kv_compress_trigger_len > 0 && + khdr->current_seq_len >= (size_t)cfg.kv_compress_trigger_len; + if (fire) + roll_compact_once(reinterpret_cast(kbuf.data() + kHeaderBytes), + reinterpret_cast(vbuf.data() + kHeaderBytes), khdr, + reinterpret_cast(vbuf.data()), kv_heads, head_dim, theta, + cfg.kv_compress_target_len); + }; + + // Gate closed: cache byte-identical, length unchanged. + cactus::engine::Config off; off.kv_compress_trigger_len = 0; off.kv_compress_target_len = 0; + std::vector k_off = fresh_cache(), v_off(k_off.size(), 0), k_off0 = k_off, v_off0 = v_off; + roll_if_gated(off, k_off, v_off); + if (reinterpret_cast(k_off.data())->current_seq_len != n) return false; + if (std::memcmp(k_off.data(), k_off0.data(), k_off.size()) != 0) return false; + if (std::memcmp(v_off.data(), v_off0.data(), v_off.size()) != 0) return false; + + // Gate open: compacts to target_len and rewrites bytes. + cactus::engine::Config on; on.kv_compress_trigger_len = 4096; on.kv_compress_target_len = 2048; + std::vector k_on = fresh_cache(), v_on(k_on.size(), 0), k_on0 = k_on; + roll_if_gated(on, k_on, v_on); + if (reinterpret_cast(k_on.data())->current_seq_len != 2048) return false; + if (std::memcmp(k_on.data(), k_on0.data(), k_on.size()) == 0) return false; + return true; +} + +bool test_degenerate_rolling_config_disabled() { + // validate_kv_compress requires 0 < target_len < trigger_len; bad configs reset both to 0. + auto disabled = [](int trig, int targ) { + cactus::engine::Config c; + c.kv_compress = true; + c.kv_compress_trigger_len = trig; + c.kv_compress_target_len = targ; + c.validate_kv_compress(); + return c.kv_compress_trigger_len == 0 && c.kv_compress_target_len == 0; + }; + if (!disabled(4096, 0)) return false; + if (!disabled(4096, 4096)) return false; + if (!disabled(2048, 4096)) return false; + if (!disabled(4096, -1)) return false; + cactus::engine::Config ok; + ok.kv_compress = true; + ok.kv_compress_trigger_len = 4096; + ok.kv_compress_target_len = 2048; + ok.validate_kv_compress(); + return ok.kv_compress_trigger_len == 4096 && ok.kv_compress_target_len == 2048; +} + +bool test_env_override_parse() { + cactus::engine::Config off; + if (off.parse_kv_compress_override(nullptr, nullptr)) return false; + if (!off.kv_compress || off.kv_compress_trigger_len != 4096 || off.kv_compress_target_len != 2048) return false; + + cactus::engine::Config both; + if (!both.parse_kv_compress_override("3000", "1000")) return false; + if (!both.kv_compress || both.kv_compress_trigger_len != 3000 || both.kv_compress_target_len != 1000) return false; + + // Setting only one keeps the other's default. + cactus::engine::Config one; + if (!one.parse_kv_compress_override(nullptr, "1500")) return false; + if (one.kv_compress_trigger_len != 4096 || one.kv_compress_target_len != 1500) return false; + + cactus::engine::Config disabled; + if (!disabled.parse_kv_compress_override("0", nullptr)) return false; + if (disabled.kv_compress) return false; + + cactus::engine::Config bad; + if (!bad.parse_kv_compress_override("2048", "4096")) return false; + if (bad.kv_compress_trigger_len != 0 || bad.kv_compress_target_len != 0) return false; + return true; +} + +bool test_rerope_recent_fp16_uniform() { + // After a delta shift each recent row reads as rope(pre, t+delta); sink rows stay byte-identical. + const size_t n = 40, max_seq = 64, kv_heads = 2, head_dim = 16, sink = 4; + const double theta = 10000.0, delta = -7.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + std::vector sink_before(krows, krows + sink * kv_heads * head_dim); + + rerope_recent_fp16(krows, kv_heads, head_dim, sink, n, theta, delta); + + if (std::memcmp(krows, sink_before.data(), sink_before.size() * sizeof(uint16_t)) != 0) return false; + for (size_t t = sink; t < n; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector exp = rope_reference(pre[h][t], (double)t + delta, theta); + const uint16_t* r = krows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(f16_to_f32(r[d]) - exp[d]) > 5e-2f) return false; + } + return true; +} + +namespace { +void build_int8_key_cache(size_t n, size_t kv_heads, size_t head_dim, double theta, + std::vector>>& pre, + std::vector& k_i8, std::vector& k_sc) { + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + pre.assign(kv_heads, std::vector>(n)); + k_i8.assign(n * int8_stride, 0); k_sc.assign(n * scale_stride, 0.f); + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim); + for (size_t d = 0; d < head_dim; ++d) k[d] = std::sin(0.1 * (h + 1) * (t + 1) + 0.2 * d); + pre[h][t] = k; + std::vector kr = rope_reference(k, (double)t, theta); + quantize_row(kr, k_i8.data() + t * int8_stride + h * head_dim, + k_sc.data() + t * scale_stride + h * groups, head_dim, kGroupSize); + } +} +} // namespace + +bool test_rerope_recent_int8() { + const size_t n = 40, kv_heads = 2, head_dim = 32, sink = 4; + const double theta = 10000.0, delta = -9.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + std::vector>> pre; + std::vector k_i8; std::vector k_sc; + build_int8_key_cache(n, kv_heads, head_dim, theta, pre, k_i8, k_sc); + + std::vector sink_i8(k_i8.begin(), k_i8.begin() + sink * int8_stride); + std::vector sink_sc(k_sc.begin(), k_sc.begin() + sink * scale_stride); + rerope_recent_int8(k_i8.data(), k_sc.data(), kv_heads, head_dim, kGroupSize, sink, n, theta, delta); + + if (std::memcmp(k_i8.data(), sink_i8.data(), sink_i8.size()) != 0) return false; + if (std::memcmp(k_sc.data(), sink_sc.data(), sink_sc.size() * sizeof(float)) != 0) return false; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = sink; t < n; ++t) { + std::vector exp = rope_reference(pre[h][t], (double)t + delta, theta); + const int8_t* kd = k_i8.data() + t * int8_stride + h * head_dim; + const float* ks = k_sc.data() + t * scale_stride + h * groups; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs((float)kd[d] * ks[d / kGroupSize] - exp[d]) > 0.05f) return false; + } + return true; +} + +bool test_rotate_int8_row_matches_inline() { + // rotate_int8_row must be bitwise-identical to the inline dequant -> rotate -> requant sequence. + const size_t head_dim = 32; const double theta = 10000.0, delta = -9.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + std::vector pre(head_dim); + for (size_t d = 0; d < head_dim; ++d) pre[d] = std::sin(0.3 * d) + 0.1 * d; + std::vector stored = rope_reference(pre, 21.0, theta); + std::vector a_i8(head_dim), b_i8(head_dim); + std::vector a_sc(groups), b_sc(groups); + quantize_row(stored, a_i8.data(), a_sc.data(), head_dim, kGroupSize); + b_i8 = a_i8; b_sc = a_sc; + + std::vector row(head_dim); + for (size_t d = 0; d < head_dim; ++d) row[d] = (float)a_i8[d] * a_sc[d / kGroupSize]; + rope_rotate_row(row.data(), head_dim, theta, delta); + quantize_row(row, a_i8.data(), a_sc.data(), head_dim, kGroupSize); + rotate_int8_row(b_i8.data(), b_sc.data(), head_dim, kGroupSize, theta, delta); + + return std::memcmp(a_i8.data(), b_i8.data(), head_dim) == 0 + && std::memcmp(a_sc.data(), b_sc.data(), groups * sizeof(float)) == 0; +} + +bool test_sliding_plus_global_rolling() { + // A global cache renumbers survivors to 0..B-1 while a sliding cache re-ropes its recent rows by + // -Δ (Δ = old_total - B), keeping both in the same compressed frame across 3 cycles. + const size_t kv_heads = 2, head_dim = 16, sink = 4; + const double gtheta = 1000000.0, ltheta = 10000.0; + + { + const size_t n = 60, max_seq = 96; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, gtheta, pre, vals); + auto vbuf = make_fp16_value_cache(n, max_seq, kv_heads, head_dim, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + Params p; p.recent_frac = 0.3f; p.sink = sink; p.abs_budget = 16; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, gtheta, p); + size_t B = kept[0].size(); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, gtheta); + for (auto& k : kept) if (k.size() != B) return false; + for (size_t h = 0; h < kv_heads; ++h) + for (size_t rank = 0; rank < B; ++rank) { + std::vector exp = rope_reference(pre[h][kept[h][rank]], (double)rank, gtheta); + const uint16_t* r = krows + (rank * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(f16_to_f32(r[d]) - exp[d]) > 5e-2f) return false; + } + } + + const size_t n = 16, max_seq = 32; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, ltheta, pre, vals); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + std::vector sink_before(krows, krows + sink * kv_heads * head_dim); + + const size_t old_total = 30, B = 12; + const double Delta = (double)(old_total - B); + double accum = 0.0; + for (int cycle = 0; cycle < 3; ++cycle) { + rerope_recent_fp16(krows, kv_heads, head_dim, sink, n, ltheta, -Delta); + accum -= Delta; + } + if (std::memcmp(krows, sink_before.data(), sink_before.size() * sizeof(uint16_t)) != 0) return false; + for (size_t t = sink; t < n; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector exp = rope_reference(pre[h][t], (double)t + accum, ltheta); + const uint16_t* r = krows + (t * kv_heads + h) * head_dim; + for (size_t d = 0; d < head_dim; ++d) + if (std::abs(f16_to_f32(r[d]) - exp[d]) > 8e-2f) return false; // looser for 3-cycle drift + } + return true; +} + +bool test_rerope_local_theta_used() { + // The sliding re-rope must use the LOCAL theta; the wrong (global) theta misses the position. + const size_t n = 24, max_seq = 32, kv_heads = 1, head_dim = 16, sink = 4; + const double ltheta = 10000.0, gtheta = 1000000.0, delta = -5.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, ltheta, pre, vals); + std::vector kbuf2 = kbuf; + auto* a = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* b = reinterpret_cast(kbuf2.data() + kHeaderBytes); + + rerope_recent_fp16(a, kv_heads, head_dim, sink, n, ltheta, delta); // correct theta + rerope_recent_fp16(b, kv_heads, head_dim, sink, n, gtheta, delta); // wrong theta + + bool correct_ok = true, wrong_differs = false; + for (size_t t = sink; t < n; ++t) { + std::vector exp = rope_reference(pre[0][t], (double)t + delta, ltheta); + const uint16_t* ra = a + (t * kv_heads) * head_dim; + const uint16_t* rb = b + (t * kv_heads) * head_dim; + for (size_t d = 0; d < head_dim; ++d) { + if (std::abs(f16_to_f32(ra[d]) - exp[d]) > 5e-2f) correct_ok = false; + if (std::abs(f16_to_f32(rb[d]) - exp[d]) > 5e-2f) wrong_differs = true; + } + } + return correct_ok && wrong_differs; +} + +bool test_rerope_zero_delta_noop() { + const size_t n = 20, max_seq = 32, kv_heads = 2, head_dim = 16, sink = 4; + const double theta = 10000.0; + std::vector>> pre, vals; + auto kbuf = make_fp16_cache(n, max_seq, kv_heads, head_dim, theta, pre, vals); + std::vector before = kbuf; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + rerope_recent_fp16(krows, kv_heads, head_dim, sink, n, theta, 0.0); // zero delta + rerope_recent_fp16(krows, kv_heads, head_dim, n, n, theta, -3.0); // empty range hi<=lo + return kbuf == before; +} + +// dequant_row's NEON path must match the scalar fallback, via int8 compact and the int8 keep-set fill. +bool test_dequant_simd_matches_scalar() { + const size_t n = 64, max_seq = 96, kv_heads = 4, head_dim = 128; + const double theta = 1000000.0; + size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + Params p; p.recent_frac = 0.30f; p.sink = 4; p.abs_budget = 32; + + std::vector k8(max_seq * int8_stride, 0), v8(max_seq * int8_stride, 0); + std::vector ks(max_seq * scale_stride, 0.f), vs(max_seq * scale_stride, 0.f); + for (size_t h = 0; h < kv_heads; ++h) + for (size_t t = 0; t < n; ++t) { + std::vector k(head_dim), v(head_dim); + for (size_t d = 0; d < head_dim; ++d) { k[d] = std::sin(0.1 * (h + 1) * (t + 1) + 0.2 * d); v[d] = std::cos(0.05 * (t + 1) + 0.1 * d); } + std::vector kr = rope_reference(k, (double)t, theta); + quantize_row(kr, k8.data() + t * int8_stride + h * head_dim, ks.data() + t * scale_stride + h * groups, head_dim, kGroupSize); + quantize_row(v, v8.data() + t * int8_stride + h * head_dim, vs.data() + t * scale_stride + h * groups, head_dim, kGroupSize); + } + std::vector> kept = {{0, 5, 11, 22}, {1, 4, 13, 26}, {2, 8, 17, 30}, {3, 9, 19, 40}}; + + auto compact = [&](bool simd) { + kv_set_simd(simd); + std::vector a = k8; + std::vector b = ks; + compact_int8(a.data(), b.data(), kv_heads, head_dim, kGroupSize, kept, theta, true); + a.insert(a.end(), reinterpret_cast(b.data()), reinterpret_cast(b.data() + b.size())); + return a; + }; + auto keepsets = [&](bool simd) { + kv_set_simd(simd); + return keepsets_from_int8(k8.data(), ks.data(), n, kv_heads, head_dim, kGroupSize, theta, p); + }; + bool ok = compact(true) == compact(false) && keepsets(true) == keepsets(false); + kv_set_simd(true); + return ok; +} + +bool test_keepset_protect_union() { + // A low-score position KeyDiff would evict is force-kept via Params::protect. + const size_t n = 64; + std::vector scores(n); + for (size_t i = 0; i < n; ++i) scores[i] = 0.01f * static_cast(i); + Params p; p.sink = 4; p.recent_frac = 0.25f; p.abs_budget = 24; + const int special = 9; + + auto base = keepset_for_head(scores.data(), n, p); + if (std::find(base.begin(), base.end(), special) != base.end()) return false; + + p.protect = {special}; + auto prot = keepset_for_head(scores.data(), n, p); + if (std::find(prot.begin(), prot.end(), special) == prot.end()) return false; + if (prot.size() != 24) return false; + for (int i = 0; i < 4; ++i) + if (std::find(prot.begin(), prot.end(), i) == prot.end()) return false; + return true; +} + +bool test_rolling_protect_survives() { + // A token KeyDiff would evict survives each compaction cycle when re-protected. + const int trigger_len = 256, target_len = 128; + const size_t kv_heads = 2, head_dim = 16, max_seq = trigger_len + 8; + const double theta = 1000000.0; + + std::vector kbuf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + std::vector vbuf(kbuf.size(), 0); + auto* khdr = reinterpret_cast(kbuf.data()); + auto* vhdr = reinterpret_cast(vbuf.data()); + *khdr = Header{0, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + *vhdr = *khdr; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + std::vector special_v(head_dim); + for (size_t d = 0; d < head_dim; ++d) special_v[d] = (d % 2 ? -3.0f : 3.0f); + + auto append = [&](size_t from, size_t to, long special_abs) { + for (size_t t = from; t < to; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector k(head_dim); + for (size_t d = 0; d < head_dim; ++d) k[d] = std::sin(0.05 * (h + 1) + 0.3 * d); + std::vector kr = rope_reference(k, (double)t, theta); + std::vector vr(head_dim); + if ((long)t == special_abs) vr = special_v; + else for (size_t d = 0; d < head_dim; ++d) vr[d] = std::cos(0.07 * (t + 1) + 0.2 * d); + for (size_t d = 0; d < head_dim; ++d) { + krows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(kr[d]); + vrows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(vr[d]); + } + } + }; + + auto special_rank = [&]() -> long { + for (size_t r = 0; r < khdr->current_seq_len; ++r) { + const uint16_t* v = vrows + (r * kv_heads + 0) * head_dim; + double dot = 0, na = 0, nb = 0; + for (size_t d = 0; d < head_dim; ++d) { + float x = f16_to_f32(v[d]); dot += x * special_v[d]; na += x * x; nb += special_v[d] * special_v[d]; + } + if (dot / (std::sqrt(na) * std::sqrt(nb) + 1e-9) > 0.99) return (long)r; + } + return -1; + }; + + long special_pos = trigger_len / 2; + for (int c = 0; c < 3; ++c) { + size_t start = khdr->current_seq_len; + append(start, (size_t)trigger_len, c == 0 ? special_pos : -1); + khdr->current_seq_len = trigger_len; + vhdr->current_seq_len = trigger_len; + + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = target_len; + p.protect = {(int)special_pos}; + auto kept = keepsets_from_fp16(krows, (size_t)trigger_len, kv_heads, head_dim, theta, p); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + khdr->current_seq_len = kept[0].size(); + vhdr->current_seq_len = kept[0].size(); + + special_pos = special_rank(); + if (special_pos < 0) return false; + } + return true; +} + +bool test_cache_starts_small_and_grows() { + CactusGraph gb; + const size_t kv_heads = 2, head_dim = 64, ceiling = 100000, chunk = 100; + const size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + const size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + size_t new_kv = gb.input({chunk, kv_heads, head_dim}, Precision::FP16); + size_t state = gb.kv_cache_state(ceiling, kv_heads, head_dim, /*window*/0, /*sink*/4); + size_t append = gb.kv_cache_append(new_kv, state, /*window*/0, /*sink*/4); + gb.retain_outputs({static_cast(state), static_cast(append)}); + + bool ok = true; + std::vector expect; + auto val = [](size_t row, size_t c) { return std::sin(0.013 * row + 0.07 * c); }; + for (size_t step = 0; step < 3; ++step) { + std::vector in(chunk * int8_stride); + for (size_t r = 0; r < chunk; ++r) + for (size_t c = 0; c < int8_stride; ++c) { + float v = static_cast(val(step * chunk + r, c)); + in[r * int8_stride + c] = f32_to_f16(v); + expect.push_back(v); + } + gb.set_input(new_kv, in.data(), Precision::FP16); + gb.execute(); + if (step == 0 && reinterpret_cast(gb.get_output(state))->max_seq_len != 256) ok = false; + } + + const auto* hdr = reinterpret_cast(gb.get_output(state)); + if (hdr->max_seq_len != 512) ok = false; + if (hdr->current_seq_len != 300) ok = false; + + const auto* base = static_cast(gb.get_output(state)); + const int8_t* i8 = reinterpret_cast(base + kHeaderBytes); + const float* sc = reinterpret_cast(base + kHeaderBytes + 512 * int8_stride); + for (size_t r = 0; r < 300 && ok; ++r) + for (size_t h = 0; h < kv_heads; ++h) + for (size_t d = 0; d < head_dim; ++d) { + float dq = static_cast(i8[r * int8_stride + h * head_dim + d]) * + sc[r * scale_stride + h * groups + d / kGroupSize]; + if (std::abs(dq - expect[r * int8_stride + h * head_dim + d]) > 0.05f) ok = false; + } + return ok; +} + +bool test_sliding_layer_fixed_capacity() { + CactusGraph gb; + const size_t kv_heads = 2, head_dim = 64, ceiling = 100000, window = 512, sink = 4; + size_t state = gb.kv_cache_state(ceiling, kv_heads, head_dim, window, sink); + gb.retain_outputs({static_cast(state)}); + gb.execute(); + auto* hdr = reinterpret_cast(gb.get_output(state)); + return hdr->max_seq_len == window + sink + 1; +} + +bool test_special_rows_remap_through_kept() { + std::vector rows = {3, 7, 10}; + std::vector kept = {0, 3, 5, 7, 9, 10, 12}; + return remap_rows_through_kept(rows, kept) == std::vector{1, 3, 5}; +} + +bool test_per_head_protect_keeps_specials() { + const double theta = 1000000.0; + const size_t n = 64, kv_heads = 3, head_dim = 16; + std::vector>> pre, val; + auto kbuf = make_fp16_cache(n, n, kv_heads, head_dim, theta, pre, val); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = 24; + auto unrope = unrope_table(n, head_dim, theta); + std::vector> protect(kv_heads); + for (size_t h = 0; h < kv_heads; ++h) protect[h] = {static_cast(20 + h)}; + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, unrope, p, protect); + if (kept.size() != kv_heads) return false; + for (size_t h = 0; h < kv_heads; ++h) + if (std::find(kept[h].begin(), kept[h].end(), static_cast(20 + h)) == kept[h].end()) return false; + return true; +} + +bool test_preflight_specials_exceed_budget() { + SpecialRowTracker tracked; + std::vector rows; + for (int r = 0; r < 50; ++r) rows.push_back(r); + tracked.add_appended(0, 2, rows); + SpecialRowTracker untracked; + return tracked.max_reserved(0, 4, {}) == 50 && tracked.max_reserved(0, 4, {}) > 40 && + untracked.max_reserved(0, 4, rows) == 50 && untracked.max_reserved(0, 4, rows) > 40; +} + +bool test_empty_protect_per_head_uses_params_fallback() { + const double theta = 1000000.0; + const size_t n = 64, kv_heads = 3, head_dim = 16; + std::vector>> pre, val; + auto kbuf = make_fp16_cache(n, n, kv_heads, head_dim, theta, pre, val); + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = 24; p.protect = {25}; + auto unrope = unrope_table(n, head_dim, theta); + auto kept = keepsets_from_fp16(krows, n, kv_heads, head_dim, unrope, p, /*protect_per_head=*/{}); + if (kept.size() != kv_heads) return false; + for (size_t h = 0; h < kv_heads; ++h) + if (std::find(kept[h].begin(), kept[h].end(), 25) == kept[h].end()) return false; + return true; +} + +bool test_all_heads_keep_special_across_cycles() { + const int trigger_len = 256, target_len = 128; + const size_t kv_heads = 4, head_dim = 16, max_seq = trigger_len + 8; + const double theta = 1000000.0; + std::vector kbuf(kHeaderBytes + max_seq * kv_heads * head_dim * sizeof(uint16_t), 0); + std::vector vbuf(kbuf.size(), 0); + auto* khdr = reinterpret_cast(kbuf.data()); + auto* vhdr = reinterpret_cast(vbuf.data()); + *khdr = Header{0, max_seq, kv_heads, head_dim, 4, {0, 0, 0}}; + *vhdr = *khdr; + auto* krows = reinterpret_cast(kbuf.data() + kHeaderBytes); + auto* vrows = reinterpret_cast(vbuf.data() + kHeaderBytes); + + auto special_val = [&](size_t h) { + std::vector v(head_dim); + for (size_t d = 0; d < head_dim; ++d) v[d] = (d % 2 ? -2.0f : 2.0f) * static_cast(h + 1); + return v; + }; + const long special_abs = trigger_len / 2; // a MIDDLE position (not sink, not recent) + auto append = [&](size_t from, size_t to) { + for (size_t t = from; t < to; ++t) + for (size_t h = 0; h < kv_heads; ++h) { + std::vector k(head_dim); + // Special key is non-distinctive (KeyDiff would evict it); other rows differ per (h,t). + for (size_t d = 0; d < head_dim; ++d) + k[d] = (static_cast(t) == special_abs) ? 0.01f + : std::sin(0.11 * (h + 1) * (t + 1) + 0.3 * d) + 0.05 * d; + std::vector kr = rope_reference(k, static_cast(t), theta); + std::vector vr(head_dim); + if (static_cast(t) == special_abs) vr = special_val(h); + else for (size_t d = 0; d < head_dim; ++d) vr[d] = std::cos(0.07 * (t + 1) + 0.2 * d); + for (size_t d = 0; d < head_dim; ++d) { + krows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(kr[d]); + vrows[(t * kv_heads + h) * head_dim + d] = f32_to_f16(vr[d]); + } + } + }; + auto special_rank = [&](size_t h) -> long { + std::vector sv = special_val(h); + for (size_t r = 0; r < khdr->current_seq_len; ++r) { + const uint16_t* v = vrows + (r * kv_heads + h) * head_dim; + double dot = 0, na = 0, nb = 0; + for (size_t d = 0; d < head_dim; ++d) { float x = f16_to_f32(v[d]); dot += x * sv[d]; na += x * x; nb += sv[d] * sv[d]; } + if (dot / (std::sqrt(na) * std::sqrt(nb) + 1e-9) > 0.99) return static_cast(r); + } + return -1; + }; + + SpecialRowTracker tracker; + Params p; p.sink = 4; p.recent_frac = 0.30f; p.abs_budget = target_len; + auto unrope = unrope_table(trigger_len, head_dim, theta); + for (int c = 0; c < 3; ++c) { + size_t start = khdr->current_seq_len; + append(start, static_cast(trigger_len)); + khdr->current_seq_len = trigger_len; vhdr->current_seq_len = trigger_len; + std::vector appended; + if (c == 0) appended.push_back(static_cast(special_abs)); + tracker.add_appended(0, kv_heads, appended); + auto kept = keepsets_from_fp16(krows, static_cast(trigger_len), kv_heads, head_dim, + unrope, p, tracker.protect(0)); + compact_fp16(krows, vrows, kv_heads, head_dim, kept, theta); + tracker.remap(0, kept); + size_t B = kept.empty() ? 0 : kept[0].size(); + khdr->current_seq_len = B; vhdr->current_seq_len = B; + tracker.set_tracked_len(B); + } + long r0 = special_rank(0); + if (r0 < 0) return false; + bool diverged = false; + for (size_t h = 0; h < kv_heads; ++h) { + long r = special_rank(h); + if (r < 0) return false; + if (r != r0) diverged = true; // ranks differ across heads -> a real per-head scenario, not trivially aligned + } + return diverged; +} + +bool test_shrink_cache_buffer_preserves_rows() { + CactusGraph gb; + const size_t kv_heads = 2, head_dim = 64, ceiling = 100000, chunk = 300; + const size_t groups = (head_dim + kGroupSize - 1) / kGroupSize; + const size_t int8_stride = kv_heads * head_dim, scale_stride = kv_heads * groups; + size_t new_kv = gb.input({chunk, kv_heads, head_dim}, Precision::FP16); + size_t state = gb.kv_cache_state(ceiling, kv_heads, head_dim, /*window*/0, /*sink*/4); + size_t append = gb.kv_cache_append(new_kv, state, /*window*/0, /*sink*/4); + gb.retain_outputs({static_cast(state), static_cast(append)}); + + std::vector expect; + auto val = [](size_t row, size_t c) { return std::sin(0.013 * row + 0.07 * c); }; + std::vector in(chunk * int8_stride); + for (size_t r = 0; r < chunk; ++r) + for (size_t c = 0; c < int8_stride; ++c) { + float v = static_cast(val(r, c)); + in[r * int8_stride + c] = f32_to_f16(v); + expect.push_back(v); + } + gb.set_input(new_kv, in.data(), Precision::FP16); + gb.execute(); // appends 300 -> grows 256->512, occupancy 300 + + auto* hdr = reinterpret_cast(gb.get_output(state)); + if (hdr->max_seq_len != 512 || hdr->current_seq_len != 300) return false; + + gb.shrink_cache_buffer(state, 256); // clamped up to occupancy 300 + hdr = reinterpret_cast(gb.get_output(state)); + if (hdr->max_seq_len != 300 || hdr->current_seq_len != 300) return false; + + const auto* base = static_cast(gb.get_output(state)); + const int8_t* i8 = reinterpret_cast(base + kHeaderBytes); + const float* sc = reinterpret_cast(base + kHeaderBytes + 300 * int8_stride); + bool ok = true; + for (size_t r = 0; r < 300 && ok; ++r) + for (size_t h = 0; h < kv_heads; ++h) + for (size_t d = 0; d < head_dim; ++d) { + float dq = static_cast(i8[r * int8_stride + h * head_dim + d]) * + sc[r * scale_stride + h * groups + d / kGroupSize]; + if (std::abs(dq - expect[r * int8_stride + h * head_dim + d]) > 0.05f) ok = false; + } + return ok; +} + +int main() { + TestUtils::TestRunner runner("KV Compress Free-Function Tests"); + runner.run_test("cache_starts_small_and_grows", test_cache_starts_small_and_grows()); + runner.run_test("sliding_layer_fixed_capacity", test_sliding_layer_fixed_capacity()); + runner.run_test("compact_fp16_cache", test_compact_fp16_cache()); + runner.run_test("dense_check_full_budget", test_dense_check_full_budget()); + runner.run_test("rope_renumber_contiguous", test_rope_renumber_contiguous()); + runner.run_test("fp16_storage_round_trip", test_fp16_storage_round_trip()); + runner.run_test("gemma_layer_selection", test_gemma_layer_selection()); + runner.run_test("compact_int8_cache", test_compact_int8_cache()); + runner.run_test("rerope_recent_fp16_uniform", test_rerope_recent_fp16_uniform()); + runner.run_test("rerope_recent_int8", test_rerope_recent_int8()); + runner.run_test("rotate_int8_row_matches_inline", test_rotate_int8_row_matches_inline()); + runner.run_test("sliding_plus_global_rolling", test_sliding_plus_global_rolling()); + runner.run_test("rerope_local_theta_used", test_rerope_local_theta_used()); + runner.run_test("rerope_zero_delta_noop", test_rerope_zero_delta_noop()); + runner.run_test("rolling_bounded_compaction", test_rolling_bounded_compaction()); + runner.run_test("config_parse_rolling_fields", test_config_parse_rolling_fields()); + runner.run_test("trigger_zero_gates_rolling", test_trigger_zero_gates_rolling()); + runner.run_test("degenerate_rolling_config_disabled", test_degenerate_rolling_config_disabled()); + runner.run_test("env_override_parse", test_env_override_parse()); + runner.run_test("dequant_simd_matches_scalar", test_dequant_simd_matches_scalar()); + runner.run_test("keepset_protect_union", test_keepset_protect_union()); + runner.run_test("rolling_protect_survives", test_rolling_protect_survives()); + runner.run_test("special_rows_remap_through_kept", test_special_rows_remap_through_kept()); + runner.run_test("per_head_protect_keeps_specials", test_per_head_protect_keeps_specials()); + runner.run_test("preflight_specials_exceed_budget", test_preflight_specials_exceed_budget()); + runner.run_test("empty_protect_per_head_uses_params_fallback", test_empty_protect_per_head_uses_params_fallback()); + runner.run_test("all_heads_keep_special_across_cycles", test_all_heads_keep_special_across_cycles()); + runner.run_test("shrink_cache_buffer_preserves_rows", test_shrink_cache_buffer_preserves_rows()); + runner.print_summary(); + return runner.all_passed() ? 0 : 1; +} diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp index 56382474d..a7fe93cb0 100644 --- a/cactus-engine/tests/test_llm.cpp +++ b/cactus-engine/tests/test_llm.cpp @@ -1,4 +1,5 @@ #include "test_utils.h" +#include "../src/utils.h" #include #include #include @@ -14,9 +15,39 @@ #endif using namespace EngineTestUtils; +using namespace cactus::engine; +using cactus::ffi::partition_thinking_response; static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); +static bool check_partition(const std::string& input, + const std::string& expected_thinking, + const std::string& expected_content) { + std::string thinking, content; + partition_thinking_response(input, thinking, content); + if (thinking != expected_thinking) { + std::cerr << " thinking: '" << thinking << "' != '" << expected_thinking << "'\n"; + return false; + } + if (content != expected_content) { + std::cerr << " content: '" << content << "' != '" << expected_content << "'\n"; + return false; + } + return true; +} + +static cactus_model_t load_gemma4_or_skip() { + if (!g_model_path) { std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; return nullptr; } + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { std::cout << " [WARN] Could not load model; skipping\n"; return nullptr; } + if (static_cast(model)->model->get_config().model_type != Config::ModelType::GEMMA4) { + std::cout << " [WARN] chosen model is not Gemma4; skipping\n"; + cactus_destroy(model); + return nullptr; + } + return model; +} + static const char* g_options = R"({ "max_tokens": 256, "stop_sequences": ["<|im_end|>", ""], @@ -559,6 +590,118 @@ bool test_1k_context() { }, nullptr, 100); } +bool test_partition_thinking_response() { + return check_partition("<|channel>reasonanswer", "reason", "answer") + && check_partition("<|channel>\n reason\n\n\nanswer", "reason", "answer") + && check_partition("no tags here", "", "no tags here") + && check_partition("reasonanswer", "reason", "answer") + && check_partition("<|channel>thought1text1<|channel>thought2text2", + "thought1\nthought2", "text1text2"); +} + +bool test_prompt_gemma4_retains_thinking() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tok = handle->model->get_tokenizer(); + + std::vector msgs = { + {"user", "hello", "", {}, {}, 0, {}}, + {"assistant", "<|channel>internal reasoningvisible response", "", {}, {}, 0, {}}, + {"user", "followup", "", {}, {}, 0, {}} + }; + + std::string prompt = tok->format_chat_prompt(msgs, true, "", true); + cactus_destroy(model); + + bool has_visible = prompt.find("visible response") != std::string::npos; + bool has_reasoning = prompt.find("internal reasoning") != std::string::npos; + bool has_channel_tags = prompt.find("<|channel>") != std::string::npos + && prompt.find("") != std::string::npos; + + if (!has_visible) std::cerr << " missing visible response in prompt\n"; + if (!has_reasoning) std::cerr << " thinking content not retained in assistant turn\n"; + if (!has_channel_tags) std::cerr << " channel tags not retained in prompt\n"; + + return has_visible && has_reasoning && has_channel_tags; +} + +bool test_complete_gemma4_thinking_api_clean() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + const char* msgs = R"([{"role": "user", "content": "What is 2+2?"}])"; + char buf[8192]; + + int r = cactus_complete(model, msgs, buf, sizeof(buf), + R"({"max_tokens":128,"enable_thinking_if_supported":true,"telemetry_enabled":false})", + nullptr, nullptr, nullptr, nullptr, 0); + std::string resp(buf); + cactus_destroy(model); + + std::string response = EngineTestUtils::json_string(resp, "response"); + bool ok = r > 0 + && resp.find("\"success\":true") != std::string::npos + && response.find("<|channel>") == std::string::npos + && response.find("") == std::string::npos; + if (!ok) std::cerr << " thinking-enabled completion api not clean: " << resp << "\n"; + return ok; +} + +bool test_multiturn_thinking_persist() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tokenizer = handle->model->get_tokenizer(); + const char* options = R"({"max_tokens":128,"temperature":0,"top_k":1,"enable_thinking_if_supported":true,"telemetry_enabled":false,"auto_handoff":false})"; + const char* turn1_msgs = R"([{"role": "user", "content": "My name is Alice. Please remember this."}])"; + char buf[16384]; + + int r1 = cactus_complete(model, turn1_msgs, buf, sizeof(buf), options, nullptr, nullptr, nullptr, nullptr, 0); + if (r1 <= 0) { std::cerr << " Turn 1 failed\n"; cactus_destroy(model); return false; } + + std::vector processed_after_t1 = handle->processed_tokens; + std::string turn1_json(buf); + std::string context_response = EngineTestUtils::json_string(turn1_json, "context_response"); + if (context_response.empty()) { + std::cerr << " context_response missing from turn 1 result\n"; + cactus_destroy(model); + return false; + } + + std::vector t2_chat = { + {"user", "My name is Alice. Please remember this.", "", {}, {}, 0, {}}, + {"assistant", context_response, "", {}, {}, 0, {}}, + {"user", "What is my name?", "", {}, {}, 0, {}} + }; + std::vector t2_prompt_tokens = tokenizer->encode(tokenizer->format_chat_prompt(t2_chat, true, "", true)); + + bool prefix_ok = (t2_prompt_tokens.size() >= processed_after_t1.size()) && + std::equal(processed_after_t1.begin(), processed_after_t1.end(), t2_prompt_tokens.begin()); + std::cout << " Prefix match (cache reuse): " << (prefix_ok ? "YES" : "NO") << "\n"; + + std::string escaped = EngineTestUtils::escape_json(context_response); + std::string turn2_json = R"([{"role": "user", "content": "My name is Alice. Please remember this."},{"role": "assistant", "content": ")" + + escaped + R"("},{"role": "user", "content": "What is my name?"}])"; + + int r2 = cactus_complete(model, turn2_json.c_str(), buf, sizeof(buf), options, nullptr, nullptr, nullptr, nullptr, 0); + if (r2 <= 0) { std::cerr << " Turn 2 failed\n"; cactus_destroy(model); return false; } + + std::string turn2_result(buf); + std::string turn2_response = EngineTestUtils::json_string(turn2_result, "response"); + bool mentions_alice = turn2_response.find("Alice") != std::string::npos + || turn2_response.find("alice") != std::string::npos; + std::cout << " Turn 2 mentions Alice: " << (mentions_alice ? "YES" : "NO") << "\n"; + + cactus_destroy(model); + + if (!prefix_ok) std::cerr << " FAIL: re-rendered history did not prefix-match cache\n"; + if (!mentions_alice) std::cerr << " FAIL: turn 2 did not recall the name\n"; + return prefix_ok && mentions_alice; +} + int main() { TestUtils::TestRunner runner("LLM Tests"); runner.run_test("1k_context", test_1k_context()); @@ -570,6 +713,10 @@ int main() { runner.run_test("tool_calls", test_tool_call()); runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations()); runner.run_test("tool_calls_with_three_tools", test_tool_call_with_three_tools()); + runner.run_test("partition_thinking_response", test_partition_thinking_response()); + runner.run_test("prompt_retains_thinking", test_prompt_gemma4_retains_thinking()); + runner.run_test("complete_thinking_api_clean", test_complete_gemma4_thinking_api_clean()); + runner.run_test("multiturn_thinking_persist", test_multiturn_thinking_persist()); runner.print_summary(); return runner.all_passed() ? 0 : 1; } diff --git a/cactus-engine/tests/test_utils.h b/cactus-engine/tests/test_utils.h index 4a99f4f3f..19aba3819 100644 --- a/cactus-engine/tests/test_utils.h +++ b/cactus-engine/tests/test_utils.h @@ -31,6 +31,19 @@ class TestRunner { namespace EngineTestUtils { +// Double-precision rotate_half RoPE oracle for kv_compress tests (mirrors norms_rope.cpp). +inline std::vector rope_reference(const std::vector& v, double pos, double theta) { + size_t d = v.size(), half = d / 2; + std::vector o(d); + for (size_t i = 0; i < half; ++i) { + double inv = std::pow(theta, -(2.0 * (double)i) / (double)d); + double a = pos * inv, c = std::cos(a), s = std::sin(a); + o[i] = (float)(v[i] * c - v[i + half] * s); + o[i + half] = (float)(v[i + half] * c + v[i] * s); + } + return o; +} + struct Timer { std::chrono::high_resolution_clock::time_point start; Timer(); diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 3d7db196e..1cf4e87cc 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -707,6 +707,9 @@ class CactusGraph { const std::vector& output_shape, const OpParams& params = {}); const BufferDesc& get_output_buffer(size_t node_id) const; OpType get_node_op_type(size_t node_id) const; + size_t get_node_window_size(size_t node_id) const; + void steal_cache_buffer(size_t dst_node, CactusGraph& src, size_t src_node); + void shrink_cache_buffer(size_t node_id, size_t new_capacity); void allocate_buffers(); size_t get_node_count() const; diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index bf6c2f491..0ae9abdfc 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -1243,6 +1243,10 @@ OpType CactusGraph::get_node_op_type(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->op_type; } +size_t CactusGraph::get_node_window_size(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->params.window_size; +} + size_t CactusGraph::persistent(size_t source_node) { const auto& source_buffer = get_output_buffer(source_node); OpParams params; diff --git a/cactus-graph/src/ops_cache.cpp b/cactus-graph/src/ops_cache.cpp index 7df2a588a..2ed462a23 100644 --- a/cactus-graph/src/ops_cache.cpp +++ b/cactus-graph/src/ops_cache.cpp @@ -4,6 +4,7 @@ #include #include #include +#include namespace { @@ -87,6 +88,50 @@ inline bool use_fp16_kv_cache() { return cached; } +constexpr size_t kInitialCacheEntries = 256; + +inline bool resize_cache_buffer(BufferDesc& buf, size_t new_max) { + auto* meta = get_meta(buf); + size_t cur = meta->max_seq_len; + const size_t current_seq = meta->current_seq_len; + if (new_max == cur || new_max < current_seq) return false; + + const size_t kv_heads = meta->num_kv_heads; + const size_t hdim = meta->head_dim; + const bool fp16_cache = buf.precision == Precision::FP16; + + size_t total = fp16_cache ? fp16_cache_elements(new_max, kv_heads, hdim) + : cache_buffer_size(new_max, kv_heads, hdim); + BufferDesc resized({total}, fp16_cache ? Precision::FP16 : Precision::INT8); + resized.allocate(); + std::memset(resized.get_data(), 0, resized.byte_size); + + std::memcpy(resized.get_data(), buf.get_data(), sizeof(CacheMetadata)); + if (fp16_cache) { + std::memcpy(get_fp16_data(resized), get_fp16_data(buf), + current_seq * kv_heads * hdim * sizeof(__fp16)); + } else { + std::memcpy(get_int8_data(resized), get_int8_data(buf), current_seq * kv_heads * hdim); + const size_t groups = (hdim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + std::memcpy(get_scales(resized, new_max, kv_heads, hdim), + get_scales(buf, cur, kv_heads, hdim), + current_seq * kv_heads * groups * sizeof(float)); + } + get_meta(resized)->max_seq_len = new_max; + buf = std::move(resized); + return true; +} + +inline bool grow_cache_buffer(BufferDesc& buf, size_t needed, size_t ceiling) { + size_t cur = get_meta(buf)->max_seq_len; + if (needed <= cur || cur >= ceiling) return false; + size_t new_max = cur; + while (new_max < needed) new_max <<= 1; + if (new_max > ceiling) new_max = ceiling; + if (new_max <= cur) return false; + return resize_cache_buffer(buf, new_max); +} + } // namespace void compute_kv_cache_state_node( @@ -96,7 +141,11 @@ void compute_kv_cache_state_node( if (node.output_buffer.get_data()) return; - size_t max_seq = node.params.max_cache_seq_len; + size_t ceiling = node.params.max_cache_seq_len; + size_t window = node.params.window_size; + bool sliding = window > 0 && window < ceiling; + size_t max_seq = sliding ? std::min(ceiling, window + node.params.cache_sink_size + 1) + : std::min(ceiling, kInitialCacheEntries); size_t kv_heads = node.params.num_kv_heads; size_t hdim = node.params.head_dim; const bool fp16_cache = use_fp16_kv_cache(); @@ -136,12 +185,20 @@ void compute_kv_cache_append_node( size_t int8_stride = kv_heads * hdim; size_t scale_stride = kv_heads * num_groups; + size_t ceiling = nodes[node_index_map.at(node.input_ids[1])]->params.max_cache_seq_len; + bool sliding = node.params.window_size > 0 && node.params.window_size < ceiling; + if (!sliding && current_len + new_seq_len > max_len) { + if (grow_cache_buffer(cache_buf, current_len + new_seq_len, ceiling)) { + meta = get_meta(cache_buf); + max_len = meta->max_seq_len; + } + } + if (cache_buf.precision == Precision::FP16) { size_t stride = kv_heads * hdim; __fp16* fp16_base = get_fp16_data(cache_buf); const __fp16* source = new_kv.data_as<__fp16>(); - size_t window = node.params.window_size; - if (window == 0) window = max_len; + size_t window = sliding ? node.params.window_size : max_len; size_t new_total = current_len + new_seq_len; bool needs_eviction = new_total > window; @@ -185,8 +242,7 @@ void compute_kv_cache_append_node( int8_t* int8_base = get_int8_data(cache_buf); float* scale_base = get_scales(cache_buf, max_len, kv_heads, hdim); - size_t window = node.params.window_size; - if (window == 0) window = max_len; + size_t window = sliding ? node.params.window_size : max_len; size_t new_total = current_len + new_seq_len; bool needs_eviction = new_total > window; @@ -467,3 +523,21 @@ void compute_recurrent_cache_write_node( auto& cache_buf = nodes[node_index_map.at(node.input_ids[1])]->output_buffer; std::memcpy(cache_buf.get_data(), src.get_data(), src.byte_size); } + +void CactusGraph::steal_cache_buffer(size_t dst_node, CactusGraph& src, size_t src_node) { + auto& dst = nodes_[node_index_map_.at(dst_node)]; + auto& s = src.nodes_[src.node_index_map_.at(src_node)]; + // Buffer carries its own runtime precision (may be fp16); only op_type is invariant pre-move. + assert(dst->op_type == s->op_type); + dst->output_buffer = std::move(s->output_buffer); +} + +void CactusGraph::shrink_cache_buffer(size_t node_id, size_t new_capacity) { + auto& buf = nodes_[node_index_map_.at(node_id)]->output_buffer; + if (!buf.get_data()) return; + auto* meta = get_meta(buf); + size_t target = std::max(new_capacity, meta->current_seq_len); + if (target >= meta->max_seq_len) return; + resize_cache_buffer(buf, target); +} + diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index 59f5acb7a..ac8708fc6 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -234,6 +234,8 @@ def create_parser(): help="Initial prompt to send immediately") run_parser.add_argument("--input-ids", default=None, help="Comma-separated token ids for transpiled causal-LM bundles") + run_parser.add_argument("--input-ids-file", default=None, + help="File containing token ids for transpiled causal-LM bundles") run_parser.add_argument("--max-new-tokens", type=_positive_int, default=None, help="Maximum tokens to generate for transpiled causal-LM bundles") run_parser.add_argument("--result-json", default=None, @@ -381,6 +383,8 @@ def create_parser(): help="NPU audio encoder quant: 0=fp16, 4=int4, 8=int8 (default: 8)") transpile_parser.add_argument("--npu-vision-quantize", type=int, choices=[0, 4, 8], default=None, help="NPU vision encoder quant: 0=fp16, 4=int4, 8=int8 (default: 0; int4 degrades Gemma4 vision)") + transpile_parser.add_argument("--cache-context-length", default=None, + help="KV cache context length for cached decode graphs (default: model config)") return parser diff --git a/python/cactus/cli/convert.py b/python/cactus/cli/convert.py index 8c86e337d..03fdc1fb1 100644 --- a/python/cactus/cli/convert.py +++ b/python/cactus/cli/convert.py @@ -147,6 +147,8 @@ def cmd_transpile(args): extra_args.extend(["--npu-audio-quantize", str(args.npu_audio_quantize)]) if args.npu_vision_quantize is not None: extra_args.extend(["--npu-vision-quantize", str(args.npu_vision_quantize)]) + if args.cache_context_length is not None: + extra_args.extend(["--cache-context-length", str(args.cache_context_length)]) return run_transpile( args.model_id, diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index 6071d46bb..a24e4d6d0 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -177,6 +177,7 @@ class TranspileOptions: npu_quantize: int | None = None npu_audio_quantize: int | None = None npu_vision_quantize: int | None = None + cache_context_length: str | int | None = None def ensure_bundle(model_id, *, bits=4, token=None, @@ -289,6 +290,8 @@ def ensure_bundle(model_id, *, bits=4, token=None, extra_args.extend(["--npu-audio-quantize", str(int(opts.npu_audio_quantize))]) if opts.npu_vision_quantize is not None: extra_args.extend(["--npu-vision-quantize", str(int(opts.npu_vision_quantize))]) + if opts.cache_context_length is not None: + extra_args.extend(["--cache-context-length", str(opts.cache_context_length)]) rc = run_transpile(model_id, extra_args=extra_args) if rc != 0: diff --git a/python/cactus/cli/run.py b/python/cactus/cli/run.py index 52c32ac4d..f69399039 100644 --- a/python/cactus/cli/run.py +++ b/python/cactus/cli/run.py @@ -39,6 +39,7 @@ def cmd_run(args) -> int: ("--image", args.image), ("--audio", args.audio), ("--input-ids", args.input_ids), + ("--input-ids-file", args.input_ids_file), ("--max-new-tokens", args.max_new_tokens), ("--result-json", args.result_json), ("--confidence-threshold", args.confidence_threshold), diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py index 60394ebc6..332096a31 100644 --- a/python/cactus/transpile/component_pipeline.py +++ b/python/cactus/transpile/component_pipeline.py @@ -16,6 +16,7 @@ from cactus.transpile.lower import transpile_preoptimized_ir from cactus.transpile.optimize_graph import FusionConfig from cactus.transpile.optimize_graph import optimize_graph +from cactus.transpile.optimize_graph import precompute_rope_tables @dataclass @@ -96,8 +97,12 @@ def capture_component_spec( raw_ir_graph = copy.deepcopy(captured.ir_graph) optimized_ir_graph = copy.deepcopy(captured.ir_graph) canonicalize_exported_graph(optimized_ir_graph) - optimize_graph(optimized_ir_graph, config=fusion_config) - transpiled_graph = transpile_preoptimized_ir(copy.deepcopy(optimized_ir_graph)) + optimize_graph(optimized_ir_graph, config=fusion_config, precompute_rope=False) + # Bake rope tables only into the lowered graph.cactus; the saved optimized IR stays un-baked. + lowered_ir = copy.deepcopy(optimized_ir_graph) + if precompute_rope_tables(lowered_ir): + canonicalize_exported_graph(lowered_ir) + transpiled_graph = transpile_preoptimized_ir(lowered_ir) if len(spec.output_keys) != len(transpiled_graph.outputs): raise ValueError( f"component {spec.component} declared {len(spec.output_keys)} output keys but lowered " diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index 5c57ac9f3..558a2756f 100644 --- a/python/cactus/transpile/component_plan.py +++ b/python/cactus/transpile/component_plan.py @@ -28,6 +28,14 @@ def _plan_from_profile( if config is not None and profile.needs_image: has_vision = _has_dict_config(config, "vision_config", "visual_config", "image_config") if not has_vision: + architectures = [str(a).lower() for a in (config.get("architectures") or ())] + if any("qwen3forcausallm" in a for a in architectures): + return ComponentPlan( + task="causal_lm_logits", + components=("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk", "decoder_media_step"), + force_component_pipeline=True, + ) return ComponentPlan( task="causal_lm_logits", components=("decoder", "decoder_step"), diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index 7fe99854a..d13b4cfea 100644 --- a/python/cactus/transpile/hf_model.py +++ b/python/cactus/transpile/hf_model.py @@ -2962,6 +2962,11 @@ def main() -> int: "so the saved bundle can do greedy autoregressive decoding." ), ) + parser.add_argument( + "--cache-context-length", + default="auto", + help="KV cache context length for cached decode graphs. Use auto to read the model config.", + ) parser.add_argument( "--audio-file", default="", @@ -3267,6 +3272,7 @@ def main() -> int: weights_dir=weights_dir, inputs_metadata=prepared.metadata, components=requested_components, + cache_context_length=args.cache_context_length, ) use_component_pipeline = False if args.component_pipeline == "on": diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index 5f2c3d1e2..e16af4f9b 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -238,7 +238,8 @@ def _lower_attention_with_internal_kv_cache( ) if layer_key not in cache_states: default_cache_len = max(512, int(key_shape[1])) - max_cache_seq_len = int(ir.meta.get("max_cache_seq_len", default_cache_len) or default_cache_len) + requested_cache_len = node.meta.get("max_cache_seq_len", ir.meta.get("max_cache_seq_len", default_cache_len)) + max_cache_seq_len = max(default_cache_len, int(requested_cache_len or default_cache_len)) sink_size = int(ir.meta.get("cache_sink_size", 4) or 4) window_size = int(node.attrs.get("window_size", 0) or 0) k_cache = g.kv_cache_state( diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index b9aefe0cc..0d427dc9f 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -63,12 +63,20 @@ def _model_name_or_path(model: torch.nn.Module) -> str: def _transpile_graph_meta(model: torch.nn.Module, *, adapter_family: str, adapter_type: str, input_names: tuple[str, ...]) -> dict[str, object]: - return { + meta: dict[str, object] = { "adapter_family": adapter_family, "adapter_type": adapter_type, "model_name_or_path": _model_name_or_path(model), "input_names": input_names, } + # Rope-table precompute falls back to this on components that don't reserve a KV cache. + for candidate in _config_candidates_for_context_length(model): + for key in ("max_position_embeddings", "context_length", "model_max_length", "n_positions"): + value = _config_int_value(candidate, key) + if value is not None: + meta["max_position_embeddings"] = value + return meta + return meta def _extract_tensor_output(output: object, *, preferred_field: str | None = None) -> torch.Tensor: @@ -290,6 +298,97 @@ def _gemma4_text_config(multimodal_backbone: torch.nn.Module) -> object: return config +def _parse_cache_context_length(value: str | int | None) -> int | None: + if value is None: + return None + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"", "auto"}: + return None + parsed = int(normalized) + else: + parsed = int(value) + if parsed <= 0: + raise ValueError("--cache-context-length must be a positive integer or auto") + return parsed + + +def _config_int_value(config: object, key: str) -> int | None: + if config is None: + return None + if isinstance(config, dict): + raw_value = config.get(key) + else: + raw_value = getattr(config, key, None) + if raw_value is None: + return None + parsed = int(raw_value) + return parsed if parsed > 0 else None + + +def _config_candidates_for_context_length(model: torch.nn.Module) -> tuple[object, ...]: + candidates: list[object] = [] + seen: set[int] = set() + + def _add(candidate: object) -> None: + if candidate is None: + return + candidate_id = id(candidate) + if candidate_id in seen: + return + seen.add(candidate_id) + candidates.append(candidate) + + for owner in ( + model, + getattr(model, "model", None), + getattr(getattr(model, "model", None), "language_model", None), + ): + config = getattr(owner, "config", None) + _add(config) + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + _add(get_text_config()) + for attr in ("text_config", "llm_config", "language_config", "decoder_config"): + _add(getattr(config, attr, None)) + + return tuple(candidates) + + +def _cache_context_length( + model: torch.nn.Module, + *, + input_seq_len: int, + cache_context_length: str | int | None, + fallback_extra_tokens: int, +) -> int: + explicit_length = _parse_cache_context_length(cache_context_length) + if explicit_length is not None: + return explicit_length + + for candidate in _config_candidates_for_context_length(model): + for key in ("max_position_embeddings", "context_length", "model_max_length", "n_positions"): + value = _config_int_value(candidate, key) + if value is not None: + return value + return input_seq_len + int(fallback_extra_tokens) + + +def _max_cache_seq_len( + model: torch.nn.Module, + input_ids: torch.Tensor, + cache_context_length: str | int | None, + *, + fallback_extra_tokens: int, +) -> int: + return max(1024, _cache_context_length( + model, + input_seq_len=int(input_ids.shape[1]), + cache_context_length=cache_context_length, + fallback_extra_tokens=fallback_extra_tokens, + )) + + def _gemma4_special_token_ids(multimodal_backbone: torch.nn.Module) -> tuple[int, int, int]: config = getattr(multimodal_backbone, "config", None) text_config = _gemma4_text_config(multimodal_backbone) @@ -407,9 +506,9 @@ def _gemma4_text_attention_forward( query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2) query_states = query_states.transpose(1, 2) - kv_shared_layer_index = getattr(attn, "kv_shared_layer_index", None) - if bool(getattr(attn, "is_kv_shared_layer", False)) and kv_shared_layer_index in shared_kv_states: - key_states, value_states = shared_kv_states[int(kv_shared_layer_index)] + kv_shared_key = getattr(attn, "layer_type", getattr(attn, "kv_shared_layer_index", None)) + if bool(getattr(attn, "is_kv_shared_layer", False)) and kv_shared_key in shared_kv_states: + key_states, value_states = shared_kv_states[kv_shared_key] key_states = key_states.to(query_states.device) value_states = value_states.to(query_states.device) else: @@ -424,10 +523,9 @@ def _gemma4_text_attention_forward( value_states = value_states.transpose(1, 2) layer_idx = getattr(attn, "layer_idx", None) - if layer_idx is not None and ( - bool(getattr(attn, "store_full_length_kv", False)) - or not bool(getattr(attn, "is_kv_shared_layer", False)) - ): + if bool(getattr(attn, "store_full_length_kv", False)): + shared_kv_states[kv_shared_key] = (key_states, value_states) + elif layer_idx is not None and not bool(getattr(attn, "is_kv_shared_layer", False)): shared_kv_states[int(layer_idx)] = (key_states, value_states) attention_interface = eager_attention_forward @@ -3249,6 +3347,7 @@ def _build_gemma4_multimodal_component_specs( named_tensors: dict[str, torch.Tensor], weights_dir: str | None, components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec]: pixel_values = named_tensors["pixel_values"] pixel_position_ids = named_tensors.get("pixel_position_ids") @@ -3467,7 +3566,7 @@ def _require(component: str) -> None: graph_meta={**common_graph_meta, "component": "lm_encoder"}, metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, )) - cache_seq_len = max(1024, int(input_ids.shape[1]) + 256) + cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=256) prefill_chunk_size = max(1, int(os.environ.get("CACTUS_GEMMA4_PREFILL_CHUNK", "128") or "128")) prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) if "decoder" in expanded_components: @@ -4177,12 +4276,99 @@ def get_transpile_metadata(self): } +class Qwen3LMEncoderStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.backbone = model.model + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + return inputs_embeds, position_ids.to(dtype=torch.int64) + + def get_transpile_metadata(self): + return { + "graph": { + **_transpile_graph_meta( + self.backbone, + adapter_family="qwen3", + adapter_type=type(self).__name__, + input_names=("input_ids", "position_ids"), + ), + } + } + + +class Qwen3LMEncoderTextChunkAdapter(Qwen3LMEncoderStepAdapter): + pass + + +class Qwen3EmbedsCausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + self.backbone = model.model + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + seq_len = int(inputs_embeds.shape[1]) + allowed_positions = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + ).view(1, 1, seq_len, seq_len) + allowed_values = torch.zeros( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + blocked_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * torch.finfo(inputs_embeds.dtype).min + causal_mask = torch.where(allowed_positions, allowed_values, blocked_values) + + hidden_states = inputs_embeds + text_position_ids = position_ids.to(dtype=torch.int64) + position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) + for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=text_position_ids, + past_key_values=None, + use_cache=False, + ) + hidden_states = self.backbone.norm(hidden_states) + return self.model.lm_head(hidden_states[:, -1:, :]) + + def get_transpile_metadata(self): + sliding_window = getattr(self.backbone.config, "sliding_window", None) + layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return { + "graph": { + **_transpile_graph_meta( + self.model, + adapter_family="qwen3", + adapter_type=type(self).__name__, + input_names=("inputs_embeds", "position_ids"), + ), + "num_hidden_layers": int(self.backbone.config.num_hidden_layers), + "layer_types": tuple(layer_types), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + +class Qwen3EmbedsCausalLMPrefillChunkAdapter(Qwen3EmbedsCausalLMStepAdapter): + pass + + def _build_qwen_causal_lm_component_specs( model: torch.nn.Module, *, named_tensors: dict[str, torch.Tensor], weights_dir: str | None = None, components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec] | None: input_ids = named_tensors.get("input_ids") if input_ids is None: @@ -4205,6 +4391,14 @@ def _build_qwen_causal_lm_component_specs( "task": "causal_lm_logits", "adapter_family": family, } + + chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_media_step", "decoder_prefill_chunk"} + wants_chunked = bool(chunk_components & requested_set) + if wants_chunked and family != "qwen3": + raise RuntimeError( + f"text chunked-prefill components are only supported for the qwen3 family, got {family}" + ) + specs: list[ComponentModuleSpec] = [] if "decoder" in requested_set: specs.append(ComponentModuleSpec( @@ -4219,7 +4413,7 @@ def _build_qwen_causal_lm_component_specs( if "decoder_step" in requested_set: step_input_ids = input_ids[:, :1] step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) specs.append(ComponentModuleSpec( component="decoder_step", module=decoder_step, @@ -4237,6 +4431,85 @@ def _build_qwen_causal_lm_component_specs( }, metadata={"family": family, "task": "causal_lm_logits"}, )) + + if wants_chunked: + lm_encoder_step = Qwen3LMEncoderStepAdapter(model).eval() + lm_encoder_text_chunk = Qwen3LMEncoderTextChunkAdapter(model).eval() + decoder_media_step = Qwen3EmbedsCausalLMStepAdapter(model).eval() + decoder_prefill_chunk = Qwen3EmbedsCausalLMPrefillChunkAdapter(model).eval() + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) + prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) + chunk_input_ids = input_ids[:, :prefill_chunk_size].contiguous() + chunk_position_ids = torch.arange( + prefill_chunk_size, + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + with torch.no_grad(): + chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + with torch.no_grad(): + step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) + + if "lm_encoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "lm_encoder_text_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_text_chunk"}, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill_chunk, + example_inputs=(chunk_embeds, chunk_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_prefill_chunk", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_media_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_media_step", + module=decoder_media_step, + example_inputs=(step_embeds, step_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_media_step", + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + }, + metadata={"family": family, "task": "causal_lm_logits"}, + )) return specs @@ -4246,6 +4519,7 @@ def _build_qwen3_5_multimodal_component_specs( named_tensors: dict[str, torch.Tensor], weights_dir: str | None, components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec] | None: input_ids = named_tensors.get("input_ids") attention_mask = named_tensors.get("attention_mask") @@ -4330,6 +4604,7 @@ def _require(component: str) -> None: "task": "multimodal_causal_lm_logits", "adapter_family": "qwen3_5", } + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) specs: list[ComponentModuleSpec] = [] if "vision_encoder" in expanded_components: specs.append(ComponentModuleSpec( @@ -4390,7 +4665,6 @@ def _require(component: str) -> None: if "decoder_prefill_chunk" in expanded_components: if prefill_decoder_inputs is None: raise RuntimeError("Qwen3.5 decoder_prefill_chunk spec requires precomputed text chunk decoder inputs") - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) specs.append(ComponentModuleSpec( component="decoder_prefill_chunk", module=decoder_prefill_chunk, @@ -4414,7 +4688,6 @@ def _require(component: str) -> None: raise RuntimeError("Qwen3.5 decoder_media_step spec requires precomputed decoder inputs") step_inputs_embeds = decoder_inputs[0][:, :1, :] step_position_ids = decoder_inputs[2][:, :, :1] - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) specs.append(ComponentModuleSpec( component="decoder_media_step", module=decoder_media_step, @@ -4435,7 +4708,6 @@ def _require(component: str) -> None: if "decoder_step" in expanded_components: step_input_ids = input_ids[:, :1] step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) specs.append(ComponentModuleSpec( component="decoder_step", module=decoder_step, @@ -5011,6 +5283,7 @@ def _build_lfm2_causal_lm_component_specs( named_tensors: dict[str, torch.Tensor], weights_dir: str | None = None, components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec] | None: input_ids = named_tensors.get("input_ids") if input_ids is None: @@ -5040,7 +5313,7 @@ def _build_lfm2_causal_lm_component_specs( if "decoder_step" in requested_set: step_input_ids = input_ids[:, :1] step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) specs.append(ComponentModuleSpec( component="decoder_step", module=decoder_step, @@ -5067,6 +5340,7 @@ def _build_lfm2_vl_multimodal_component_specs( named_tensors: dict[str, torch.Tensor], weights_dir: str | None, components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec]: input_ids = named_tensors["input_ids"] attention_mask = named_tensors["attention_mask"] @@ -5168,6 +5442,7 @@ def _require(component: str) -> None: "task": "multimodal_causal_lm_logits", "adapter_family": "lfm2_vl", } + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) specs: list[ComponentModuleSpec] = [] if "vision_encoder" in expanded_components: specs.append(ComponentModuleSpec( @@ -5260,7 +5535,6 @@ def _require(component: str) -> None: if image_features is None: image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) specs.append(ComponentModuleSpec( component="decoder_prefill_chunk", module=decoder_last_token, @@ -5282,7 +5556,6 @@ def _require(component: str) -> None: if image_features is None: image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) - max_cache_seq_len = max(1024, int(input_ids.shape[1]) + 512) specs.append(ComponentModuleSpec( component="decoder_step", module=decoder_last_token, @@ -5310,6 +5583,7 @@ def build_component_module_specs( weights_dir: str | None = None, inputs_metadata: dict[str, object] | None = None, components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec] | None: family = _family_key(model) if family == "qwen3_5" and task == "multimodal_causal_lm_logits": @@ -5318,6 +5592,7 @@ def build_component_module_specs( named_tensors=named_tensors, weights_dir=weights_dir, components=components, + cache_context_length=cache_context_length, ) if family in {"qwen3", "qwen3_5"} and task == "causal_lm_logits": return _build_qwen_causal_lm_component_specs( @@ -5325,6 +5600,7 @@ def build_component_module_specs( named_tensors=named_tensors, weights_dir=weights_dir, components=components, + cache_context_length=cache_context_length, ) if family == "gemma4" and task == "multimodal_causal_lm_logits": return _build_gemma4_multimodal_component_specs( @@ -5332,6 +5608,7 @@ def build_component_module_specs( named_tensors=named_tensors, weights_dir=weights_dir, components=components, + cache_context_length=cache_context_length, ) if family == "lfm2_vl" and task == "multimodal_causal_lm_logits": return _build_lfm2_vl_multimodal_component_specs( @@ -5339,6 +5616,7 @@ def build_component_module_specs( named_tensors=named_tensors, weights_dir=weights_dir, components=components, + cache_context_length=cache_context_length, ) if family == "lfm2" and task == "causal_lm_logits": return _build_lfm2_causal_lm_component_specs( @@ -5346,6 +5624,7 @@ def build_component_module_specs( named_tensors=named_tensors, weights_dir=weights_dir, components=components, + cache_context_length=cache_context_length, ) if family == "parakeet_tdt" and task == "tdt_transcription": from cactus.transpile.tdt_runtime import build_parakeet_tdt_component_specs diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index 430fd0e7d..c700c7128 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/python/cactus/transpile/optimize_graph.py @@ -20,8 +20,10 @@ from cactus.transpile.fusion import match_rope from cactus.transpile.fusion import match_self_attention_block from cactus.transpile.fusion.common import producer +from cactus.transpile.fusion.common import strip_layout_passthrough from cactus.transpile.fusion.common import strip_passthrough from cactus.transpile.fusion.linear import match_linear +from cactus.transpile.fusion.rope import _extract_rope_angle_source, _find_constant_ancestor from cactus.transpile.graph_ir import IRGraph from cactus.transpile.graph_ir import IRNode from cactus.transpile.graph_ir import IRValue @@ -55,7 +57,8 @@ class FusionConfig: enable_dense_mlp_tq_fused: bool = True -def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None) -> IRGraph: +def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None, + precompute_rope: bool = True) -> IRGraph: config = config or FusionConfig() verify_ir(graph) canonicalize_exported_graph(graph) @@ -106,6 +109,9 @@ def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig break canonicalize_exported_graph(graph) + if precompute_rope and precompute_rope_tables(graph): + canonicalize_exported_graph(graph) + annotate_gold_patterns(graph) _prune_unused_inputs(graph) verify_ir(graph) @@ -361,6 +367,122 @@ def fuse_rope(graph: IRGraph) -> bool: return changed +def precompute_rope_tables(graph: IRGraph) -> bool: + """Replace the runtime cos/sin(position * inv_freq) rope angle with an fp64-precomputed fp16 table + gathered by position id, fixing fp16's inability to represent angles past ~2048 in long context.""" + + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"scalar_cos", "scalar_sin"} or len(node.inputs) != 1: + continue + # Only the text decoder's rope. Vision/audio encoders carry cos/sin with + # the same shape but an axial/2D position this scalar table cannot model. + if node.meta.get("component") != "decoder": + continue + + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + if output_value is None or output_value.shape is None: + continue + index_shape = tuple(int(dim) for dim in output_value.shape[:-1]) + + match = _match_rope_angle_table_source(graph, node.inputs[0], index_shape) + if match is None: + continue + inv_freq_value_id, position_value_id = match + + inv_freq = graph.constants[inv_freq_value_id].detach().cpu().to(torch.float64).reshape(-1) + head_dim = int(inv_freq.numel()) * 2 + max_seq = _rope_table_max_seq(graph) + positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) + freqs = positions * inv_freq.reshape(1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + table = (torch.cos(emb) if node.op == "scalar_cos" else torch.sin(emb)).to(torch.float16) + assert table.shape[-1] == head_dim + + table_value_id = _materialize_rope_table_constant(graph, node.id, table) + output_value.dtype = "fp16" + node.op = "embedding" + node.inputs = [table_value_id, position_value_id] + node.attrs = {} + node.kind = "generic" + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +def _match_rope_angle_table_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> tuple[str, str] | None: + cat_node = producer(graph, strip_passthrough(graph, value_id)) + if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: + return None + if strip_layout_passthrough(graph, cat_node.inputs[0]) != strip_layout_passthrough(graph, cat_node.inputs[1]): + return None + angle_info = _extract_rope_angle_source(graph, cat_node.inputs[0]) + if angle_info is None: + return None + angle_node = angle_info["matmul_node"] + if len(angle_node.inputs) != 2: + return None + + inv_freq_value_id: str | None = None + position_value_id: str | None = None + for input_id in angle_node.inputs: + const_id = _find_constant_ancestor(graph, input_id) + if const_id is not None: + inv_freq_value_id = const_id + continue + position_value_id = _trace_integer_position_source(graph, input_id, index_shape) + + if inv_freq_value_id is None or position_value_id is None: + return None + return inv_freq_value_id, position_value_id + + +def _trace_integer_position_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> str | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + value = graph.values.get(current) + if ( + value is not None + and str(value.dtype) in {"int32", "int64", "i32", "i64"} + and value.shape is not None + and tuple(int(dim) for dim in value.shape) == index_shape + ): + return current + node = producer(graph, current) + if node is None or node.op not in {"precision_cast", "view", "reshape", "unsqueeze", "squeeze", "contiguous"} or len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _rope_table_max_seq(graph: IRGraph) -> int: + max_seq = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) + if max_seq is None or max_seq <= 0: + max_seq = _coerce_optional_int(graph.meta.get("max_position_embeddings")) + if max_seq is None or max_seq <= 0: + raise ValueError("rope table precompute requires max_cache_seq_len or max_position_embeddings in graph.meta") + return int(max_seq) + + +def _materialize_rope_table_constant(graph: IRGraph, node_id: str, table: torch.Tensor) -> str: + value_id = f"c_rope_table_{node_id}" + graph.constants[value_id] = table + graph.values[value_id] = IRValue( + id=value_id, + shape=tuple(int(dim) for dim in table.shape), + dtype=dtype_to_ir(table.dtype), + producer=None, + users=[], + ) + return value_id + + def fuse_attention(graph: IRGraph) -> bool: changed = False for node_id in list(graph.order): @@ -701,6 +823,9 @@ def _assign_gemma4_decoder_attention_hints_from_graph_meta(graph: IRGraph) -> bo if len(attention_nodes) != len(layer_types): return False + use_internal_cache = bool(graph.meta.get("use_internal_kv_cache", False)) + sliding_window = _graph_sliding_window(graph) + graph_cache_seq_len = _graph_max_cache_seq_len(graph) changed = False for layer_index, (node, layer_type) in enumerate(zip(attention_nodes, layer_types, strict=True)): if "attention_layer_type" not in node.meta: @@ -709,6 +834,16 @@ def _assign_gemma4_decoder_attention_hints_from_graph_meta(graph: IRGraph) -> bo if "attention_layer_index" not in node.meta: node.meta["attention_layer_index"] = int(layer_index) changed = True + if use_internal_cache: + normalized_layer_type = str(layer_type).strip().lower() + node_cache_seq_len = None + if normalized_layer_type in {"sliding", "sliding_attention"} and sliding_window is not None and sliding_window > 0: + node_cache_seq_len = int(sliding_window) + elif normalized_layer_type in {"global", "full", "full_attention"} and graph_cache_seq_len is not None: + node_cache_seq_len = int(graph_cache_seq_len) + if node_cache_seq_len is not None and node.meta.get("max_cache_seq_len") != node_cache_seq_len: + node.meta["max_cache_seq_len"] = node_cache_seq_len + changed = True return changed @@ -1449,6 +1584,22 @@ def _graph_sliding_window(graph: IRGraph) -> int | None: return None +def _graph_max_cache_seq_len(graph: IRGraph) -> int | None: + max_cache_seq_len = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) + if max_cache_seq_len is not None: + return max_cache_seq_len + + providers = graph.meta.get("transpile_metadata_providers") + if isinstance(providers, dict): + for provider_meta in providers.values(): + if not isinstance(provider_meta, dict): + continue + max_cache_seq_len = _coerce_optional_int(provider_meta.get("max_cache_seq_len")) + if max_cache_seq_len is not None: + return max_cache_seq_len + return None + + def _coerce_string_tuple(value: object) -> tuple[str, ...]: if not isinstance(value, (tuple, list)): return () diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index 2db25d2d1..92f5ecfe4 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -244,6 +244,7 @@ def test_defaults(self): assert opts.system_prompt is None assert opts.trust_remote_code is False assert opts.local_files_only is False + assert opts.cache_context_length is None def test_custom_values(self): from cactus.cli.model import TranspileOptions @@ -252,11 +253,13 @@ def test_custom_values(self): prompt="Hello", max_new_tokens=256, trust_remote_code=True, + cache_context_length="131072", ) assert opts.task == "causal_lm_logits" assert opts.prompt == "Hello" assert opts.max_new_tokens == 256 assert opts.trust_remote_code is True + assert opts.cache_context_length == "131072" def test_frozen(self): from cactus.cli.model import TranspileOptions @@ -297,6 +300,7 @@ def test_run_command_chunked_bundle_flags(self): "--audio", "audio.wav", "--image", "image.png", "--input-ids", "1,2,3", + "--input-ids-file", "tokens.txt", "--max-new-tokens", "4", "--result-json", "result.json", ]) @@ -304,6 +308,7 @@ def test_run_command_chunked_bundle_flags(self): assert args.audio == "audio.wav" assert args.image == "image.png" assert args.input_ids == "1,2,3" + assert args.input_ids_file == "tokens.txt" assert args.max_new_tokens == 4 assert args.result_json == "result.json" diff --git a/python/tests/test_cli_run.py b/python/tests/test_cli_run.py index f63772dfc..f3722d138 100644 --- a/python/tests/test_cli_run.py +++ b/python/tests/test_cli_run.py @@ -20,8 +20,10 @@ def test_cmd_run_forwards_chunked_bundle_flags(monkeypatch, tmp_path: Path) -> N image_file = tmp_path / "image.png" audio_file = tmp_path / "audio.wav" result_json = tmp_path / "result.json" + input_ids_file = tmp_path / "tokens.txt" image_file.write_bytes(b"image") audio_file.write_bytes(b"audio") + input_ids_file.write_text("1 2 3", encoding="utf-8") calls = [] @@ -44,6 +46,7 @@ def fake_subprocess_run(cmd): image=str(image_file), audio=str(audio_file), input_ids="1,2,3", + input_ids_file=str(input_ids_file), max_new_tokens=4, result_json=str(result_json), confidence_threshold=None, @@ -60,5 +63,6 @@ def fake_subprocess_run(cmd): assert cmd[cmd.index("--image") + 1] == str(image_file) assert cmd[cmd.index("--audio") + 1] == str(audio_file) assert cmd[cmd.index("--input-ids") + 1] == "1,2,3" + assert cmd[cmd.index("--input-ids-file") + 1] == str(input_ids_file) assert cmd[cmd.index("--max-new-tokens") + 1] == "4" assert cmd[cmd.index("--result-json") + 1] == str(result_json) diff --git a/python/tests/test_cli_transpile_defaults.py b/python/tests/test_cli_transpile_defaults.py index e3c7b78c9..fe9021911 100644 --- a/python/tests/test_cli_transpile_defaults.py +++ b/python/tests/test_cli_transpile_defaults.py @@ -1,9 +1,71 @@ from __future__ import annotations from pathlib import Path +from types import SimpleNamespace +import pytest from cactus import cli from cactus.cli import convert as convert_cli +from cactus.transpile.model_adapters import _cache_context_length + + +class _ConfigWithText: + def get_text_config(self): + return SimpleNamespace(max_position_embeddings=128000) + + +def test_cache_context_length_uses_explicit_value() -> None: + model = SimpleNamespace(config=SimpleNamespace(max_position_embeddings=40960)) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length="32768", + fallback_extra_tokens=512, + ) == 32768 + + +def test_cache_context_length_reads_top_level_config() -> None: + model = SimpleNamespace(config=SimpleNamespace(max_position_embeddings=40960)) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length=None, + fallback_extra_tokens=512, + ) == 40960 + + +def test_cache_context_length_reads_text_config() -> None: + model = SimpleNamespace(config=_ConfigWithText()) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length="auto", + fallback_extra_tokens=512, + ) == 128000 + + +def test_cache_context_length_falls_back_to_capture_size() -> None: + model = SimpleNamespace(config=SimpleNamespace()) + + assert _cache_context_length( + model, + input_seq_len=2048, + cache_context_length=None, + fallback_extra_tokens=512, + ) == 2560 + + +def test_cache_context_length_rejects_non_positive_explicit_value() -> None: + with pytest.raises(ValueError): + _cache_context_length( + SimpleNamespace(config=SimpleNamespace()), + input_seq_len=2048, + cache_context_length="0", + fallback_extra_tokens=512, + ) def test_cmd_convert_does_not_transpile(monkeypatch, tmp_path: Path) -> None: diff --git a/python/tests/test_optimize_gemma4_attention.py b/python/tests/test_optimize_gemma4_attention.py index 843e3f65c..36fddc19d 100644 --- a/python/tests/test_optimize_gemma4_attention.py +++ b/python/tests/test_optimize_gemma4_attention.py @@ -1,7 +1,11 @@ +import torch + +from cactus.transpile.canonicalize.utils import rebuild_graph from cactus.transpile.graph_ir import IRGraph from cactus.transpile.graph_ir import IRNode from cactus.transpile.graph_ir import IRValue from cactus.transpile.optimize_graph import normalize_gemma4_decoder_attention_semantics +from cactus.transpile.optimize_graph import precompute_rope_tables def test_normalize_gemma4_full_attention_uses_sequence_window_compat() -> None: @@ -178,6 +182,57 @@ def test_normalize_gemma4_attention_recovers_layer_hints_from_graph_meta() -> No assert graph.nodes["attn_1"].attrs["window_size"] == 800 +def test_normalize_gemma4_cached_attention_sets_per_layer_cache_lengths() -> None: + graph = IRGraph( + values={ + "query_0": IRValue(id="query_0", shape=(1, 8, 1, 256), dtype="fp16"), + "key_0": IRValue(id="key_0", shape=(1, 8, 1, 256), dtype="fp16"), + "value_0": IRValue(id="value_0", shape=(1, 8, 1, 256), dtype="fp16"), + "out_0": IRValue(id="out_0", shape=(1, 8, 1, 256), dtype="fp16", producer="attn_0"), + "query_1": IRValue(id="query_1", shape=(1, 8, 1, 256), dtype="fp16"), + "key_1": IRValue(id="key_1", shape=(1, 8, 1, 256), dtype="fp16"), + "value_1": IRValue(id="value_1", shape=(1, 8, 1, 256), dtype="fp16"), + "out_1": IRValue(id="out_1", shape=(1, 8, 1, 256), dtype="fp16", producer="attn_1"), + }, + nodes={ + "attn_0": IRNode( + id="attn_0", + op="attention", + inputs=["query_0", "key_0", "value_0"], + outputs=["out_0"], + attrs={"is_causal": False, "window_size": 0}, + ), + "attn_1": IRNode( + id="attn_1", + op="attention", + inputs=["query_1", "key_1", "value_1"], + outputs=["out_1"], + attrs={"is_causal": False, "window_size": 0}, + ), + }, + order=["attn_0", "attn_1"], + inputs=["query_0", "key_0", "value_0", "query_1", "key_1", "value_1"], + outputs=["out_1"], + constants={}, + meta={ + "adapter_family": "gemma4", + "component": "decoder_step", + "use_internal_kv_cache": True, + "sliding_window": 512, + "max_cache_seq_len": 131072, + "layer_types": ("sliding_attention", "full_attention"), + }, + ) + + changed = normalize_gemma4_decoder_attention_semantics(graph) + + assert changed is True + assert graph.nodes["attn_0"].attrs["window_size"] == 512 + assert graph.nodes["attn_0"].meta["max_cache_seq_len"] == 512 + assert graph.nodes["attn_1"].attrs["window_size"] == 0 + assert graph.nodes["attn_1"].meta["max_cache_seq_len"] == 131072 + + def test_normalize_gemma4_attention_elides_logical_fp16_mask() -> None: graph = IRGraph( values={ @@ -222,3 +277,52 @@ def test_normalize_gemma4_attention_elides_logical_fp16_mask() -> None: assert graph.nodes["attn"].inputs == ["query", "key", "value"] assert graph.nodes["attn"].attrs["is_causal"] is True assert graph.nodes["attn"].attrs["window_size"] == 512 + + +def test_precompute_rope_tables_replaces_runtime_angle_with_fp16_lookup() -> None: + head_dim = 8 + max_seq = 4096 + theta = 10000.0 + inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + + graph = IRGraph( + values={ + "positions": IRValue(id="positions", shape=(1, 4), dtype="int32"), + "angle": IRValue(id="angle", shape=(1, 4, head_dim // 2), dtype="fp32", producer="matmul"), + "emb": IRValue(id="emb", shape=(1, 4, head_dim), dtype="fp32", producer="cat"), + "cos": IRValue(id="cos", shape=(1, 4, head_dim), dtype="fp32", producer="cos_node"), + }, + nodes={ + "matmul": IRNode(id="matmul", op="matmul", inputs=["positions", "inv_freq"], outputs=["angle"]), + "cat": IRNode(id="cat", op="cat", inputs=["angle", "angle"], outputs=["emb"], attrs={"dim": -1}), + "cos_node": IRNode( + id="cos_node", op="scalar_cos", inputs=["emb"], outputs=["cos"], + meta={"component": "decoder"}, + ), + }, + order=["matmul", "cat", "cos_node"], + inputs=["positions"], + outputs=["cos"], + constants={"inv_freq": inv_freq}, + meta={"component": "decoder", "max_cache_seq_len": max_seq}, + ) + rebuild_graph(graph) + + assert precompute_rope_tables(graph) is True + + cos_node = graph.nodes["cos_node"] + assert cos_node.op == "embedding" + table_id, position_id = cos_node.inputs + assert table_id.startswith("c_rope_table_") + assert position_id == "positions" + + table = graph.constants[table_id] + assert table.dtype == torch.float16 + assert tuple(table.shape) == (max_seq, head_dim) + assert graph.values["cos"].dtype == "fp16" + + # Positions past the fp16-angle range (>2048) stay distinct/representable in the table. + positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) + freqs = positions * inv_freq.to(torch.float64).reshape(1, -1) + expected = torch.cos(torch.cat((freqs, freqs), dim=-1)).to(torch.float16) + assert torch.equal(table[[0, 1, 2049, max_seq - 1]], expected[[0, 1, 2049, max_seq - 1]])