diff --git a/kv_cache_manager/optimizer/docs/optimizer_architecture.md b/kv_cache_manager/optimizer/docs/optimizer_architecture.md index b07ab9a14..25280448f 100644 --- a/kv_cache_manager/optimizer/docs/optimizer_architecture.md +++ b/kv_cache_manager/optimizer/docs/optimizer_architecture.md @@ -43,7 +43,7 @@ KVCacheManager Optimizer is a standalone cache optimization analysis module. It The Optimizer currently contains two execution paths: offline trace replay and online optimization service: - The offline replay path uses `OptimizerManager`, `OptIndexerManager`, `OptEvictionManager`, and `RadixTreeIndex`, depending on replay configuration, trace loader/converter, and data storage types. -- The online service path uses `OnlineOptimizerManager`, LiteHit, `CacheIndexerFactory`, `LruCacheIndexer`/`TtlCacheIndexerWrapper`, and service/protobuf interfaces, depending on online instance group/instance configuration and registry. Full-attention multi-capacity LRU is directly held by `InstanceState` as LiteHit; linear attention is still simulated by the generic `CacheIndexer`. +- The online service path uses `OnlineOptimizerManager`, the unified LiteHit runtime, and service/protobuf interfaces, depending on online instance group/instance configuration and registry. Each `InstanceState` holds one `TtlLiteHit` decorator with a single `LiteHit` core inside; full attention uses the Full-only policy, and linear attention additionally enables the Linear state policy. The two paths share the hit-rate modeling code within the optimizer module, but the runtime, configuration targets, and service boundaries remain independent, avoiding the online service introducing offline replay/data storage dependencies. @@ -78,9 +78,10 @@ OnlineOptimizerManager (online runtime coordinator) ├── OptimizerRegistryManager (instance group/instance persistence) └── InstanceState (instance_id-isolated runtime state) ↓ - ├── LiteHit (full-attention multi-capacity LRU) - └── CacheIndexerFactory (linear attention) - └── LruCacheIndexer / TtlCacheIndexerWrapper + └── LiteHit (unified weighted LRU core) + ├── Full-only (optional block-RLE) + └── Linear state policy (mixed Full + Linear charges) + both can layer a fixed TTL ``` ### Directory Structure @@ -97,10 +98,12 @@ kv_cache_manager/optimizer/ │ └── online_runtime/ # online runtime │ └── online_optimizer_manager.h/cc # online instance registration, TraceQuery, and statistics ├── index/ # index layer -│ ├── radix_tree_index.h/cc # offline Radix tree index -│ └── online/ # linear attention online capacity/TTL indexer -├── liteHit/ # lightweight multi-capacity full-attention LRU hit-rate core -│ ├── lite_hit.h/cc # multi-capacity LRU hit-rate analyzer +│ └── radix_tree_index.h/cc # offline Radix tree index +├── liteHit/ # lightweight multi-capacity hit-rate core +│ ├── lite_hit.h/cc # capacity-independent weighted LRU core shared by Full/Mamba +│ ├── lite_hit_linear.h/cc # Linear state restore points and write policy (no separate LRU state) +│ ├── lite_hit_ttl.h/cc # fixed-TTL decorator wrapping LiteHit (epochs/watermark) +│ ├── weighted_lru_pool.h/cc # byte-weighted LRU pool (typed keys: Full / Linear) │ └── dynamic_fenwick_tree.h/cc # order-statistics Fenwick for reuse-distance ├── eviction_policy/ # eviction policy layer │ ├── base.h # policy base class @@ -156,7 +159,9 @@ kv_cache_manager/optimizer/ The online service protocol is defined in `kv_cache_manager/protocol/protobuf/optimizer_service.proto`, and is converted by the service layer into optimizer online config/runtime objects. -A full-attention TraceQuery only puts complete blocks into `block_keys`, and passes the original input length including trailing tokens via `input_token_len`. The Online Manager uses the fixed byte charge of the full location spec group to floor `capacity_gb` to block capacity, then hands the same request to LiteHit; both per-request and cumulative hit rates are `prefix_hit_blocks * block_size_tokens / input_tokens`. When old clients lack the length, the compatibility assumption is that there are no trailing tokens. When a full-attention group config has `ttl_seconds != 0`, a fixed TTL is layered on top of LiteHit with wall-clock time (metrics consistent with `TtlCacheIndexerWrapper`); the statistics metrics and TTL behavior of linear attention keep the legacy path unchanged. +A full-attention TraceQuery only puts complete blocks into `block_keys`, and passes the original input length including trailing tokens via `input_token_len`. Like linear attention, the Online Manager keeps only byte capacities; the core's default byte-step can be losslessly compressed into Full-only block RLE, and the fixed byte charge of the full location spec group performs a single floor only at final projection. Both per-request and cumulative hit rates are `prefix_hit_blocks * block_size_tokens / input_tokens`. When old clients lack the length, the compatibility assumption is that there are no trailing tokens. + +Linear attention (`linear_step > 0`) enables the Linear state policy on the same LiteHit: a Full block and a Linear state are two independent objects that share **one recency order and one total byte budget**; the capacity axis is total bytes rather than a block count, and the hit semantics is "resume from a Linear state". `LiteHit` itself has no notion of time; the outer `TtlLiteHit` decorator filters both Full and Linear states with a shared epoch/watermark, expiring an object from the weighted LRU's visible set once its age reaches the group TTL. Online uses wall-clock time, offline uses trace timestamps, and with TTL 0 the decorator is a transparent pass-through. --- diff --git a/kv_cache_manager/optimizer/docs/optimizer_architecture_zh.md b/kv_cache_manager/optimizer/docs/optimizer_architecture_zh.md index 4419e3764..e84893fa6 100644 --- a/kv_cache_manager/optimizer/docs/optimizer_architecture_zh.md +++ b/kv_cache_manager/optimizer/docs/optimizer_architecture_zh.md @@ -43,7 +43,7 @@ KVCacheManager Optimizer 是一个独立的缓存优化分析模块，通过回 Optimizer 当前包含离线 trace 回放和在线优化服务两条运行路径： - 离线回放路径使用 `OptimizerManager`、`OptIndexerManager`、`OptEvictionManager` 和 `RadixTreeIndex`，依赖 replay 配置、trace loader/converter 和 data storage 类型。 -- 在线服务路径使用 `OnlineOptimizerManager`、LiteHit、`CacheIndexerFactory`、`LruCacheIndexer`/`TtlCacheIndexerWrapper` 和 service/protobuf 接口，依赖 online 实例组/实例配置与 registry。full-attention 多容量 LRU 由 `InstanceState` 直接持有 LiteHit；linear attention 仍由通用 `CacheIndexer` 模拟。 +- 在线服务路径使用 `OnlineOptimizerManager`、统一的 LiteHit 运行时和 service/protobuf 接口，依赖 online 实例组/实例配置与 registry。每个 `InstanceState` 持有一个 `TtlLiteHit` decorator，内部只有一个 `LiteHit` core；full-attention 使用 Full-only 策略，linear attention 额外启用 Linear state 策略。 两条路径共享 optimizer 模块内的命中率建模代码，但运行时、配置 target 和服务边界保持独立，避免在线服务引入离线 replay/data storage 依赖。 @@ -78,9 +78,10 @@ OnlineOptimizerManager (在线运行时协调器) ├── OptimizerRegistryManager (实例组/实例持久化) └── InstanceState (instance_id 隔离的运行状态) ↓ - ├── LiteHit (full-attention 多容量 LRU) - └── CacheIndexerFactory (linear attention) - └── LruCacheIndexer / TtlCacheIndexerWrapper + └── LiteHit (统一 weighted LRU core) + ├── Full-only（可选 block-RLE） + └── Linear state policy（Full + Linear 混合 charge） + 两者均可叠固定 TTL ``` ### 目录结构 @@ -97,10 +98,12 @@ kv_cache_manager/optimizer/ │ └── online_runtime/ # 在线运行时 │ └── online_optimizer_manager.h/cc # 在线实例注册、TraceQuery 和统计 ├── index/ # 索引层 -│ ├── radix_tree_index.h/cc # 离线 Radix 树索引 -│ └── online/ # linear attention 在线容量/TTL 索引器 -├── liteHit/ # 轻量多容量 full-attention LRU 命中率核心 -│ ├── lite_hit.h/cc # 多容量 LRU 命中率分析器 +│ └── radix_tree_index.h/cc # 离线 Radix 树索引 +├── liteHit/ # 轻量多容量命中率核心 +│ ├── lite_hit.h/cc # Full/Mamba 共享的容量无关 weighted LRU 核心 +│ ├── lite_hit_linear.h/cc # Linear state 恢复点与写入策略（无独立 LRU 状态） +│ ├── lite_hit_ttl.h/cc # 包裹 LiteHit 的固定 TTL decorator（epoch/watermark） +│ ├── weighted_lru_pool.h/cc # 字节加权 LRU 池（typed key：Full / Linear） │ └── dynamic_fenwick_tree.h/cc # reuse-distance 用的 order-statistics Fenwick ├── eviction_policy/ # 驱逐策略层 │ ├── base.h # 策略基类 @@ -156,7 +159,9 @@ kv_cache_manager/optimizer/ 在线服务协议定义在 `kv_cache_manager/protocol/protobuf/optimizer_service.proto`，由 service 层转换为 optimizer online config/runtime 对象。 -full-attention TraceQuery 只把完整 block 放入 `block_keys`，并用 `input_token_len` 传入包含尾部 token 的原始输入长度。Online Manager 用 full location spec group 的固定字节 charge 将 `capacity_gb` 向下取整为 block 容量，再把同一请求交给 LiteHit；请求级与累计命中率均为 `prefix_hit_blocks * block_size_tokens / input_tokens`。旧客户端缺少长度时兼容假设没有尾部 token。full-attention 组配置 `ttl_seconds != 0` 时，固定 TTL 以墙钟时间叠加在 LiteHit 之上（口径与 `TtlCacheIndexerWrapper` 一致）；linear attention 的统计口径和 TTL 行为保持 legacy 路径不变。 +full-attention TraceQuery 只把完整 block 放入 `block_keys`，并用 `input_token_len` 传入包含尾部 token 的原始输入长度。Online Manager 与 linear attention 一样只保存 byte 容量；核心默认 byte-step 可无损压成 Full-only block-RLE，最终投影时再用 full location spec group 的固定字节 charge 做一次 floor。请求级与累计命中率均为 `prefix_hit_blocks * block_size_tokens / input_tokens`。旧客户端缺少长度时兼容假设没有尾部 token。 + +linear attention（`linear_step > 0`）在同一个 LiteHit 上启用 Linear state 策略：Full block 与 Linear state 是两个独立对象，但共享**同一个 recency 和总字节预算**；容量轴是总字节而非 block 数，命中语义是“从某个 Linear state 恢复”。`LiteHit` 本身没有时间概念；外层 `TtlLiteHit` decorator 用共享 epoch/watermark 同时过滤 Full 与 Linear state，age 达到组 TTL 时从 weighted LRU 可见集合中失效。在线使用墙钟，离线使用 trace 时间戳，TTL 为 0 时 decorator 透明透传。 --- diff --git a/kv_cache_manager/optimizer/index/BUILD b/kv_cache_manager/optimizer/index/BUILD index dc89b6943..54a8a73b8 100644 --- a/kv_cache_manager/optimizer/index/BUILD +++ b/kv_cache_manager/optimizer/index/BUILD @@ -16,48 +16,3 @@ cc_library( "//kv_cache_manager/optimizer/analysis", ], ) - -cc_library( - name = "cache_indexer", - srcs = [ - "online/lru_cache_indexer.cc", - ], - hdrs = [ - "online/cache_indexer.h", - "online/lru_cache_indexer.h", - ], - deps = [ - "//kv_cache_manager/common:logger", - "//kv_cache_manager/common/cache", - ], -) - -cc_library( - name = "ttl_cache_indexer_wrapper", - srcs = [ - "online/ttl_cache_indexer_wrapper.cc", - ], - hdrs = [ - "online/ttl_cache_indexer_wrapper.h", - ], - deps = [ - ":cache_indexer", - "//kv_cache_manager/common:timestamp_util", - ], -) - -cc_library( - name = "cache_indexer_factory", - srcs = [ - "online/cache_indexer_factory.cc", - ], - hdrs = [ - "online/cache_indexer_factory.h", - ], - deps = [ - ":cache_indexer", - ":ttl_cache_indexer_wrapper", - "//kv_cache_manager/common:env_util", - "//kv_cache_manager/common:logger", - ], -) diff --git a/kv_cache_manager/optimizer/index/online/cache_indexer.h b/kv_cache_manager/optimizer/index/online/cache_indexer.h deleted file mode 100644 index d3b8b2c80..000000000 --- a/kv_cache_manager/optimizer/index/online/cache_indexer.h +++ /dev/null @@ -1,79 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -namespace kv_cache_manager { - -// Per-bucket hit count for cache age distribution. -// Each bucket covers hits whose age (now - last_access_time) falls within -// [0, threshold[0]), [threshold[0], threshold[1]), ..., [threshold[N-1], +inf). -struct HitAgeBucketInfo { - int64_t threshold_seconds; // upper bound of this bucket (0 means "+inf") - int64_t hit_count; -}; - -class CacheIndexer { -public: - virtual ~CacheIndexer() = default; - - CacheIndexer() = default; - CacheIndexer(const CacheIndexer &) = delete; - CacheIndexer &operator=(const CacheIndexer &) = delete; - - // Initialize the indexer with capacity and size parameters. - // capacity_gb: capacity tiers in GB. - // size_full_only: byte size of a full-only block. - // size_full_linear: byte size of a full+linear block. - // linear_step: linear step factor (>=0). - virtual void Init(const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step) = 0; - - // Process a batch of key accesses and compute per-capacity prefix hit count. - // keys: the block keys in query order. - // hit_count: output vector sized to number of capacity tiers, filled with - // the count of contiguous prefix hits for each tier. - // key_hits: optional output, one entry per input key. When provided, it is - // filled with whether the key was resident before this query updated the - // indexer. This is separate from prefix hit_count. - virtual void ProcessKeys(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits = nullptr) = 0; - - virtual int64_t unique_count() const = 0; - - // Number of keys evicted from the indexer. - virtual int64_t eviction_count() const = 0; - - // Estimated memory usage in bytes of internal data structures. - virtual int64_t memory_usage_bytes() const = 0; - - // Estimated total kv cache size in bytes for all keys currently tracked. - virtual int64_t kv_cache_usage_bytes() const = 0; - - // Resident unique key count for each configured capacity. Implementations - // that do not maintain per-capacity state may return an empty vector. - virtual std::vector capacity_unique_counts() const { return {}; } - - // Remove a specific key from the indexer. - // Returns true if the key existed and was removed. - virtual bool RemoveKey(int64_t key) { return false; } - - // Number of keys evicted due to TTL expiration. - virtual int64_t ttl_eviction_count() const { return 0; } - - // Called after processing all keys in a query batch. - // Subclasses may perform eviction, compaction, etc. - virtual void PostQueryMaintenance() {} - - // Return per-bucket hit counts for cache age distribution. - // Default returns empty (no age tracking). - virtual std::vector GetHitAgeBuckets() const { return {}; } -}; - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/index/online/cache_indexer_factory.cc b/kv_cache_manager/optimizer/index/online/cache_indexer_factory.cc deleted file mode 100644 index 6f13655d2..000000000 --- a/kv_cache_manager/optimizer/index/online/cache_indexer_factory.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "kv_cache_manager/optimizer/index/online/cache_indexer_factory.h" - -#include -#include -#include -#include -#include -#include - -#include "kv_cache_manager/common/env_util.h" -#include "kv_cache_manager/common/logger.h" -#include "kv_cache_manager/optimizer/index/online/lru_cache_indexer.h" -#include "kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.h" - -namespace kv_cache_manager { - -static void ApplyHitAgeBucketThresholdsFromEnv(TtlCacheIndexerWrapper *wrapper) { - std::string env_value = EnvUtil::GetEnv("KVCM_HIT_AGE_BUCKET_THRESHOLDS", std::string("")); - if (env_value.empty()) { - return; - } - std::vector thresholds; - std::istringstream stream(env_value); - std::string token; - while (std::getline(stream, token, ',')) { - try { - int64_t value = std::stoll(token); - if (value > 0) { - thresholds.push_back(value); - } - } catch (...) { KVCM_LOG_WARN("Invalid token in KVCM_HIT_AGE_BUCKET_THRESHOLDS: [%s]", token.c_str()); } - } - if (!thresholds.empty()) { - std::sort(thresholds.begin(), thresholds.end()); - wrapper->SetHitAgeBucketThresholds(thresholds); - KVCM_LOG_INFO("Applied custom hit age bucket thresholds from env, count=%zu", thresholds.size()); - } -} - -std::unique_ptr CacheIndexerFactory::CreateCacheIndexer(const std::string &eviction_policy, - bool enable_theoretical_max_cache, - const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step, - int64_t ttl_seconds) { - if (eviction_policy != "lru") { - KVCM_LOG_ERROR("CreateCacheIndexer: unsupported eviction_policy[%s]", eviction_policy.c_str()); - return nullptr; - } - if (ttl_seconds < 0) { - KVCM_LOG_ERROR("CreateCacheIndexer: ttl_seconds must be non-negative[%ld]", ttl_seconds); - return nullptr; - } - - std::unique_ptr indexer = std::make_unique(enable_theoretical_max_cache); - indexer->Init(capacity_gb, size_full_only, size_full_linear, linear_step); - - if (ttl_seconds > 0) { - auto ttl_wrapper = std::make_unique(std::move(indexer), ttl_seconds); - ApplyHitAgeBucketThresholdsFromEnv(ttl_wrapper.get()); - return ttl_wrapper; - } - - return indexer; -} - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/index/online/cache_indexer_factory.h b/kv_cache_manager/optimizer/index/online/cache_indexer_factory.h deleted file mode 100644 index 18e64d5bb..000000000 --- a/kv_cache_manager/optimizer/index/online/cache_indexer_factory.h +++ /dev/null @@ -1,23 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -#include "kv_cache_manager/optimizer/index/online/cache_indexer.h" - -namespace kv_cache_manager { - -class CacheIndexerFactory { -public: - static std::unique_ptr CreateCacheIndexer(const std::string &eviction_policy, - bool enable_theoretical_max_cache, - const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step, - int64_t ttl_seconds = 0); -}; - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/index/online/lru_cache_indexer.cc b/kv_cache_manager/optimizer/index/online/lru_cache_indexer.cc deleted file mode 100644 index 175380293..000000000 --- a/kv_cache_manager/optimizer/index/online/lru_cache_indexer.cc +++ /dev/null @@ -1,290 +0,0 @@ -#include "kv_cache_manager/optimizer/index/online/lru_cache_indexer.h" - -#include -#include -#include - -#include "kv_cache_manager/common/cache/cache.h" - -namespace kv_cache_manager { - -const Cache::CacheItemHelper LruCacheIndexer::kHelper(CacheEntryRole::kMisc, nullptr); - -LruCacheIndexer::LruCacheIndexer(bool enable_theoretical_max_cache) - : enable_theoretical_max_cache_(enable_theoretical_max_cache) {} - -void LruCacheIndexer::Init(const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step) { - size_full_only_ = size_full_only; - size_full_linear_ = size_full_linear; - linear_step_ = std::max(linear_step, int32_t(0)); - - capacity_bytes_.resize(capacity_gb.size()); - for (size_t i = 0; i < capacity_gb.size(); i++) { - capacity_bytes_[i] = static_cast(capacity_gb[i] * 1024.0 * 1024.0 * 1024.0); - } - - RebuildCaches(); - - if (enable_theoretical_max_cache_) { - constexpr size_t kTheoreticalMaxCacheBytes = std::numeric_limits::max() / 4; - max_cache_ = NewLRUCache(kTheoreticalMaxCacheBytes, - 0, - false, - false, - 0.0, - nullptr, - kDefaultToAdaptiveMutex, - kDontChargeCacheMetadata); - } -} - -void LruCacheIndexer::RebuildCaches() { - caches_.clear(); - caches_.reserve(capacity_bytes_.size()); - for (int64_t cap_bytes : capacity_bytes_) { - auto cache = NewLRUCache(static_cast(cap_bytes), - 0, - false, - false, - 0.0, - nullptr, - kDefaultToAdaptiveMutex, - kDontChargeCacheMetadata); - caches_.push_back(std::move(cache)); - } - max_cache_.reset(); - unique_count_ = 0; - eviction_count_ = 0; -} - -bool LruCacheIndexer::LookupAndInsert(Cache *cache, std::string_view key_sv, bool &is_new_key) { - auto *handle = cache->Lookup(key_sv); - if (handle) { - is_new_key = false; - cache->Release(handle); - return true; - } - cache->Insert(key_sv, nullptr, &kHelper, static_cast(size_full_only_)); - return false; -} - -bool LruCacheIndexer::LookupAndInsert(Cache *cache, - std::string_view key_sv, - bool is_linear, - int64_t desired_charge, - bool &is_new_key, - bool &is_checkpoint) { - is_checkpoint = false; - auto *handle = cache->Lookup(key_sv); - if (handle) { - is_new_key = false; - size_t stored_charge = cache->GetCharge(handle); - cache->Release(handle); - is_checkpoint = (static_cast(stored_charge) == size_full_linear_); - if (is_linear && !is_checkpoint) { - cache->Erase(key_sv); - cache->Insert(key_sv, nullptr, &kHelper, static_cast(desired_charge)); - } - return true; - } - cache->Insert(key_sv, nullptr, &kHelper, static_cast(desired_charge)); - return false; -} - -void LruCacheIndexer::ProcessKeys(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits) { - if (linear_step_ == 0) { - ProcessKeysFullAttention(keys, hit_count, max_hit_count, key_hits); - } else { - ProcessKeysLinearAttention(keys, hit_count, max_hit_count, key_hits); - } -} - -void LruCacheIndexer::ProcessKeysFullAttention(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits) { - const size_t num_caps = caches_.size(); - const int64_t total_keys = static_cast(keys.size()); - hit_count.assign(num_caps, total_keys); - max_hit_count = max_cache_ ? total_keys : -1; - if (key_hits) { - key_hits->assign(keys.size(), false); - } - - for (int64_t i = 0; i < total_keys; i++) { - int64_t key = keys[i]; - std::string_view key_sv(reinterpret_cast(&key), sizeof(key)); - - bool is_new_key = true; - bool largest_cache_hit = false; - for (size_t j = 0; j < num_caps; j++) { - const bool hit = LookupAndInsert(caches_[j].get(), key_sv, is_new_key); - if (j + 1 == num_caps) { - largest_cache_hit = hit; - } - if (!hit) { - if (i < hit_count[j]) { - hit_count[j] = i; - } - } - } - bool max_cache_hit = false; - if (max_cache_) { - max_cache_hit = LookupAndInsert(max_cache_.get(), key_sv, is_new_key); - if (!max_cache_hit) { - if (i < max_hit_count) { - max_hit_count = i; - } - } - } - if (key_hits) { - (*key_hits)[i] = largest_cache_hit; - } - - if (is_new_key) { - unique_count_++; - } - } -} - -void LruCacheIndexer::ProcessKeysLinearAttention(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits) { - const size_t num_caps = caches_.size(); - const int64_t total_keys = static_cast(keys.size()); - std::vector first_miss(num_caps, total_keys); - max_hit_count = max_cache_ ? 0 : -1; - int64_t max_first_miss = total_keys; - if (key_hits) { - key_hits->assign(keys.size(), false); - } - - std::vector last_checkpoint(num_caps, -1); - int64_t max_last_checkpoint = -1; - - for (int64_t i = 0; i < total_keys; i++) { - int64_t key = keys[i]; - std::string_view key_sv(reinterpret_cast(&key), sizeof(key)); - const bool is_linear = (((i + 1) % linear_step_) == 0) || (i == total_keys - 1); - const int64_t desired_charge = is_linear ? size_full_linear_ : size_full_only_; - - bool is_new_key = true; - bool largest_cache_hit = false; - for (size_t j = 0; j < num_caps; j++) { - bool is_checkpoint = false; - const bool hit = - LookupAndInsert(caches_[j].get(), key_sv, is_linear, desired_charge, is_new_key, is_checkpoint); - if (j + 1 == num_caps) { - largest_cache_hit = hit; - } - if (hit) { - if (is_checkpoint && i < first_miss[j]) { - last_checkpoint[j] = i; - } - } else { - if (i < first_miss[j]) { - first_miss[j] = i; - } - } - } - bool max_cache_hit = false; - if (max_cache_) { - bool is_checkpoint = false; - max_cache_hit = - LookupAndInsert(max_cache_.get(), key_sv, is_linear, desired_charge, is_new_key, is_checkpoint); - if (max_cache_hit) { - if (is_checkpoint && i < max_first_miss) { - max_last_checkpoint = i; - } - } else { - if (i < max_first_miss) { - max_first_miss = i; - } - } - } - if (key_hits) { - (*key_hits)[i] = largest_cache_hit; - } - - if (is_new_key) { - unique_count_++; - } - } - - hit_count.assign(num_caps, 0); - for (size_t j = 0; j < num_caps; j++) { - if (last_checkpoint[j] >= 0 && last_checkpoint[j] < first_miss[j]) { - hit_count[j] = last_checkpoint[j] + 1; - } - } - if (max_cache_) { - if (max_last_checkpoint >= 0 && max_last_checkpoint < max_first_miss) { - max_hit_count = max_last_checkpoint + 1; - } - } -} - -void LruCacheIndexer::PostQueryMaintenance() { - Cache *largest = max_cache_ ? max_cache_.get() : (caches_.empty() ? nullptr : caches_.back().get()); - if (largest) { - int64_t occupancy = static_cast(largest->GetOccupancyCount()); - if (unique_count_ > occupancy) { - eviction_count_ += (unique_count_ - occupancy); - unique_count_ = occupancy; - } - } -} - -bool LruCacheIndexer::RemoveKey(int64_t key) { - std::string_view key_sv(reinterpret_cast(&key), sizeof(key)); - bool found = false; - for (auto &cache : caches_) { - if (cache->Erase(key_sv)) - found = true; - } - if (max_cache_ && max_cache_->Erase(key_sv)) - found = true; - if (found) { - unique_count_--; - eviction_count_++; - } - return found; -} - -int64_t LruCacheIndexer::kv_cache_usage_bytes() const { - if (max_cache_) { - return static_cast(max_cache_->GetUsage()); - } - if (caches_.empty()) - return 0; - return static_cast(caches_.back()->GetUsage()); -} - -int64_t LruCacheIndexer::memory_usage_bytes() const { - int64_t total = 0; - for (const auto &cache : caches_) { - total += static_cast(cache->GetOccupancyCount()) * kEstimatedCacheEntryOverheadBytes; - } - if (max_cache_) { - total += static_cast(max_cache_->GetOccupancyCount()) * kEstimatedCacheEntryOverheadBytes; - } - return total; -} - -std::vector LruCacheIndexer::capacity_unique_counts() const { - std::vector result; - result.reserve(caches_.size()); - for (const auto &cache : caches_) { - result.push_back(static_cast(cache->GetOccupancyCount())); - } - return result; -} - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/index/online/lru_cache_indexer.h b/kv_cache_manager/optimizer/index/online/lru_cache_indexer.h deleted file mode 100644 index 44fdc2162..000000000 --- a/kv_cache_manager/optimizer/index/online/lru_cache_indexer.h +++ /dev/null @@ -1,72 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "kv_cache_manager/common/cache/advanced_cache.h" -#include "kv_cache_manager/optimizer/index/online/cache_indexer.h" - -namespace kv_cache_manager { - -class LruCacheIndexer : public CacheIndexer { -public: - explicit LruCacheIndexer(bool enable_theoretical_max_cache = false); - - void Init(const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step) override; - - void ProcessKeys(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits = nullptr) override; - - int64_t unique_count() const override { return unique_count_; } - int64_t eviction_count() const override { return eviction_count_; } - int64_t memory_usage_bytes() const override; - int64_t kv_cache_usage_bytes() const override; - std::vector capacity_unique_counts() const override; - - void PostQueryMaintenance() override; - bool RemoveKey(int64_t key) override; - -private: - // Coarse bookkeeping estimate for the online simulation metadata per cached key. - static constexpr int64_t kEstimatedCacheEntryOverheadBytes = 200; - - void RebuildCaches(); - void ProcessKeysFullAttention(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits); - void ProcessKeysLinearAttention(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits); - - bool LookupAndInsert(Cache *cache, std::string_view key_sv, bool &is_new_key); - bool LookupAndInsert(Cache *cache, - std::string_view key_sv, - bool is_linear, - int64_t desired_charge, - bool &is_new_key, - bool &is_checkpoint); - - bool enable_theoretical_max_cache_; - int64_t unique_count_ = 0; - int64_t eviction_count_ = 0; - - int64_t size_full_only_ = 0; - int64_t size_full_linear_ = 0; - int32_t linear_step_ = 1; - - std::vector capacity_bytes_; - - static const Cache::CacheItemHelper kHelper; - std::vector> caches_; - std::shared_ptr max_cache_; -}; - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.cc b/kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.cc deleted file mode 100644 index b94e8b8e8..000000000 --- a/kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.cc +++ /dev/null @@ -1,132 +0,0 @@ -#include "kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.h" - -#include - -#include "kv_cache_manager/common/timestamp_util.h" - -namespace kv_cache_manager { - -TtlCacheIndexerWrapper::TtlCacheIndexerWrapper(std::unique_ptr inner, int64_t ttl_seconds) - : TtlCacheIndexerWrapper(std::move(inner), ttl_seconds, []() { return TimestampUtil::GetCurrentTimeSec(); }) {} - -TtlCacheIndexerWrapper::TtlCacheIndexerWrapper(std::unique_ptr inner, - int64_t ttl_seconds, - ClockFunc clock) - : inner_(std::move(inner)) - , ttl_seconds_(ttl_seconds) - , clock_(std::move(clock)) - , hit_age_bucket_counts_(hit_age_thresholds_.size() + 1, 0) {} - -void TtlCacheIndexerWrapper::Init(const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step) { - inner_->Init(capacity_gb, size_full_only, size_full_linear, linear_step); -} - -void TtlCacheIndexerWrapper::ProcessKeys(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits) { - int64_t now = clock_(); - - HarvestExpired(now); - - std::vector inner_key_hits; - inner_->ProcessKeys(keys, hit_count, max_hit_count, &inner_key_hits); - if (key_hits) { - *key_hits = inner_key_hits; - } - - for (size_t i = 0; i < keys.size(); ++i) { - int64_t key = keys[i]; - auto it = key_access_time_.find(key); - if (it != key_access_time_.end()) { - if (i < inner_key_hits.size() && inner_key_hits[i]) { - int64_t age_seconds = now - it->second; - size_t bucket_index = FindAgeBucket(age_seconds); - hit_age_bucket_counts_[bucket_index]++; - } - - expire_set_.erase({it->second + ttl_seconds_, key}); - it->second = now; - expire_set_.insert({now + ttl_seconds_, key}); - } else { - key_access_time_[key] = now; - expire_set_.insert({now + ttl_seconds_, key}); - } - } -} - -void TtlCacheIndexerWrapper::HarvestExpired(int64_t now) { - while (!expire_set_.empty()) { - auto it = expire_set_.begin(); - if (it->first > now) { - break; - } - int64_t key = it->second; - expire_set_.erase(it); - key_access_time_.erase(key); - if (inner_->RemoveKey(key)) { - ttl_eviction_count_++; - } - } -} - -int64_t TtlCacheIndexerWrapper::unique_count() const { return inner_->unique_count(); } - -int64_t TtlCacheIndexerWrapper::eviction_count() const { return inner_->eviction_count(); } - -int64_t TtlCacheIndexerWrapper::ttl_eviction_count() const { return ttl_eviction_count_; } - -int64_t TtlCacheIndexerWrapper::memory_usage_bytes() const { - constexpr int64_t kHashMapEntryBytes = 56; - constexpr int64_t kSetNodeBytes = 48; - int64_t ttl_bytes = static_cast(key_access_time_.size()) * kHashMapEntryBytes + - static_cast(expire_set_.size()) * kSetNodeBytes; - return inner_->memory_usage_bytes() + ttl_bytes; -} - -int64_t TtlCacheIndexerWrapper::kv_cache_usage_bytes() const { return inner_->kv_cache_usage_bytes(); } - -std::vector TtlCacheIndexerWrapper::capacity_unique_counts() const { return inner_->capacity_unique_counts(); } - -void TtlCacheIndexerWrapper::PostQueryMaintenance() { inner_->PostQueryMaintenance(); } - -bool TtlCacheIndexerWrapper::RemoveKey(int64_t key) { - auto it = key_access_time_.find(key); - if (it != key_access_time_.end()) { - expire_set_.erase({it->second + ttl_seconds_, key}); - key_access_time_.erase(it); - } - return inner_->RemoveKey(key); -} - -std::vector TtlCacheIndexerWrapper::GetHitAgeBuckets() const { - std::vector result; - result.reserve(hit_age_bucket_counts_.size()); - for (size_t i = 0; i < hit_age_thresholds_.size(); i++) { - result.push_back({hit_age_thresholds_[i], hit_age_bucket_counts_[i]}); - } - // The last bucket covers [last_threshold, +inf), threshold=0 means infinity - result.push_back({0, hit_age_bucket_counts_.back()}); - return result; -} - -void TtlCacheIndexerWrapper::SetHitAgeBucketThresholds(const std::vector &thresholds) { - hit_age_thresholds_ = thresholds; - std::sort(hit_age_thresholds_.begin(), hit_age_thresholds_.end()); - hit_age_bucket_counts_.assign(hit_age_thresholds_.size() + 1, 0); -} - -size_t TtlCacheIndexerWrapper::FindAgeBucket(int64_t age_seconds) const { - // Find the first threshold >= age_seconds. - // Buckets: [0, t0] (t0+1, t1] ... (t_{n-1}, +inf) - // age=3 with thresholds {5,30} → lower_bound → 5 → bucket 0 ("5s") - // age=10 with thresholds {5,30} → lower_bound → 30 → bucket 1 ("30s") - // age=100 with thresholds {5,30} → lower_bound → end → bucket 2 ("+inf") - auto it = std::lower_bound(hit_age_thresholds_.begin(), hit_age_thresholds_.end(), age_seconds); - return static_cast(it - hit_age_thresholds_.begin()); -} - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.h b/kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.h deleted file mode 100644 index 33676fd8a..000000000 --- a/kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.h +++ /dev/null @@ -1,69 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -#include "kv_cache_manager/optimizer/index/online/cache_indexer.h" - -namespace kv_cache_manager { - -class TtlCacheIndexerWrapper : public CacheIndexer { -public: - using ClockFunc = std::function; - - TtlCacheIndexerWrapper(std::unique_ptr inner, int64_t ttl_seconds); - - TtlCacheIndexerWrapper(std::unique_ptr inner, int64_t ttl_seconds, ClockFunc clock); - - void Init(const std::vector &capacity_gb, - int64_t size_full_only, - int64_t size_full_linear, - int32_t linear_step) override; - - void ProcessKeys(const std::vector &keys, - std::vector &hit_count, - int64_t &max_hit_count, - std::vector *key_hits = nullptr) override; - - int64_t unique_count() const override; - int64_t eviction_count() const override; - int64_t ttl_eviction_count() const override; - int64_t memory_usage_bytes() const override; - int64_t kv_cache_usage_bytes() const override; - std::vector capacity_unique_counts() const override; - - void PostQueryMaintenance() override; - bool RemoveKey(int64_t key) override; - std::vector GetHitAgeBuckets() const override; - - // Configure the age bucket thresholds (in seconds, ascending order). - // The last bucket implicitly covers [last_threshold, +inf). - // Default: {5, 30, 60, 120, 300, 600, 1800, 3600, 7200}. - void SetHitAgeBucketThresholds(const std::vector &thresholds); - -private: - void HarvestExpired(int64_t now); - size_t FindAgeBucket(int64_t age_seconds) const; - - std::unique_ptr inner_; - int64_t ttl_seconds_; - ClockFunc clock_; - - std::unordered_map key_access_time_; - std::set> expire_set_; - int64_t ttl_eviction_count_ = 0; - - // Age bucket thresholds (sorted ascending). Each value is the upper bound - // of a bucket. A final "+inf" bucket is always appended. - std::vector hit_age_thresholds_ = {5, 30, 60, 120, 300, 600, 1800, 3600, 7200}; - // hit_age_bucket_counts_ has size = hit_age_thresholds_.size() + 1 - // buckets: [0,5) [5,30) [30,60) ... [7200, +inf) - std::vector hit_age_bucket_counts_; -}; - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/BUILD b/kv_cache_manager/optimizer/liteHit/BUILD index 612efe13f..6eb63c5b4 100644 --- a/kv_cache_manager/optimizer/liteHit/BUILD +++ b/kv_cache_manager/optimizer/liteHit/BUILD @@ -41,13 +41,37 @@ cc_library( hdrs = ["trace_router.h"], ) +cc_library( + name = "weighted_lru_pool", + srcs = ["weighted_lru_pool.cc"], + hdrs = ["weighted_lru_pool.h"], + deps = [":dynamic_fenwick_tree"], +) + +cc_library( + name = "lite_hit_linear", + srcs = ["lite_hit_linear.cc"], + hdrs = ["lite_hit_linear.h"], + deps = [":weighted_lru_pool"], +) + +cc_library( + name = "lite_hit_ttl", + srcs = ["lite_hit_ttl.cc"], + hdrs = ["lite_hit_ttl.h"], + deps = [ + ":lite_hit", + ":weighted_lru_pool", + ], +) + cc_library( name = "lite_hit", srcs = ["lite_hit.cc"], hdrs = ["lite_hit.h"], deps = [ - ":dynamic_fenwick_tree", ":hit_curve", + ":lite_hit_linear", + ":weighted_lru_pool", ], ) - diff --git a/kv_cache_manager/optimizer/liteHit/README.md b/kv_cache_manager/optimizer/liteHit/README.md index 5f2b99b77..2d5d71657 100644 --- a/kv_cache_manager/optimizer/liteHit/README.md +++ b/kv_cache_manager/optimizer/liteHit/README.md @@ -2,7 +2,7 @@ > [中文](README_zh.md) | English -LiteHit is a lightweight, exact LRU hit-rate analyzer for full-attention KVCache blocks. The core replays a trace only once, producing **capacity-independent facts** for each request (`RequestFact`, a hit curve encoded as arithmetic-segment RLE); the hit count for any capacity is derived after the fact from the facts by the stateless projector `HitCurveProjector`. The core never receives a capacity list, nor does it accumulate any per-capacity results. +LiteHit is a lightweight, exact LRU hit-rate analyzer for full-attention and linear/Mamba KVCache. The two models share one weighted recency core that replays a trace only once, producing **capacity-independent facts** for each request; the hit count for any capacity is derived after the fact from the facts by the stateless projector `HitCurveProjector`. The core never receives a capacity list, nor does it accumulate any per-capacity results. The core produces a total-byte-axis `RequestFact` by default; only Full-only instances losslessly compress it into a block-axis `FullRequestFact` RLE. The problem LiteHit solves is: @@ -16,12 +16,12 @@ how to replay the trace only once and precisely answer, for "any" LRU capacity: Capacity no longer needs to be given before the analysis starts. ``` -The first phase only supports the following model: +The currently supported model is: -- full attention; -- every complete block has the same KVCache charge (equal charge, see 6.4); +- full attention (every complete block has the same charge, see 6.4); +- linear/Mamba (Full blocks and Linear states carry different charges, sharing one weighted LRU pool); - exact LRU, with in-request reverse-order submission; -- does not handle linear attention / Mamba, blocks of different sizes, admission, prefetch, or multi-level cache policies; TTL support is "one fixed TTL per group" layered on top of LRU (see §7 TTL Replay), and does not support query-time scanning of arbitrary TTLs. +- does not handle admission, prefetch, or multi-level cache policies; TTL support is "one fixed TTL per group", applying to both Full blocks and Linear states (see §7 TTL Replay), and does not support query-time scanning of arbitrary TTLs. ## 0. Architecture Overview @@ -34,9 +34,10 @@ The first phase only supports the following model: └───────────────┬────────────────────────────┘ │ NormalizedRequest ┌───────────────▼────────────────────────────┐ - │ LiteHit core (capacity-independent) │ - │ ProcessRequest(block_keys) → RequestFact │ - │ state: Fenwick + last_positions │ + │ TtlLiteHit decorator (fixed TTL; 0=pass-through)│ + ├────────────────────────────────────────────┤ + │ LiteHit core (time-free, shared Full/Mamba) │ + │ WeightedLruPool + Linear state policy │ └───────┬───────────────────────┬────────────┘ │ RequestFact │ RequestFact Online path │ │ Offline path @@ -54,9 +55,9 @@ The first phase only supports the following model: Three inviolable layering constraints: -1. **Core is capacity-independent**: `LiteHit::ProcessRequest` only receives block keys and returns a `RequestFact`; there is no capacity parameter. +1. **The core is capacity-independent**: `LiteHit::ProcessRequest` only receives block keys and returns a `RequestFact`; there is no capacity parameter. 2. **Projection is the sole entry point**: all "capacity → hit block count" conversions must go through `HitCurveProjector`; Online and facts query share the same implementation, and no component is allowed to implement boundary logic on its own. -3. **Byte conversion happens only at the projection boundary**: the core and facts are all in block units; `ProjectBytes` performs a single floor division using `block_bytes` at projection time. +3. **byte-step is the default fact**: the core first produces the `RequestFact` on the total-byte axis; only the Full-only equal-charge case uses `ProcessFullRequest` to convert it into block-axis RLE. --- @@ -171,7 +172,7 @@ The Fenwick is not a cache; it is merely an order-statistics representation of t --- -## 5. RequestFact: Hit Curve Encoded as Arithmetic-Segment RLE +## 5. RequestFact: byte-step by Default, Full-only Compressible to RLE ### 5.1 From Per-Block Thresholds to Hit Curve @@ -187,11 +188,11 @@ The sequence `prefix_required[1..h]` (h being the length before cold truncation) hit_blocks(C) = |{ j : prefix_required[j] <= C }| ``` -This monotonic step function is the **hit curve** of this request — it is the entire fact of this request. +This monotonic step function is the **hit curve** of this request — it is the entire fact of this request. The core records each threshold directly as `{min_total_capacity_bytes, hit_blocks}`, forming the default byte-step `RequestFact`. -### 5.2 Under the Contract, Thresholds Are Strictly Increasing ⇒ Arithmetic-Segment RLE Is Lossless +### 5.2 Full-only: Under the Contract, Thresholds Are Strictly Increasing ⇒ Arithmetic-Segment RLE Is Lossless -Under reverse-order submission, the minimum hit capacity of a later block on the chain is **strictly greater** than that of an earlier block (each block deeper adds at least its parent key within the snapshot interval), so for contract inputs `prefix_required` is strictly increasing. A stronger structural property is: adjacent blocks on the chain occupy **consecutive** positions in the global LRU, and the threshold only jumps at "queue-jumping" points (positions where sibling branches interleave). So consecutive `+1` thresholds are compressed into one arithmetic segment: +Under reverse-order submission, the minimum hit capacity of a later block on the chain is **strictly greater** than that of an earlier block (each block deeper adds at least its parent key within the snapshot interval), so for contract inputs `prefix_required` is strictly increasing. In Full-only every object carries the same charge, so byte thresholds divide exactly by `full_charge_bytes`; consecutive `+1` block thresholds are then compressed into one arithmetic segment: ```text HitCurveSegment { start_required_blocks, run_length } @@ -214,11 +215,11 @@ For contract inputs this defense is always a no-op; for non-contract inputs it o ### 5.4 HitCurveProjector ```cpp -ProjectBlocks(fact, capacity_blocks) // linear scan along segments: - // hits += min(run_length, C - start + 1), until start > C -ProjectBytes(fact, capacity_bytes, block_bytes) - // = ProjectBlocks(fact, floor(bytes / block_bytes)) -ProjectInfinite(fact) // = Σ run_length (no capacity miss, only cold miss) +ProjectBytes(fact, capacity_bytes) // default byte-step +ProjectInfinite(fact) // last point of the byte-step curve +ProjectFullBlocks(full_fact, capacity_blocks) // Full-only RLE +ProjectFullBytes(full_fact, bytes, block_bytes) // floor, then project the RLE +ProjectFullInfinite(full_fact) // Σ run_length of the RLE ``` An empty curve means the request head is cold, and hits are 0 at any capacity. @@ -250,11 +251,11 @@ capacity_blocks = floor(capacity_bytes / block_bytes) capacity_gb uses binary conversion: capacity_bytes = capacity_gb * 1024^3 ``` -`block_bytes` comes from the sum of spec.size of each spec in the full location spec group of instance registration (`size_full_only`). Each row of the facts CSV records `block_bytes`, making facts self-describing: even if the charge estimate is corrected later, historical facts can be re-projected. +Online keeps only `capacity_bytes` for both Full-only and Linear instances; no Full-specific block capacity is stored ahead of time. `block_bytes` comes from the sum of spec.size of each spec in the full location spec group of instance registration (the charge of one Full block); only the final projection of Full RLE performs the floor above. Each row of the facts CSV records `block_bytes`, making facts self-describing: even if the charge estimate is corrected later, historical facts can be re-projected. ### 6.4 Equal-Charge Invariant (Premise of Block-Unit RLE) -The hit curve is in block units and `ProjectBytes` performs a single floor division, which is exact **if and only if** the charges of all participating blocks are exactly equal — full-only instances satisfy this (each block's charge is constantly `size_full_only`, an exact value not an average). For linear/Mamba mixed instances, the per-block charge is unequal, and charge-weighted thresholds must be used instead, which is a separate subsequent task; currently Offline rejects instances with `linear_step != 0`. +Full RLE is in block units and `ProjectFullBytes` performs a single floor division, which is exact **if and only if** the charges of all participating objects are exactly equal — full-attention instances satisfy this. For linear/Mamba the Full block and Linear state charges differ, so the default `RequestFact` is kept and projected directly on the **total byte-capacity axis**, with no average-block conversion. --- @@ -275,7 +276,7 @@ batch window = pipeline_worker_count * 256 entries **Fanout mode**: when `fanout_all_instances = true`, each request is broadcast to all lanes (each lane has independent LRU state and independent facts rows); combined with multiple instances of different `block_size`, a single replay can scan multiple analysis granities over the same trace; mutually exclusive with `override_instance_id`. The facts query summary is grouped by instance (one row per instance + one total row), and fanout results are directly readable. -**TTL replay**: when the instance group config has `ttl_seconds != 0`, the `LiteHit` core of that group's lanes layers a fixed TTL on top (consistent with the online `TtlCacheIndexerWrapper` semantics: a block survives while strictly less than TTL since its last access; expired blocks are misses at any capacity and truncate the prefix like cold blocks; every access (hit or miss) refreshes last_access; time is taken from the trace timestamp, and replay is deterministic). Age is monotonic along the LRU stack, and expired blocks do not raise the reuse distance of surviving blocks, so a single replay remains exact for the joint metric of "fixed TTL × arbitrary capacity", and facts are still ordinary hit curve rows. TTL is a replay-time parameter, taken directly from the group's `ttl_seconds`; to scan multiple TTLs, configure multiple groups each with a different `ttl_seconds` (can be combined with fanout to complete in a single replay). +**TTL replay**: each lane wraps the time-free `LiteHit` core with `TtlLiteHit`; with `ttl_seconds == 0` the decorator is a transparent pass-through. Full blocks and Linear states share the same request-time epochs; an object whose age reaches the TTL drops out of the weighted LRU's visible set. Full blocks are touched on every request, while Linear states are touched only when written at periodic positions or the request tail. Time is taken from the trace timestamp, so replay is deterministic. Expired objects do not raise the capacity threshold of surviving objects, so a single replay remains exact for "fixed TTL × arbitrary capacity". To scan multiple TTLs, configure multiple groups. **Fail-fast**: timestamp out of order, unknown instance, length validation failure, zero valid rows in the whole file — any of these fails the whole thing with a reason — facts are an all-or-nothing reconciliation ledger, and silent row loss is not allowed. @@ -287,7 +288,7 @@ batch window = pipeline_worker_count * 256 entries trace_id,instance_id,timestamp_ns,input_token_len,block_size_tokens,block_bytes,hit_curve ``` -`hit_curve` is a quoted JSON array `[[start_required_blocks, run_length], ...]`; string fields are quote-escaped per CSV rules. Each row is independently parseable, self-describing, and re-projectable. +`hit_curve` is a quoted JSON array: default byte-step rows use `bytes:[[min_capacity_bytes, hit_blocks], ...]`; Full-only rows use `rle:[[start_required_blocks, run_length], ...]`. The reader also accepts the legacy `mamba:` byte-step and unprefixed Full RLE. String fields are quote-escaped per CSV rules. ### 7.2 facts query Tool @@ -306,16 +307,17 @@ Memory is only O(number of instances × number of capacity slots) cumulative int ## 8. Online Integration -The Online Optimizer's full-attention `InstanceState` directly holds a `LiteHit` (when the group config has `ttl_seconds != 0`, a fixed TTL is layered on with wall-clock time, with metrics consistent with the linear path's `TtlCacheIndexerWrapper`). Each TraceQuery: +Each `InstanceState` of the Online Optimizer holds one `TtlLiteHit` decorator with a single `LiteHit` core inside. Whenever an instance group config has `ttl_seconds != 0`, a fixed TTL is applied with wall-clock time; with TTL 0 the decorator passes straight through to the core. Semantics match offline replay, only the time source differs. Each TraceQuery: ```text -NormalizeRequest → ProcessRequest → obtain RequestFact - ├─ for each configured capacity slot: ProjectBlocks(fact, lite_hit_capacity_blocks[i]) - ├─ theoretical upper bound: ProjectInfinite(fact) +NormalizeRequest + ├─ linear: ProcessRequest → RequestFact → ProjectBytes + ├─ Full-only: ProcessFullRequest → FullRequestFact → ProjectFullBytes + ├─ theoretical upper bound: ProjectInfinite / ProjectFullInfinite of the corresponding fact └─ update cumulative integers: total_queries / total_input_tokens / total_hits per slot ``` -Online does not persist facts (facts persistence is currently Offline-exclusive); the hit rate of `ListInstances` is derived from cumulative integers (`total_hits * block_size_tokens / total_input_tokens`). linear attention continues to go through the legacy indexer path (when prefix hash is enabled, only `ApplyPrefixHash` is performed). +Online does not persist facts (facts persistence is currently Offline-exclusive); the hit rate of `ListInstances` is derived from cumulative integers (`total_hits * block_size_tokens / total_input_tokens`). linear attention enables the Linear state policy and projects the byte-step directly; the TTL watermark filters both Full and Linear objects, and the resident bytes, unique Full blocks and TTL evictions in the statistics are all computed over the filtered working set. --- @@ -325,8 +327,9 @@ Let N = total number of block accesses, U = number of historically distinct bloc ```text ProcessRequest: O(m log U) (m is the number of request blocks) -ProjectBlocks: O(S), independent of capacity value -core persistent state: Fenwick + last_positions, O(U) +ProjectBytes / ProjectFullBytes: O(S), independent of capacity value +core persistent state: WeightedLruPool (Fenwick + typed positions), O(U) +TTL decorator: request-time epochs + alive-position watermark, O(Q) ``` **No capacity-based pruning**: capacity is only known after the fact, and any key may be used by some future large-capacity query, so the core retains all historical unique keys (this is the inherent cost of capacity independence). The compaction in Section 4 only reclaims abandoned positions, it does not delete keys. `memory_usage_bytes()` / `current_unique_blocks()` provide observability. @@ -366,7 +369,7 @@ Request 3's curve `[[1,2],[4,1]]` reads directly: capacity 2, 3 hits 2 blocks (t Unit tests (`LiteHitTest` / `LiteHitOfflineRunnerTest`) cover: -1. **oracle comparison**: under contract inputs (random tree-shaped chains + `ApplyPrefixHash`), `ProjectBlocks` is **exactly identical** to naive multi-capacity LRU (snapshot evaluation + reverse-order per-block touch) at capacities {0,1,2,4,9,∞}; +1. **oracle comparison**: under contract inputs (random tree-shaped chains + `ApplyPrefixHash`), `ProjectFullBlocks` is **exactly identical** to naive multi-capacity LRU (snapshot evaluation + reverse-order per-block touch) at capacities {0,1,2,4,9,∞}; 2. non-contract input projection ≤ oracle (monotonic defense is pessimistic, never optimistic), infinite capacity still exact; 3. RLE shape: whole-chain replay single segment, queue-jump breaks segments, adjacent segments cannot be merged; 4. projection boundaries: capacity 0, segment boundaries, byte floor conversion; @@ -401,5 +404,5 @@ The three most easily confused points: at the cost of not being able to do capacity-based pruning. 2. block_size_tokens converts tokens, block_bytes converts capacity; they cannot replace each other. 3. The losslessness of arithmetic-segment RLE depends on the prefix hash contract + reverse-order submission + equal charge; - mixed charge (linear/Mamba) must be designed separately. + mixed charge (linear/Mamba) reuses the same core but encodes an explicit step curve on the total-byte axis. ``` diff --git a/kv_cache_manager/optimizer/liteHit/README_zh.md b/kv_cache_manager/optimizer/liteHit/README_zh.md index 8fef523b7..b68b275dc 100644 --- a/kv_cache_manager/optimizer/liteHit/README_zh.md +++ b/kv_cache_manager/optimizer/liteHit/README_zh.md @@ -2,7 +2,7 @@ > 中文 | [English](README.md) -LiteHit 是面向 full-attention KVCache block 的轻量、精确 LRU 命中率分析器。核心只回放一次 trace，为每条请求产出**容量无关的事实**（`RequestFact`，一条等差段 RLE 编码的 hit curve）；任何容量的命中数都由无状态投影器 `HitCurveProjector` 事后从事实推出，核心从不接收容量列表，也不累计任何逐容量结果。 +LiteHit 是面向 full-attention 与 linear/Mamba KVCache 的轻量、精确 LRU 命中率分析器。两种模型共享一个 weighted recency 核心，只回放一次 trace，为每条请求产出**容量无关的事实**；任何容量的命中数都由无状态投影器 `HitCurveProjector` 事后从事实推出，核心从不接收容量列表，也不累计任何逐容量结果。核心默认产出总字节轴 `RequestFact`；只有 Full-only 实例会把它无损压缩成 block 轴 `FullRequestFact` RLE。 LiteHit 要解决的问题是： @@ -16,12 +16,12 @@ LiteHit 要解决的问题是： 容量不再需要在分析开始前给定。 ``` -第一阶段只支持以下模型： +当前支持以下模型： -- full attention； -- 每个完整 block 的 KVCache charge 相同（等 charge，见 6.4）； +- full attention（每个完整 block 等 charge，见 6.4）； +- linear/Mamba（Full block 与 Linear state 不同 charge，共享一个 weighted LRU pool）； - 精确 LRU，请求内倒序提交； -- 不处理 linear attention / Mamba、不同大小 block、admission、prefetch 和多级缓存策略；TTL 支持为"每组一个固定 TTL"叠加在 LRU 之上（见 §7 TTL 回放），不支持任意 TTL 的查询期扫描。 +- 不处理 admission、prefetch 和多级缓存策略；TTL 支持为“每组一个固定 TTL”，同时适用于 Full 与 Linear state（见 §7 TTL 回放），但不支持任意 TTL 的查询期扫描。 ## 0. 架构总览 @@ -33,9 +33,10 @@ LiteHit 要解决的问题是： └───────────────┬────────────────────────────┘ │ NormalizedRequest ┌───────────────▼────────────────────────────┐ - │ LiteHit 核心（容量无关） │ - │ ProcessRequest(block_keys) → RequestFact │ - │ 状态：Fenwick + last_positions │ + │ TtlLiteHit decorator（固定 TTL；0=透传） │ + ├────────────────────────────────────────────┤ + │ LiteHit 核心（无时间概念、Full/Mamba 共享） │ + │ WeightedLruPool + Linear state policy │ └───────┬───────────────────────┬────────────┘ │ RequestFact │ RequestFact Online 路径 │ │ Offline 路径 @@ -55,7 +56,7 @@ LiteHit 要解决的问题是： 1. **核心容量无关**：`LiteHit::ProcessRequest` 只接收 block key，返回 `RequestFact`；不存在容量参数。 2. **投影唯一入口**：所有"容量 → 命中块数"的换算必须经过 `HitCurveProjector`，Online 与 facts query 共用同一实现，禁止任何组件自行实现边界逻辑。 -3. **字节换算只发生在投影边界**：核心与事实全部以 block 为单位；`ProjectBytes` 在投影时用 `block_bytes` 做一次 floor 除法。 +3. **byte-step 是默认事实**：核心先在总字节轴生成 `RequestFact`；只有 Full-only 的等 charge 场景才用 `ProcessFullRequest` 转成 block 轴 RLE。 --- @@ -170,7 +171,7 @@ Fenwick 不是一套缓存，只是全局 LRU 顺序的 order-statistics 表示 --- -## 5. RequestFact：等差段 RLE 编码的 hit curve +## 5. RequestFact：默认 byte-step，Full-only 可压成 RLE ### 5.1 从逐块门槛到 hit curve @@ -186,11 +187,11 @@ prefix_required[j] = max(r1, ..., rj) hit_blocks(C) = |{ j : prefix_required[j] <= C }| ``` -这条单调阶梯函数就是本请求的 **hit curve**——它就是这条请求的全部事实。 +这条单调阶梯函数就是本请求的 **hit curve**——它就是这条请求的全部事实。核心将每个门槛直接记录为 `{min_total_capacity_bytes, hit_blocks}`，形成默认的 byte-step `RequestFact`。 -### 5.2 契约下门槛严格递增 ⇒ 等差段 RLE 无损 +### 5.2 Full-only：契约下门槛严格递增 ⇒ 等差段 RLE 无损 -倒序提交下链上后块的最小命中容量**严格大于**前块（每深一块，快照区间内至少多出它的父 key），因此契约输入的 `prefix_required` 严格递增。更强的结构性质是：链上相邻 block 在全局 LRU 中占据**连续**位置，门槛只在"插队"点（兄弟分支交错进来的位置）跳变。于是把连续 `+1` 的门槛压成一个等差段： +倒序提交下链上后块的最小命中容量**严格大于**前块（每深一块，快照区间内至少多出它的父 key），因此契约输入的 `prefix_required` 严格递增。Full-only 中每个对象 charge 相同，字节门槛可以精确除以 `full_charge_bytes`；于是把连续 `+1` 的 block 门槛压成一个等差段： ```text HitCurveSegment { start_required_blocks, run_length } @@ -213,11 +214,11 @@ encoded_threshold = max(prefix_required[j], last_encoded + 1) ### 5.4 HitCurveProjector ```cpp -ProjectBlocks(fact, capacity_blocks) // 沿段线性扫描： - // hits += min(run_length, C - start + 1)，直到 start > C -ProjectBytes(fact, capacity_bytes, block_bytes) - // = ProjectBlocks(fact, floor(bytes / block_bytes)) -ProjectInfinite(fact) // = Σ run_length（无 capacity miss，仅 cold miss） +ProjectBytes(fact, capacity_bytes) // 默认 byte-step +ProjectInfinite(fact) // byte-step 末点 +ProjectFullBlocks(full_fact, capacity_blocks) // Full-only RLE +ProjectFullBytes(full_fact, bytes, block_bytes) // floor 后投影 RLE +ProjectFullInfinite(full_fact) // RLE 的 Σ run_length ``` 空 curve 表示请求头即 cold，任何容量命中为 0。 @@ -249,11 +250,11 @@ capacity_blocks = floor(capacity_bytes / block_bytes) capacity_gb 使用二进制换算：capacity_bytes = capacity_gb * 1024^3 ``` -`block_bytes` 来自实例注册的 full location spec group 各 spec.size 之和（`size_full_only`）。facts CSV 每行记录 `block_bytes`，事实自描述：即使日后修正 charge 估计，也能对历史事实重投影。 +Online 对 Full-only 和 Linear instance 都只保存 `capacity_bytes`，不会提前保存一份 Full 专属的 block 容量。`block_bytes` 来自实例注册的 full location spec group 各 spec.size 之和（一个 Full block 的 charge）；只有 Full RLE 最终投影时才执行上述 floor。facts CSV 每行记录 `block_bytes`，事实自描述：即使日后修正 charge 估计，也能对历史事实重投影。 ### 6.4 等 charge 不变量（block 单位 RLE 的前提） -hit curve 以 block 为单位、`ProjectBytes` 做单一 floor 除法，**当且仅当**所有参与块 charge 完全相等才精确——full-only 实例满足（每块 charge 恒为 `size_full_only`，是精确值不是平均值）。linear/Mamba 混合实例每块 charge 不等，必须改用 charge 加权门槛，属于后续独立任务；当前 Offline 拒绝 `linear_step != 0` 的实例。 +Full RLE 以 block 为单位、`ProjectFullBytes` 做单一 floor 除法，**当且仅当**所有参与对象 charge 完全相等才精确——full-attention 实例满足。linear/Mamba 的 Full block 与 Linear state charge 不等，因此保留默认 `RequestFact`，直接在**总字节容量轴**上投影，不做平均 block 换算。 --- @@ -274,7 +275,7 @@ Offline runner（`lite_hit_main` + `OptimizerLiteHitConfig`）逐批处理标准 **fanout 模式**：`fanout_all_instances = true` 时每条请求广播到全部 lane（各 lane 独立 LRU 状态、独立 facts 行），配合多个不同 `block_size` 的 instance 即可一次回放对同一份 trace 扫多个分析粒度；与 `override_instance_id` 互斥。facts query 的 summary 按 instance 分组输出（每 instance 一行 + 总计一行），fanout 结果直接可读。 -**TTL 回放**：instance group 配置 `ttl_seconds != 0` 时，该组 lane 的 `LiteHit` 核心叠加固定 TTL（与 online `TtlCacheIndexerWrapper` 语义一致：块在距上次访问严格小于 TTL 内存活，过期块对任意容量都是 miss 并像冷块一样截断前缀；每次访问（命中或未命中）都刷新 last_access；时间取 trace 时间戳，回放确定性）。年龄沿 LRU 栈单调，过期块不会抬高存活块的复用距离，因此一次回放对"固定 TTL × 任意容量"的联合口径仍然精确，facts 仍是普通 hit curve 行。TTL 是回放期参数，直接取组里的 `ttl_seconds`；要扫多个 TTL 就配多个 group 各带不同 `ttl_seconds`（可配合 fanout 一次回放完成）。 +**TTL 回放**：每个 lane 用 `TtlLiteHit` 装饰无时间概念的 `LiteHit` 核心；`ttl_seconds == 0` 时 decorator 透明透传。Full 与 Linear state 使用同一个请求时间 epoch，age 达到 TTL 即从 weighted LRU 的可见集合失效；Full 每次请求都会 touch，Linear state 只在周期位置或请求尾写入时 touch。时间取 trace 时间戳，因此回放确定性。过期对象不会抬高存活对象的容量门槛，一次回放对“固定 TTL × 任意容量”仍然精确。要扫多个 TTL 就配置多个 group。 **fail-fast**：时间戳乱序、未知 instance、长度校验失败、全文件零有效行，任一发生即整体失败并给出原因——facts 是全有或全无的对账账本，不允许静默丢行。 @@ -286,7 +287,7 @@ Offline runner（`lite_hit_main` + `OptimizerLiteHitConfig`）逐批处理标准 trace_id,instance_id,timestamp_ns,input_token_len,block_size_tokens,block_bytes,hit_curve ``` -`hit_curve` 是带引号的 JSON 数组 `[[start_required_blocks, run_length], ...]`；字符串字段按 CSV 规则引用转义。每行独立可解析、自描述、可重投影。 +`hit_curve` 是带引号的 JSON 数组：默认 byte-step 行使用 `bytes:[[min_capacity_bytes, hit_blocks], ...]`；Full-only 使用 `rle:[[start_required_blocks, run_length], ...]`。读取端还兼容旧的 `mamba:` byte-step 与无前缀 Full RLE。字符串字段按 CSV 规则引用转义。 ### 7.2 facts query 工具 @@ -305,16 +306,17 @@ trace_id,instance_id,timestamp_ns,input_token_len,block_size_tokens,block_bytes, ## 8. Online 集成 -Online Optimizer 的 full-attention `InstanceState` 直接持有 `LiteHit`（组配置 `ttl_seconds != 0` 时以墙钟时间叠加固定 TTL，口径与 linear 路径的 `TtlCacheIndexerWrapper` 一致）。每次 TraceQuery： +Online Optimizer 的每个 `InstanceState` 持有一个 `TtlLiteHit` decorator，内部只有一个 `LiteHit` core。任一实例组配置 `ttl_seconds != 0` 时都以墙钟时间应用固定 TTL；TTL 为 0 时直接透传到 core。语义与离线回放一致，只是时间源不同。每次 TraceQuery： ```text -NormalizeRequest → ProcessRequest → 得到 RequestFact - ├─ 对每个配置容量 slot：ProjectBlocks(fact, lite_hit_capacity_blocks[i]) - ├─ 理论上界：ProjectInfinite(fact) +NormalizeRequest + ├─ linear：ProcessRequest → RequestFact → ProjectBytes + ├─ Full-only：ProcessFullRequest → FullRequestFact → ProjectFullBytes + ├─ 理论上界：对应 fact 的 ProjectInfinite / ProjectFullInfinite └─ 更新累计整数：total_queries / total_input_tokens / 各 slot total_hits ``` -Online 不持久化 facts（facts 落盘当前是 Offline 专属）；`ListInstances` 的命中率从累计整数推导（`total_hits * block_size_tokens / total_input_tokens`）。linear attention 继续走 legacy indexer 路径（启用 prefix hash 时仅做 `ApplyPrefixHash`）。 +Online 不持久化 facts（facts 落盘当前是 Offline 专属）；`ListInstances` 的命中率从累计整数推导（`total_hits * block_size_tokens / total_input_tokens`）。linear attention 启用 Linear state 策略并直接投影 byte-step；TTL 水位线同时过滤 Full 与 Linear 对象，统计中的 resident bytes、unique Full blocks 和 TTL eviction 都按过滤后的 working set 计算。 --- @@ -324,8 +326,9 @@ Online 不持久化 facts（facts 落盘当前是 Offline 专属）；`ListInsta ```text ProcessRequest：O(m log U)（m 为请求块数） -ProjectBlocks： O(S)，与容量值无关 -核心持久状态： Fenwick + last_positions，O(U) +ProjectBytes / ProjectFullBytes：O(S)，与容量值无关 +核心持久状态： WeightedLruPool（Fenwick + typed positions），O(U) +TTL decorator：请求时间 epoch + 存活位置 watermark，O(Q) ``` **没有基于容量的剪枝**：容量事后才知道，任何 key 都可能被将来某个大容量查询用到，因此核心保留全部历史 unique key（这是容量无关性的固有代价）。第 4 节的 compaction 只回收废弃位置，不删除 key。`memory_usage_bytes()` / `current_unique_blocks()` 提供观测。 @@ -365,7 +368,7 @@ block_size_tokens = 4，block_bytes = 1024 单元测试（`LiteHitTest` / `LiteHitOfflineRunnerTest`）覆盖： -1. **oracle 对拍**：契约输入（随机树状链 + `ApplyPrefixHash`）下，`ProjectBlocks` 与朴素多容量 LRU（快照评估 + 倒序逐块 touch）在容量 {0,1,2,4,9,∞} 上**完全一致**； +1. **oracle 对拍**：契约输入（随机树状链 + `ApplyPrefixHash`）下，`ProjectFullBlocks` 与朴素多容量 LRU（快照评估 + 倒序逐块 touch）在容量 {0,1,2,4,9,∞} 上**完全一致**； 2. 非契约输入投影 ≤ oracle（单调防御悲观、绝不乐观），无限容量仍精确； 3. RLE 形态：整链重放单段、插队断段、相邻段不可合并； 4. 投影边界：容量 0、段边界、字节 floor 换算； @@ -400,5 +403,5 @@ block_size_tokens = 4，block_bytes = 1024 代价是不能做基于容量的剪枝。 2. block_size_tokens 换 token，block_bytes 换容量，不能互相替代。 3. 等差段 RLE 的无损性依赖前缀 hash 契约 + 倒序提交 + 等 charge； - 混合 charge（linear/Mamba）必须另行设计。 + 混合 charge（linear/Mamba）复用同一核心，但用总字节轴显式阶梯编码。 ``` diff --git a/kv_cache_manager/optimizer/liteHit/facts_csv.cc b/kv_cache_manager/optimizer/liteHit/facts_csv.cc index a56c986a4..85248b838 100644 --- a/kv_cache_manager/optimizer/liteHit/facts_csv.cc +++ b/kv_cache_manager/optimizer/liteHit/facts_csv.cc @@ -3,6 +3,7 @@ #include #include #include +#include namespace kv_cache_manager { @@ -93,8 +94,8 @@ bool ParseInt64Field(const std::string &field, int64_t &value, const char *name, } // Parses a JSON array of [uint, uint] pairs, e.g. "[[1,2],[4,1]]". -bool ParseHitCurve(const std::string &field, std::vector &segments, std::string &error) { - segments.clear(); +bool ParsePairArray(const std::string &field, std::vector> &pairs, std::string &error) { + pairs.clear(); std::size_t i = 0; const auto skip_spaces = [&] { while (i < field.size() && std::isspace(static_cast(field[i]))) { @@ -140,8 +141,8 @@ bool ParseHitCurve(const std::string &field, std::vector &segme return false; } ++i; - HitCurveSegment segment; - if (!parse_uint(segment.start_required_blocks)) { + std::pair pair; + if (!parse_uint(pair.first)) { error = "hit_curve segment start is not a valid integer"; return false; } @@ -151,7 +152,7 @@ bool ParseHitCurve(const std::string &field, std::vector &segme return false; } ++i; - if (!parse_uint(segment.run_length)) { + if (!parse_uint(pair.second)) { error = "hit_curve segment length is not a valid integer"; return false; } @@ -161,7 +162,7 @@ bool ParseHitCurve(const std::string &field, std::vector &segme return false; } ++i; - segments.push_back(segment); + pairs.push_back(pair); skip_spaces(); if (i < field.size() && field[i] == ',') { ++i; @@ -183,19 +184,72 @@ bool ParseHitCurve(const std::string &field, std::vector &segme return true; } +constexpr const char *kByteStepPrefix = "bytes:"; +constexpr const char *kFullRlePrefix = "rle:"; +constexpr const char *kLegacyMambaPrefix = "mamba:"; + +bool ParseHitCurveField(const std::string &field, LiteHitFactRecord &record, std::string &error) { + std::vector> pairs; + std::size_t byte_prefix_size = 0; + if (field.rfind(kByteStepPrefix, 0) == 0) { + byte_prefix_size = std::string(kByteStepPrefix).size(); + } else if (field.rfind(kLegacyMambaPrefix, 0) == 0) { + byte_prefix_size = std::string(kLegacyMambaPrefix).size(); + } + if (byte_prefix_size > 0) { + record.is_full_rle = false; + record.full_rle_fact.hit_curve.clear(); + if (!ParsePairArray(field.substr(byte_prefix_size), pairs, error)) { + return false; + } + record.fact.points.clear(); + record.fact.points.reserve(pairs.size()); + for (const auto &[capacity, hits] : pairs) { + record.fact.points.push_back(ByteStepPoint{capacity, hits}); + } + return true; + } + + record.is_full_rle = true; + record.fact.points.clear(); + const std::string rle_field = + field.rfind(kFullRlePrefix, 0) == 0 ? field.substr(std::string(kFullRlePrefix).size()) : field; + if (!ParsePairArray(rle_field, pairs, error)) { + return false; + } + record.full_rle_fact.hit_curve.clear(); + record.full_rle_fact.hit_curve.reserve(pairs.size()); + for (const auto &[start, length] : pairs) { + record.full_rle_fact.hit_curve.push_back(HitCurveSegment{start, length}); + } + return true; +} + } // namespace std::string SerializeLiteHitFactRow(const LiteHitFactRecord &record) { std::ostringstream curve; - curve << '['; - for (std::size_t i = 0; i < record.fact.hit_curve.size(); ++i) { - const HitCurveSegment &segment = record.fact.hit_curve[i]; - if (i > 0) { - curve << ','; + if (!record.is_full_rle) { + curve << kByteStepPrefix << '['; + for (std::size_t i = 0; i < record.fact.points.size(); ++i) { + const ByteStepPoint &point = record.fact.points[i]; + if (i > 0) { + curve << ','; + } + curve << '[' << point.min_total_capacity_bytes << ',' << point.hit_blocks << ']'; + } + curve << ']'; + } else { + curve << kFullRlePrefix << '['; + for (std::size_t i = 0; i < record.full_rle_fact.hit_curve.size(); ++i) { + const HitCurveSegment &segment = record.full_rle_fact.hit_curve[i]; + if (i > 0) { + curve << ','; + } + curve << '[' << segment.start_required_blocks << ',' << segment.run_length << ']'; } - curve << '[' << segment.start_required_blocks << ',' << segment.run_length << ']'; + curve << ']'; } - curve << ']'; std::ostringstream row; row << QuoteCsvField(record.trace_id) << ',' << QuoteCsvField(record.instance_id) << ',' << record.timestamp_ns @@ -219,7 +273,7 @@ bool ParseLiteHitFactRow(const std::string &line, LiteHitFactRecord &record, std ParseUint64Field(fields[3], record.input_token_len, "input_token_len", error) && ParseUint64Field(fields[4], record.block_size_tokens, "block_size_tokens", error) && ParseUint64Field(fields[5], record.block_bytes, "block_bytes", error) && - ParseHitCurve(fields[6], record.fact.hit_curve, error); + ParseHitCurveField(fields[6], record, error); } } // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/facts_csv.h b/kv_cache_manager/optimizer/liteHit/facts_csv.h index c9d9a65e4..71a96a675 100644 --- a/kv_cache_manager/optimizer/liteHit/facts_csv.h +++ b/kv_cache_manager/optimizer/liteHit/facts_csv.h @@ -8,9 +8,10 @@ namespace kv_cache_manager { -// One row of the Full-only facts CSV. The row is a capacity-independent, -// recomputable fact: any capacity can be projected from hit_curve afterwards -// without replaying the trace. +// One row of the facts CSV. The row is a capacity-independent, recomputable +// fact: any capacity can be projected from hit_curve afterwards without +// replaying the trace. The default fact is an explicit byte-step curve; +// Full-only rows use the smaller arithmetic-run RLE on the block axis. struct LiteHitFactRecord { std::string trace_id; std::string instance_id; @@ -19,9 +20,12 @@ struct LiteHitFactRecord { uint64_t block_size_tokens = 0; // Per-block byte charge used at projection boundaries. Recording it per // row keeps facts self-describing so a corrected charge estimate can - // still reproject historical facts. + // still reproject historical facts. For byte-step rows this is the Full + // object charge for observability; projection already uses the byte axis. uint64_t block_bytes = 0; - RequestFact fact; + bool is_full_rle = false; + RequestFact fact; // default byte-step rows + FullRequestFact full_rle_fact; // Full-only rows }; inline constexpr const char *kLiteHitFactsCsvHeader = @@ -31,7 +35,9 @@ inline constexpr const char *kLiteHitFactsFileName = "litehit_facts.csv"; // Serializes one record to a CSV line (without trailing newline). String // fields are quoted and escaped when needed; hit_curve is always a quoted -// JSON array of [start_required_blocks, run_length] segments. +// JSON array. New rows use "bytes:" for byte-step facts and "rle:" for +// Full-only facts. The parser also accepts legacy "mamba:" byte steps and +// unprefixed Full RLE rows. std::string SerializeLiteHitFactRow(const LiteHitFactRecord &record); // Parses one CSV line produced by SerializeLiteHitFactRow. Returns false on diff --git a/kv_cache_manager/optimizer/liteHit/facts_query.cc b/kv_cache_manager/optimizer/liteHit/facts_query.cc index 001589bbb..58db4fab3 100644 --- a/kv_cache_manager/optimizer/liteHit/facts_query.cc +++ b/kv_cache_manager/optimizer/liteHit/facts_query.cc @@ -137,10 +137,15 @@ bool RunLiteHitFactsQuery(const std::string &facts_csv_path, totals.hit_tokens.assign(slots.size(), 0); } for (std::size_t i = 0; i < slots.size(); ++i) { - const uint64_t hits = - slots[i].infinite - ? HitCurveProjector::ProjectInfinite(record.fact) - : HitCurveProjector::ProjectBytes(record.fact, slots[i].capacity_bytes, record.block_bytes); + uint64_t hits = 0; + if (record.is_full_rle) { + hits = slots[i].infinite ? HitCurveProjector::ProjectFullInfinite(record.full_rle_fact) + : HitCurveProjector::ProjectFullBytes( + record.full_rle_fact, slots[i].capacity_bytes, record.block_bytes); + } else { + hits = slots[i].infinite ? HitCurveProjector::ProjectInfinite(record.fact) + : HitCurveProjector::ProjectBytes(record.fact, slots[i].capacity_bytes); + } hit_blocks[i] = hits; const uint64_t hit_tokens = hits * record.block_size_tokens; hit_rates[i] = record.input_token_len == 0 diff --git a/kv_cache_manager/optimizer/liteHit/hit_curve.cc b/kv_cache_manager/optimizer/liteHit/hit_curve.cc index 01d7bcb4f..3f9bdd979 100644 --- a/kv_cache_manager/optimizer/liteHit/hit_curve.cc +++ b/kv_cache_manager/optimizer/liteHit/hit_curve.cc @@ -4,7 +4,22 @@ namespace kv_cache_manager { -uint64_t HitCurveProjector::ProjectBlocks(const RequestFact &fact, uint64_t capacity_blocks) { +uint64_t HitCurveProjector::ProjectBytes(const RequestFact &fact, uint64_t total_capacity_bytes) { + uint64_t hits = 0; + for (const ByteStepPoint &point : fact.points) { + if (point.min_total_capacity_bytes > total_capacity_bytes) { + break; + } + hits = point.hit_blocks; + } + return hits; +} + +uint64_t HitCurveProjector::ProjectInfinite(const RequestFact &fact) { + return fact.points.empty() ? 0 : fact.points.back().hit_blocks; +} + +uint64_t HitCurveProjector::ProjectFullBlocks(const FullRequestFact &fact, uint64_t capacity_blocks) { uint64_t hits = 0; for (const HitCurveSegment &segment : fact.hit_curve) { if (segment.start_required_blocks > capacity_blocks) { @@ -16,11 +31,12 @@ uint64_t HitCurveProjector::ProjectBlocks(const RequestFact &fact, uint64_t capa return hits; } -uint64_t HitCurveProjector::ProjectBytes(const RequestFact &fact, uint64_t capacity_bytes, uint64_t block_bytes) { - return ProjectBlocks(fact, capacity_bytes / block_bytes); +uint64_t +HitCurveProjector::ProjectFullBytes(const FullRequestFact &fact, uint64_t capacity_bytes, uint64_t block_bytes) { + return ProjectFullBlocks(fact, capacity_bytes / block_bytes); } -uint64_t HitCurveProjector::ProjectInfinite(const RequestFact &fact) { +uint64_t HitCurveProjector::ProjectFullInfinite(const FullRequestFact &fact) { uint64_t hits = 0; for (const HitCurveSegment &segment : fact.hit_curve) { hits += segment.run_length; diff --git a/kv_cache_manager/optimizer/liteHit/hit_curve.h b/kv_cache_manager/optimizer/liteHit/hit_curve.h index 75082e574..68123a898 100644 --- a/kv_cache_manager/optimizer/liteHit/hit_curve.h +++ b/kv_cache_manager/optimizer/liteHit/hit_curve.h @@ -20,27 +20,47 @@ struct HitCurveSegment { } }; -// Capacity-independent facts of one request replay. An empty curve means no -// capacity can hit this request (cold prefix head). -struct RequestFact { +// Full-only compact representation. Equal Full charges make consecutive byte +// thresholds an arithmetic run on the block axis. +struct FullRequestFact { std::vector hit_curve; }; +// One breakpoint of the core's default byte-axis request curve. Linear +// attention jumps between Linear states; Full-only produces one point per +// covered block before optionally being encoded as FullRequestFact. +struct ByteStepPoint { + uint64_t min_total_capacity_bytes = 0; + uint64_t hit_blocks = 0; + + bool operator==(const ByteStepPoint &other) const { + return min_total_capacity_bytes == other.min_total_capacity_bytes && hit_blocks == other.hit_blocks; + } +}; + +// Points are strictly increasing in both fields (monotone envelope). Empty +// means no byte capacity can recover this request. +struct RequestFact { + std::vector points; +}; + // Stateless projection from capacity-independent facts to hit blocks for a // concrete capacity. Both the online path and the facts post-query must use // this projector; no other component may reimplement the boundary logic. class HitCurveProjector { public: - // Hit blocks when the cache holds exactly capacity_blocks blocks. - static uint64_t ProjectBlocks(const RequestFact &fact, uint64_t capacity_blocks); - - // Hit blocks for a byte capacity; the capacity is floor-converted with the - // per-block byte charge before projection. block_bytes must be positive. - static uint64_t ProjectBytes(const RequestFact &fact, uint64_t capacity_bytes, uint64_t block_bytes); + // Default byte-step curve: largest hit_blocks whose threshold fits. + static uint64_t ProjectBytes(const RequestFact &fact, uint64_t total_capacity_bytes); - // Hit blocks with unbounded capacity: cold misses remain, capacity misses - // disappear, so this is the total run length of the curve. static uint64_t ProjectInfinite(const RequestFact &fact); + + // Full-only RLE projection on the block axis. + static uint64_t ProjectFullBlocks(const FullRequestFact &fact, uint64_t capacity_blocks); + + // Full-only byte projection. block_bytes must be positive. + static uint64_t ProjectFullBytes(const FullRequestFact &fact, uint64_t capacity_bytes, uint64_t block_bytes); + + static uint64_t ProjectFullInfinite(const FullRequestFact &fact); }; } // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/lite_hit.cc b/kv_cache_manager/optimizer/liteHit/lite_hit.cc index 8827c13d6..11641fdc5 100644 --- a/kv_cache_manager/optimizer/liteHit/lite_hit.cc +++ b/kv_cache_manager/optimizer/liteHit/lite_hit.cc @@ -1,246 +1,129 @@ #include "kv_cache_manager/optimizer/liteHit/lite_hit.h" #include -#include +#include namespace kv_cache_manager { -namespace { +LiteHit::LiteHit() : LiteHit(CacheObjectConfig{}) {} -constexpr std::size_t kCompactionSlackPositions = 4096; +LiteHit::LiteHit(const CacheObjectConfig &object_config) + : object_config_(object_config) + , pool_(object_config.full_charge_bytes, object_config.linear_charge_bytes) + , linear_policy_(object_config.linear_charge_bytes, object_config.linear_step_blocks) { + if (object_config_.full_charge_bytes == 0) { + throw std::invalid_argument("LiteHit full_charge_bytes must be positive"); + } +} -} // namespace +RequestFact LiteHit::ProcessRequest(const std::vector &block_keys) { return ProcessRequest(block_keys, 0); } -RequestFact LiteHit::ProcessRequest(const std::vector &block_keys, int64_t now_ns) { - if (ttl_ns_ > 0) { - AdvanceTtlWatermark(now_ns); - } - RequestFact fact = BuildHitCurve(block_keys); - CommitRequest(block_keys, now_ns); - MaybeCompactPositions(); +RequestFact LiteHit::ProcessRequest(const std::vector &block_keys, std::size_t alive_from_position) { + const RequestFact fact = EvaluateRecoveryCurve(block_keys, alive_from_position); + CommitRequest(block_keys, alive_from_position); return fact; } -void LiteHit::AdvanceTtlWatermark(int64_t now_ns) { - // Strict boundary: an epoch whose deadline has been reached is dead, - // matching the online TtlCacheIndexerWrapper harvest. Every position of - // a dead epoch is below the next epoch's start (or below everything when - // it is the last one). - const std::size_t old_watermark = dead_below_position_; - while (!position_epochs_.empty() && - position_epochs_.front().timestamp_ns <= now_ns - static_cast(ttl_ns_)) { - position_epochs_.pop_front(); - dead_below_position_ = position_epochs_.empty() ? fenwick_.size() + 1 : position_epochs_.front().start_position; - } - // Markers the watermark sweeps over are blocks that reached the deadline - // without a refreshing access (re-touched blocks moved their marker - // above); this equals the online wrapper's harvested-eviction count. - if (dead_below_position_ > old_watermark) { - const uint64_t already_dead = old_watermark > 1 ? fenwick_.PrefixSum(old_watermark - 1) : 0; - ttl_expired_blocks_ += fenwick_.PrefixSum(dead_below_position_ - 1) - already_dead; +FullRequestFact LiteHit::ProcessFullRequest(const std::vector &block_keys) { + return ProcessFullRequest(block_keys, 0); +} + +FullRequestFact LiteHit::ProcessFullRequest(const std::vector &block_keys, std::size_t alive_from_position) { + if (uses_linear()) { + throw std::logic_error("ProcessFullRequest requires a Full-only LiteHit"); } + return EncodeFullFact(ProcessRequest(block_keys, alive_from_position)); } -RequestFact LiteHit::BuildHitCurve(const std::vector &block_keys) const { +RequestFact LiteHit::EvaluateRecoveryCurve(const std::vector &block_keys, + std::size_t alive_from_position) const { RequestFact fact; if (block_keys.empty()) { return fact; } - // All ranks are read from one immutable request-start snapshot. Repeated - // keys reuse the same snapshot entry. A cold or expired key stops every - // capacity's prefix, so its later occurrence cannot revive the prefix. - std::unordered_map snapshot_entries; - snapshot_entries.reserve(block_keys.size()); - for (int64_t block_key : block_keys) { - auto [entry_it, inserted] = snapshot_entries.emplace(block_key, SnapshotEntry{}); - if (!inserted) { - continue; + // All ranks come from one immutable request-start snapshot. A cold Full + // block ends prefix coverage for every capacity. + std::vector prefix_full_required(block_keys.size(), 0); + std::size_t covered_blocks = 0; + uint64_t running_full_required = 0; + for (std::size_t i = 0; i < block_keys.size(); ++i) { + uint64_t required = 0; + if (!pool_.RequiredBytes({CacheObjectType::kFull, block_keys[i]}, required, alive_from_position)) { + break; } - const auto previous = last_positions_.find(block_key); - if (previous != last_positions_.end() && previous->second >= dead_below_position_) { - entry_it->second.is_resident = true; - entry_it->second.required_blocks = ReuseDistance(previous->second) + 1; + running_full_required = std::max(running_full_required, required); + prefix_full_required[i] = running_full_required; + covered_blocks = i + 1; + } + + uint64_t last_full_threshold = 0; + for (std::size_t position = 0; position < covered_blocks; ++position) { + uint64_t threshold = prefix_full_required[position]; + if (uses_linear()) { + uint64_t linear_required = 0; + if (!linear_policy_.RequiredLinearBytes( + pool_, block_keys[position], alive_from_position, linear_required)) { + continue; + } + threshold = std::max(threshold, linear_required); + + // A later restore point with no larger threshold dominates the + // earlier point. Removing it keeps a strictly monotone envelope. + while (!fact.points.empty() && fact.points.back().min_total_capacity_bytes >= threshold) { + fact.points.pop_back(); + } + } else { + // Contract-valid prefix hashes make thresholds strictly increase. + // For defensive duplicate-key input, force one Full charge of + // progress per additional hit so projection is never optimistic. + threshold = std::max(threshold, last_full_threshold + object_config_.full_charge_bytes); + last_full_threshold = threshold; } + fact.points.push_back(ByteStepPoint{threshold, static_cast(position + 1)}); } + return fact; +} - uint64_t prefix_required_blocks = 0; - uint64_t last_encoded_threshold = 0; - for (int64_t block_key : block_keys) { - const SnapshotEntry &entry = snapshot_entries.at(block_key); - if (!entry.is_resident) { - break; - } - prefix_required_blocks = std::max(prefix_required_blocks, entry.required_blocks); - // Under the prefix-hash contract thresholds strictly increase and the - // max() below is a no-op. For non-contract input (duplicate keys) it - // keeps the encoded thresholds strictly increasing so the arithmetic - // runs stay representable; the projection is then pessimistic by at - // most the number of duplicates and never optimistic. - const uint64_t threshold = std::max(prefix_required_blocks, last_encoded_threshold + 1); +FullRequestFact LiteHit::EncodeFullFact(const RequestFact &byte_fact) const { + FullRequestFact fact; + for (const ByteStepPoint &point : byte_fact.points) { + const uint64_t required_blocks = point.min_total_capacity_bytes / object_config_.full_charge_bytes; if (!fact.hit_curve.empty() && - fact.hit_curve.back().start_required_blocks + fact.hit_curve.back().run_length == threshold) { - fact.hit_curve.back().run_length++; + fact.hit_curve.back().start_required_blocks + fact.hit_curve.back().run_length == required_blocks) { + ++fact.hit_curve.back().run_length; } else { - fact.hit_curve.push_back(HitCurveSegment{threshold, 1}); + fact.hit_curve.push_back(HitCurveSegment{required_blocks, 1}); } - last_encoded_threshold = threshold; } return fact; } -void LiteHit::CommitRequest(const std::vector &block_keys, int64_t now_ns) { +WeightedLruPool::PositionRemap LiteHit::CommitRequest(const std::vector &block_keys, + std::size_t alive_from_position) { if (block_keys.empty()) { - return; + return {}; } - if (ttl_ns_ > 0) { - // Positions appended below belong to this commit's epoch. Merge with - // the previous epoch when the timestamp did not advance; clamp a - // defensively out-of-order timestamp so the deque stays monotone - // (equivalent to the age-clamp of a per-key timestamp table). - const int64_t epoch_ns = - position_epochs_.empty() ? now_ns : std::max(now_ns, position_epochs_.back().timestamp_ns); - if (position_epochs_.empty() || position_epochs_.back().timestamp_ns != epoch_ns) { - position_epochs_.push_back(PositionEpoch{fenwick_.size() + 1, epoch_ns}); - } - } - - // State commits tail-to-head: sequentially touching the request in - // reverse order produces a final LRU order determined only by each - // distinct key's last touch, which is its first occurrence visited - // back-to-front. The chain head therefore ends up most recent and the - // eviction victim is always a chain leaf. Remove old markers once, then - // append those touches in reverse request order. - std::unordered_map first_occurrence; - first_occurrence.reserve(block_keys.size()); for (std::size_t i = block_keys.size(); i > 0; --i) { - first_occurrence[block_keys[i - 1]] = i - 1; - } - - for (const auto &[block_key, _] : first_occurrence) { - const auto previous = last_positions_.find(block_key); - if (previous != last_positions_.end()) { - fenwick_.Add(previous->second, -1); - last_positions_.erase(previous); - } + const std::size_t position = i - 1; + linear_policy_.CommitLinearIfNeeded(pool_, block_keys[position], position, block_keys.size()); + pool_.Touch({CacheObjectType::kFull, block_keys[position]}); } - for (std::size_t i = block_keys.size(); i > 0; --i) { - const int64_t block_key = block_keys[i - 1]; - if (first_occurrence.at(block_key) != i - 1) { - continue; - } - fenwick_.AppendZero(); - const std::size_t current_position = fenwick_.size(); - fenwick_.Add(current_position, 1); - last_positions_[block_key] = current_position; - } + return pool_.MaybeCompactPositions(alive_from_position); } -void LiteHit::MaybeCompactPositions() { - const std::size_t active_positions = static_cast(alive_marker_count()); - if (fenwick_.size() <= kCompactionSlackPositions) { - return; - } - const std::size_t positions_over_slack = fenwick_.size() - kCompactionSlackPositions; - if (active_positions >= (positions_over_slack + 1) / 2) { - return; - } - - // Markers below the TTL watermark are a miss for every capacity forever, - // so compaction drops them together with their table entries; the state - // afterwards is bounded by the alive working set again. - std::vector> ordered_positions; - ordered_positions.reserve(active_positions); - for (auto it = last_positions_.begin(); it != last_positions_.end();) { - if (it->second < dead_below_position_) { - it = last_positions_.erase(it); - } else { - ordered_positions.emplace_back(it->second, it->first); - ++it; - } - } - std::sort(ordered_positions.begin(), ordered_positions.end()); - - // The bucket array never shrinks on erase and would otherwise stay at - // the historical peak. Shrink only when it is far above the surviving - // set (with headroom) so oscillating working sets do not rehash back and - // forth. - if (last_positions_.bucket_count() > 4 * (last_positions_.size() + 1)) { - last_positions_.rehash(2 * last_positions_.size()); - } - - // Old position boundary S maps to (surviving markers below S) + 1; the - // mapping is monotone, so epoch starts and the watermark keep their - // order and semantics. - const auto remap_boundary = [&ordered_positions](std::size_t old_position) -> std::size_t { - const auto it = std::lower_bound( - ordered_positions.begin(), - ordered_positions.end(), - old_position, - [](const std::pair &entry, std::size_t value) { return entry.first < value; }); - return static_cast(it - ordered_positions.begin()) + 1; - }; - for (PositionEpoch &epoch : position_epochs_) { - epoch.start_position = remap_boundary(epoch.start_position); - } - // Equal remapped starts prove there is no alive marker between the two - // epochs, so the earlier one governs an empty range and can be dropped - // without changing any future liveness answer. Keeping the last epoch per - // start bounds the deque by the alive working set (hot keys otherwise - // accumulate one empty epoch per request until the TTL elapses). - std::size_t deduped_size = 0; - for (const PositionEpoch &epoch : position_epochs_) { - if (deduped_size > 0 && position_epochs_[deduped_size - 1].start_position == epoch.start_position) { - position_epochs_[deduped_size - 1].timestamp_ns = epoch.timestamp_ns; - } else { - position_epochs_[deduped_size++] = epoch; - } - } - position_epochs_.resize(deduped_size); - if (dead_below_position_ > 0) { - dead_below_position_ = remap_boundary(dead_below_position_); - } - - DynamicFenwickTree compacted_fenwick; - for (const auto &[_, block_key] : ordered_positions) { - compacted_fenwick.AppendZero(); - const std::size_t compacted_position = compacted_fenwick.size(); - compacted_fenwick.Add(compacted_position, 1); - last_positions_[block_key] = compacted_position; - } - fenwick_ = std::move(compacted_fenwick); -} +uint64_t LiteHit::current_unique_blocks() const { return pool_.resident_full_count(); } -uint64_t LiteHit::ReuseDistance(std::size_t previous_position) const { - return fenwick_.PrefixSum(fenwick_.size()) - fenwick_.PrefixSum(previous_position); +uint64_t LiteHit::FullObjectsWithinTotalBytes(uint64_t total_capacity_bytes) const { + return pool_.FullObjectsWithinBytes(total_capacity_bytes); } -uint64_t LiteHit::alive_marker_count() const { - const uint64_t total = fenwick_.PrefixSum(fenwick_.size()); - if (dead_below_position_ <= 1) { - return total; - } - return total - fenwick_.PrefixSum(dead_below_position_ - 1); -} +uint64_t LiteHit::resident_bytes() const { return pool_.resident_bytes(); } -void LiteHit::Reset() { - fenwick_.Clear(); - last_positions_.clear(); - position_epochs_.clear(); - dead_below_position_ = 0; - ttl_expired_blocks_ = 0; -} +uint64_t LiteHit::memory_usage_bytes() const { return pool_.memory_usage_bytes(); } -uint64_t LiteHit::memory_usage_bytes() const { - uint64_t bytes = fenwick_.memory_usage_bytes(); - bytes += static_cast(last_positions_.bucket_count()) * sizeof(void *); - constexpr uint64_t kEstimatedHashNodeOverhead = sizeof(void *) * 2; - bytes += static_cast(last_positions_.size()) * - (sizeof(std::pair) + kEstimatedHashNodeOverhead); - bytes += static_cast(position_epochs_.size()) * sizeof(PositionEpoch); - return bytes; -} +void LiteHit::Reset() { pool_.Reset(); } } // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/lite_hit.h b/kv_cache_manager/optimizer/liteHit/lite_hit.h index 73c204354..9d020ce1e 100644 --- a/kv_cache_manager/optimizer/liteHit/lite_hit.h +++ b/kv_cache_manager/optimizer/liteHit/lite_hit.h @@ -2,123 +2,74 @@ #include #include -#include -#include #include -#include "kv_cache_manager/optimizer/liteHit/dynamic_fenwick_tree.h" #include "kv_cache_manager/optimizer/liteHit/hit_curve.h" +#include "kv_cache_manager/optimizer/liteHit/lite_hit_linear.h" +#include "kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h" namespace kv_cache_manager { -// LiteHit is an exact LRU replay core for equal-charge full-attention -// blocks. It keeps one global LRU order and, per request, emits -// capacity-independent facts (RequestFact); it never receives a capacity -// list and never accumulates per-capacity results. Any capacity is answered -// afterwards by projecting the facts with HitCurveProjector. +// One capacity-independent LRU replay core for both Full-only and Mamba +// instances. All cache objects share one weighted recency pool. Full-only is +// the uniform-charge specialization where every covered prefix position is a +// restore point; Mamba adds Linear state objects and only their resident +// positions are restore points. // -// An optional fixed TTL adds the online TtlCacheIndexerWrapper semantics on -// top of the LRU stack: a block whose age since last access reached the TTL -// is a miss for every capacity, every access (hit or miss) refreshes -// last_access, and time is the caller-provided trace timestamp so the replay -// stays deterministic. Ages are monotone along the LRU stack (younger blocks -// are always fresher), so expired blocks never inflate the reuse distance of -// alive ones and one replay stays exact for every capacity under the fixed -// TTL. -// -// The TTL needs no per-key timestamps: marker positions are assigned -// monotonically, so a block's last access time is the commit time of the -// request that assigned its current position. A small epoch deque maps -// position ranges to commit timestamps; expired epochs advance a single -// dead-below watermark and liveness is one integer comparison. The extra -// state is O(distinct commit timestamps inside one TTL window), independent -// of the number of unique blocks. +// RequestFact is always evaluated on the byte axis. Full-only callers may +// losslessly encode it afterwards as FullRequestFact block-axis RLE. Linear +// state scheduling owns no duplicate LRU state. TTL is deliberately outside +// this core and is provided by the TtlLiteHit decorator. class LiteHit { public: - // ttl_ns == 0 disables the TTL (pure LRU). - explicit LiteHit(uint64_t ttl_ns = 0) : ttl_ns_(ttl_ns) {} - - // One call is one request boundary. block_keys must be the normalized - // prefix-chained keys of all complete blocks of the request, in request - // order; input length parsing, validation, and prefix hashing happen in - // the shared preprocessing outside the core. now_ns is the request trace - // timestamp (time-sorted); it is only consulted when a TTL is configured. - // - // Phase 1 evaluates the prefix hit curve against the request-start LRU - // snapshot; with a TTL an expired block stops the prefix like a cold - // one. Phase 2 commits every complete block tail-to-head (reverse - // request order), including blocks after the first miss. With - // prefix-chained keys a chain head is therefore always newer than its - // resident descendants, so the global LRU victim is always a leaf, - // matching leaf-first eviction used by production prefix caches - // (vLLM free-queue order, SGLang radix cache). - // - // Under the prefix-hash contract per-block thresholds strictly increase, - // so the arithmetic-run RLE in RequestFact is lossless. Non-contract - // input (duplicate keys inside a request) is first defensively merged to - // its longest prefix per threshold and then encoded monotonically; the - // projection is never optimistic. - RequestFact ProcessRequest(const std::vector &block_keys, int64_t now_ns = 0); + // Describes the byte charge and placement of objects in the shared LRU + // pool. It is not a general LiteHit runtime configuration. + struct CacheObjectConfig { + uint64_t full_charge_bytes = 1; + uint64_t linear_charge_bytes = 0; // 0 = Full-only + uint64_t linear_step_blocks = 0; + }; + + // Full-only constructor; unit charge preserves the public block-axis + // FullRequestFact contract. + LiteHit(); + explicit LiteHit(const CacheObjectConfig &object_config); + + // Default path for every instance kind: explicit byte-axis step curve. + RequestFact ProcessRequest(const std::vector &block_keys); + + // Full-only specialization: losslessly encodes the default byte-step fact + // as block-axis RLE. Mixed Full/Linear charge cannot use this encoding. + FullRequestFact ProcessFullRequest(const std::vector &block_keys); void Reset(); - uint64_t ttl_ns() const { return ttl_ns_; } - - // Advances the TTL watermark to now_ns without processing a request, so - // observability reads reflect the current time instead of the last - // request's. No-op without a TTL. - void AdvanceTime(int64_t now_ns) { - if (ttl_ns_ > 0) { - AdvanceTtlWatermark(now_ns); - } - } - - // Unique blocks currently alive: expired markers below the TTL watermark - // are excluded even though their table entries linger until compaction. - uint64_t current_unique_blocks() const { return alive_marker_count(); } - - // Cumulative blocks that reached the TTL deadline without a refreshing - // access; counted when the watermark sweeps over them, matching the - // online TtlCacheIndexerWrapper eviction statistics. A revived block - // counts again on its next expiry. - uint64_t ttl_expired_blocks() const { return ttl_expired_blocks_; } - - // Coarse memory estimate for observability. It is derived from the state - // already required by the algorithm and does not retain extra trace data. - uint64_t memory_usage_bytes() const; + bool uses_linear() const { return linear_policy_.enabled(); } -private: - struct SnapshotEntry { - bool is_resident = false; - uint64_t required_blocks = 0; - }; + // Full objects currently resident. Linear state objects are excluded. + uint64_t current_unique_blocks() const; - // Positions in [start_position, next epoch's start_position) were - // assigned by a commit at timestamp_ns; timestamps are non-decreasing - // along the deque. - struct PositionEpoch { - std::size_t start_position = 0; - int64_t timestamp_ns = 0; - }; + // Resident Full objects under one total byte capacity. Linear state bytes + // consume the same budget even though Linear objects are not counted. + uint64_t FullObjectsWithinTotalBytes(uint64_t total_capacity_bytes) const; - RequestFact BuildHitCurve(const std::vector &block_keys) const; - void CommitRequest(const std::vector &block_keys, int64_t now_ns); - void MaybeCompactPositions(); - uint64_t ReuseDistance(std::size_t previous_position) const; - // Resident markers at or above the TTL watermark (all markers when no - // watermark is active). - uint64_t alive_marker_count() const; - // Pops expired epochs and advances dead_below_position_. - void AdvanceTtlWatermark(int64_t now_ns); - - uint64_t ttl_ns_ = 0; - DynamicFenwickTree fenwick_; - std::unordered_map last_positions_; - // TTL state (empty when ttl_ns_ == 0): markers below the watermark are - // expired. - std::deque position_epochs_; - std::size_t dead_below_position_ = 0; - uint64_t ttl_expired_blocks_ = 0; + // Infinite-capacity resident working set, including Linear state bytes. + uint64_t resident_bytes() const; + uint64_t memory_usage_bytes() const; + +private: + friend class TtlLiteHit; + + RequestFact ProcessRequest(const std::vector &block_keys, std::size_t alive_from_position); + FullRequestFact ProcessFullRequest(const std::vector &block_keys, std::size_t alive_from_position); + RequestFact EvaluateRecoveryCurve(const std::vector &block_keys, std::size_t alive_from_position) const; + FullRequestFact EncodeFullFact(const RequestFact &byte_fact) const; + WeightedLruPool::PositionRemap CommitRequest(const std::vector &block_keys, + std::size_t alive_from_position); + + CacheObjectConfig object_config_; + WeightedLruPool pool_; + LiteHitLinearPolicy linear_policy_; }; } // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/lite_hit_linear.cc b/kv_cache_manager/optimizer/liteHit/lite_hit_linear.cc new file mode 100644 index 000000000..161bbd171 --- /dev/null +++ b/kv_cache_manager/optimizer/liteHit/lite_hit_linear.cc @@ -0,0 +1,29 @@ +#include "kv_cache_manager/optimizer/liteHit/lite_hit_linear.h" + +namespace kv_cache_manager { + +LiteHitLinearPolicy::LiteHitLinearPolicy(uint64_t linear_charge_bytes, uint64_t step_blocks) + : linear_charge_bytes_(linear_charge_bytes), step_blocks_(step_blocks == 0 ? 1 : step_blocks) {} + +bool LiteHitLinearPolicy::RequiredLinearBytes(const WeightedLruPool &pool, + int64_t prefix_block_key, + std::size_t alive_from_position, + uint64_t &required_bytes) const { + return enabled() && + pool.RequiredBytes({CacheObjectType::kLinear, prefix_block_key}, required_bytes, alive_from_position); +} + +bool LiteHitLinearPolicy::ShouldWriteLinear(std::size_t position, std::size_t total_blocks) const { + return (position + 1) % step_blocks_ == 0 || position == total_blocks - 1; +} + +void LiteHitLinearPolicy::CommitLinearIfNeeded(WeightedLruPool &pool, + int64_t prefix_block_key, + std::size_t position, + std::size_t total_blocks) const { + if (enabled() && ShouldWriteLinear(position, total_blocks)) { + pool.Touch({CacheObjectType::kLinear, prefix_block_key}); + } +} + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/lite_hit_linear.h b/kv_cache_manager/optimizer/liteHit/lite_hit_linear.h new file mode 100644 index 000000000..e57c73c74 --- /dev/null +++ b/kv_cache_manager/optimizer/liteHit/lite_hit_linear.h @@ -0,0 +1,39 @@ +#pragma once + +#include +#include + +#include "kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h" + +namespace kv_cache_manager { + +// Linear-attention restore/write policy layered on the shared LiteHit core. +// It owns no recency state: Full blocks and Linear state objects always live +// in the core's one WeightedLruPool. +class LiteHitLinearPolicy { +public: + LiteHitLinearPolicy(uint64_t linear_charge_bytes = 0, uint64_t step_blocks = 0); + + bool enabled() const { return linear_charge_bytes_ > 0; } + + // Historical Linear states are valid read candidates regardless of + // whether their position belongs to the current request's write schedule. + bool RequiredLinearBytes(const WeightedLruPool &pool, + int64_t prefix_block_key, + std::size_t alive_from_position, + uint64_t &required_bytes) const; + + // Writes periodic Linear states and always the current request's tail. + void CommitLinearIfNeeded(WeightedLruPool &pool, + int64_t prefix_block_key, + std::size_t position, + std::size_t total_blocks) const; + +private: + bool ShouldWriteLinear(std::size_t position, std::size_t total_blocks) const; + + uint64_t linear_charge_bytes_ = 0; + uint64_t step_blocks_ = 1; +}; + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/lite_hit_ttl.cc b/kv_cache_manager/optimizer/liteHit/lite_hit_ttl.cc new file mode 100644 index 000000000..587ae498e --- /dev/null +++ b/kv_cache_manager/optimizer/liteHit/lite_hit_ttl.cc @@ -0,0 +1,116 @@ +#include "kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h" + +#include +#include + +namespace kv_cache_manager { + +uint64_t LiteHitTtlState::AgeNs(int64_t now_ns, int64_t timestamp_ns) { + if (timestamp_ns >= now_ns) { + return 0; + } + return static_cast(now_ns) - static_cast(timestamp_ns); +} + +void LiteHitTtlState::AdvanceTime(int64_t now_ns, const WeightedLruPool &pool) { + if (!enabled()) { + return; + } + + const std::size_t old_watermark = dead_below_position_; + while (!position_epochs_.empty() && AgeNs(now_ns, position_epochs_.front().timestamp_ns) >= ttl_ns_) { + position_epochs_.pop_front(); + dead_below_position_ = + position_epochs_.empty() ? pool.position_count() + 1 : position_epochs_.front().start_position; + } + + if (dead_below_position_ > old_watermark) { + expired_full_blocks_ += pool.FullObjectsBefore(dead_below_position_) - pool.FullObjectsBefore(old_watermark); + } +} + +void LiteHitTtlState::BeginCommit(int64_t now_ns, std::size_t next_position) { + if (!enabled()) { + return; + } + + const int64_t epoch_ns = position_epochs_.empty() ? now_ns : std::max(now_ns, position_epochs_.back().timestamp_ns); + if (position_epochs_.empty() || position_epochs_.back().timestamp_ns != epoch_ns) { + position_epochs_.push_back(PositionEpoch{next_position, epoch_ns}); + } +} + +void LiteHitTtlState::ApplyPositionRemap(const WeightedLruPool::PositionRemap &remap) { + if (!enabled() || !remap.compacted) { + return; + } + + for (PositionEpoch &epoch : position_epochs_) { + epoch.start_position = remap.RemapBoundary(epoch.start_position); + } + + std::size_t deduped_size = 0; + for (const PositionEpoch &epoch : position_epochs_) { + if (deduped_size > 0 && position_epochs_[deduped_size - 1].start_position == epoch.start_position) { + position_epochs_[deduped_size - 1].timestamp_ns = epoch.timestamp_ns; + } else { + position_epochs_[deduped_size++] = epoch; + } + } + position_epochs_.resize(deduped_size); + if (dead_below_position_ > 0) { + dead_below_position_ = remap.RemapBoundary(dead_below_position_); + } +} + +void LiteHitTtlState::Reset() { + position_epochs_.clear(); + dead_below_position_ = 0; + expired_full_blocks_ = 0; +} + +uint64_t LiteHitTtlState::memory_usage_bytes() const { + return static_cast(position_epochs_.size()) * sizeof(PositionEpoch); +} + +TtlLiteHit::TtlLiteHit(uint64_t ttl_ns) : TtlLiteHit(LiteHit::CacheObjectConfig{}, ttl_ns) {} + +TtlLiteHit::TtlLiteHit(const LiteHit::CacheObjectConfig &object_config, uint64_t ttl_ns) + : core_(object_config), ttl_state_(ttl_ns) {} + +RequestFact TtlLiteHit::ProcessRequest(const std::vector &block_keys, int64_t now_ns) { + ttl_state_.AdvanceTime(now_ns, core_.pool_); + const RequestFact fact = core_.EvaluateRecoveryCurve(block_keys, alive_from_position()); + if (!block_keys.empty()) { + ttl_state_.BeginCommit(now_ns, core_.pool_.position_count() + 1); + } + const WeightedLruPool::PositionRemap remap = core_.CommitRequest(block_keys, alive_from_position()); + ttl_state_.ApplyPositionRemap(remap); + return fact; +} + +FullRequestFact TtlLiteHit::ProcessFullRequest(const std::vector &block_keys, int64_t now_ns) { + if (uses_linear()) { + throw std::logic_error("ProcessFullRequest requires a Full-only LiteHit"); + } + return core_.EncodeFullFact(ProcessRequest(block_keys, now_ns)); +} + +void TtlLiteHit::AdvanceTime(int64_t now_ns) { ttl_state_.AdvanceTime(now_ns, core_.pool_); } + +void TtlLiteHit::Reset() { + core_.Reset(); + ttl_state_.Reset(); +} + +uint64_t TtlLiteHit::current_unique_blocks() const { return core_.pool_.resident_full_count(alive_from_position()); } + +uint64_t TtlLiteHit::FullObjectsWithinTotalBytes(uint64_t total_capacity_bytes) const { + return core_.pool_.FullObjectsWithinBytes(total_capacity_bytes, alive_from_position()); +} + +uint64_t TtlLiteHit::resident_bytes() const { return core_.pool_.resident_bytes(alive_from_position()); } + +uint64_t TtlLiteHit::memory_usage_bytes() const { return core_.memory_usage_bytes() + ttl_state_.memory_usage_bytes(); } + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h b/kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h new file mode 100644 index 000000000..20042f8aa --- /dev/null +++ b/kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h @@ -0,0 +1,84 @@ +#pragma once + +#include +#include +#include +#include + +#include "kv_cache_manager/optimizer/liteHit/lite_hit.h" + +namespace kv_cache_manager { + +// Optional time policy for the shared LiteHit recency core. Commit positions +// are monotone with access time, so all expired objects form one old-position +// prefix. Epochs map position ranges to request timestamps without retaining +// a timestamp per key. +class LiteHitTtlState { +public: + explicit LiteHitTtlState(uint64_t ttl_ns = 0) : ttl_ns_(ttl_ns) {} + + bool enabled() const { return ttl_ns_ > 0; } + uint64_t ttl_ns() const { return ttl_ns_; } + std::size_t alive_from_position() const { return dead_below_position_; } + uint64_t expired_full_blocks() const { return expired_full_blocks_; } + + // Advances the strict TTL deadline before a request or observability read. + void AdvanceTime(int64_t now_ns, const WeightedLruPool &pool); + + // Records the timestamp of positions about to be appended by a non-empty + // request. Out-of-order timestamps are clamped to the latest epoch. + void BeginCommit(int64_t now_ns, std::size_t next_position); + + // Applies the pool's dense-position mapping after compaction and removes + // epochs whose position ranges became empty. + void ApplyPositionRemap(const WeightedLruPool::PositionRemap &remap); + + void Reset(); + uint64_t memory_usage_bytes() const; + +private: + struct PositionEpoch { + std::size_t start_position = 0; + int64_t timestamp_ns = 0; + }; + + static uint64_t AgeNs(int64_t now_ns, int64_t timestamp_ns); + + uint64_t ttl_ns_ = 0; + std::deque position_epochs_; + std::size_t dead_below_position_ = 0; + uint64_t expired_full_blocks_ = 0; +}; + +// TTL decorator for the capacity-independent LiteHit core. The wrapped core +// owns all Full/Linear recency state and remains unaware of timestamps. This +// layer supplies a liveness boundary while evaluating and committing, and +// translates pool position compaction back into its timestamp epochs. +class TtlLiteHit { +public: + explicit TtlLiteHit(uint64_t ttl_ns = 0); + TtlLiteHit(const LiteHit::CacheObjectConfig &object_config, uint64_t ttl_ns); + + RequestFact ProcessRequest(const std::vector &block_keys, int64_t now_ns = 0); + FullRequestFact ProcessFullRequest(const std::vector &block_keys, int64_t now_ns = 0); + + void AdvanceTime(int64_t now_ns); + void Reset(); + + bool uses_linear() const { return core_.uses_linear(); } + uint64_t ttl_ns() const { return ttl_state_.ttl_ns(); } + uint64_t ttl_expired_blocks() const { return ttl_state_.expired_full_blocks(); } + + uint64_t current_unique_blocks() const; + uint64_t FullObjectsWithinTotalBytes(uint64_t total_capacity_bytes) const; + uint64_t resident_bytes() const; + uint64_t memory_usage_bytes() const; + +private: + std::size_t alive_from_position() const { return ttl_state_.alive_from_position(); } + + LiteHit core_; + LiteHitTtlState ttl_state_; +}; + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/weighted_lru_pool.cc b/kv_cache_manager/optimizer/liteHit/weighted_lru_pool.cc new file mode 100644 index 000000000..e68d588dc --- /dev/null +++ b/kv_cache_manager/optimizer/liteHit/weighted_lru_pool.cc @@ -0,0 +1,217 @@ +#include "kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h" + +#include +#include +#include + +namespace kv_cache_manager { + +namespace { + +constexpr std::size_t kCompactionSlackPositions = 4096; + +} // namespace + +std::size_t WeightedLruPool::PositionRemap::RemapBoundary(std::size_t old_position) const { + if (!compacted) { + return old_position; + } + const auto it = std::lower_bound(surviving_old_positions.begin(), surviving_old_positions.end(), old_position); + return static_cast(it - surviving_old_positions.begin()) + 1; +} + +bool WeightedLruPool::IsResident(const CacheObjectKey &key, std::size_t alive_from_position) const { + const PositionMap &map = positions(key.type); + const auto it = map.find(key.prefix_block_key); + return it != map.end() && it->second >= NormalizeAliveFrom(alive_from_position); +} + +bool WeightedLruPool::RequiredBytes(const CacheObjectKey &key, + uint64_t &required_bytes, + std::size_t alive_from_position) const { + const PositionMap &map = positions(key.type); + const auto it = map.find(key.prefix_block_key); + if (it == map.end() || it->second < NormalizeAliveFrom(alive_from_position)) { + return false; + } + const uint64_t newer_bytes = fenwick_.PrefixSum(fenwick_.size()) - fenwick_.PrefixSum(it->second); + required_bytes = newer_bytes + charge_of(key.type); + return true; +} + +void WeightedLruPool::Touch(const CacheObjectKey &key) { + PositionMap &map = positions(key.type); + const uint64_t charge = charge_of(key.type); + const bool is_full = key.type == CacheObjectType::kFull; + const auto previous = map.find(key.prefix_block_key); + if (previous != map.end()) { + fenwick_.Add(previous->second, -static_cast(charge)); + if (is_full) { + full_count_fenwick_.Add(previous->second, -1); + } + map.erase(previous); + } + fenwick_.AppendZero(); + full_count_fenwick_.AppendZero(); + const std::size_t current_position = fenwick_.size(); + fenwick_.Add(current_position, static_cast(charge)); + if (is_full) { + full_count_fenwick_.Add(current_position, 1); + } + map[key.prefix_block_key] = current_position; +} + +uint64_t WeightedLruPool::resident_full_count(std::size_t alive_from_position) const { + const std::size_t alive_from = NormalizeAliveFrom(alive_from_position); + if (alive_from > fenwick_.size()) { + return 0; + } + return full_count_fenwick_.PrefixSum(fenwick_.size()) - full_count_fenwick_.PrefixSum(alive_from - 1); +} + +uint64_t WeightedLruPool::resident_linear_count(std::size_t alive_from_position) const { + const std::size_t alive_from = NormalizeAliveFrom(alive_from_position); + if (alive_from <= 1) { + return static_cast(linear_positions_.size()); + } + uint64_t count = 0; + for (const auto &[_, position] : linear_positions_) { + if (position >= alive_from) { + ++count; + } + } + return count; +} + +uint64_t WeightedLruPool::resident_bytes(std::size_t alive_from_position) const { + const std::size_t alive_from = NormalizeAliveFrom(alive_from_position); + if (alive_from > fenwick_.size()) { + return 0; + } + return fenwick_.PrefixSum(fenwick_.size()) - fenwick_.PrefixSum(alive_from - 1); +} + +uint64_t WeightedLruPool::FullObjectsWithinBytes(uint64_t byte_budget, std::size_t alive_from_position) const { + // Object at position p is resident under the budget iff the bytes of all + // markers at positions >= p fit, i.e. total - PrefixSum(p - 1) <= budget. + // That suffix-bytes function is nonincreasing in p, so binary search the + // smallest resident position and count Full markers behind it. + const std::size_t n = fenwick_.size(); + const std::size_t alive_from = NormalizeAliveFrom(alive_from_position); + const uint64_t total_bytes = fenwick_.PrefixSum(n); + if (n == 0 || byte_budget == 0 || alive_from > n) { + return 0; + } + std::size_t lo = alive_from; + std::size_t hi = n + 1; // n + 1 means nothing fits + while (lo < hi) { + const std::size_t mid = lo + (hi - lo) / 2; + if (total_bytes - fenwick_.PrefixSum(mid - 1) <= byte_budget) { + hi = mid; + } else { + lo = mid + 1; + } + } + if (lo > n) { + return 0; + } + return full_count_fenwick_.PrefixSum(n) - full_count_fenwick_.PrefixSum(lo - 1); +} + +uint64_t WeightedLruPool::FullObjectsBefore(std::size_t boundary_position) const { + if (boundary_position <= 1 || fenwick_.size() == 0) { + return 0; + } + return full_count_fenwick_.PrefixSum(std::min(boundary_position - 1, fenwick_.size())); +} + +WeightedLruPool::PositionRemap WeightedLruPool::MaybeCompactPositions(std::size_t alive_from_position) { + PositionRemap remap; + if (fenwick_.size() <= kCompactionSlackPositions) { + return remap; + } + + const std::size_t alive_from = NormalizeAliveFrom(alive_from_position); + std::size_t active_positions = 0; + for (const PositionMap *map : {&full_positions_, &linear_positions_}) { + for (const auto &[_, position] : *map) { + if (position >= alive_from) { + ++active_positions; + } + } + } + const std::size_t positions_over_slack = fenwick_.size() - kCompactionSlackPositions; + if (active_positions >= (positions_over_slack + 1) / 2) { + return remap; + } + + struct Marker { + std::size_t position; + CacheObjectType type; + int64_t prefix_block_key; + }; + std::vector ordered_markers; + ordered_markers.reserve(active_positions); + for (const auto &[block_key, position] : full_positions_) { + if (position >= alive_from) { + ordered_markers.push_back({position, CacheObjectType::kFull, block_key}); + } + } + for (const auto &[block_key, position] : linear_positions_) { + if (position >= alive_from) { + ordered_markers.push_back({position, CacheObjectType::kLinear, block_key}); + } + } + std::sort(ordered_markers.begin(), ordered_markers.end(), [](const Marker &a, const Marker &b) { + return a.position < b.position; + }); + + remap.compacted = true; + remap.surviving_old_positions.reserve(ordered_markers.size()); + full_positions_.clear(); + linear_positions_.clear(); + + DynamicFenwickTree compacted_fenwick; + DynamicFenwickTree compacted_full_count; + for (const Marker &marker : ordered_markers) { + remap.surviving_old_positions.push_back(marker.position); + compacted_fenwick.AppendZero(); + compacted_full_count.AppendZero(); + const std::size_t compacted_position = compacted_fenwick.size(); + compacted_fenwick.Add(compacted_position, static_cast(charge_of(marker.type))); + if (marker.type == CacheObjectType::kFull) { + compacted_full_count.Add(compacted_position, 1); + } + positions(marker.type)[marker.prefix_block_key] = compacted_position; + } + fenwick_ = std::move(compacted_fenwick); + full_count_fenwick_ = std::move(compacted_full_count); + + constexpr std::size_t kBucketShrinkRatio = 4; + for (PositionMap *map : {&full_positions_, &linear_positions_}) { + if (map->bucket_count() > kBucketShrinkRatio * (map->size() + 1)) { + map->rehash(2 * map->size()); + } + } + return remap; +} + +void WeightedLruPool::Reset() { + fenwick_.Clear(); + full_count_fenwick_.Clear(); + full_positions_.clear(); + linear_positions_.clear(); +} + +uint64_t WeightedLruPool::memory_usage_bytes() const { + uint64_t bytes = fenwick_.memory_usage_bytes() + full_count_fenwick_.memory_usage_bytes(); + constexpr uint64_t kEstimatedHashNodeOverhead = sizeof(void *) * 2; + for (const PositionMap *map : {&full_positions_, &linear_positions_}) { + bytes += static_cast(map->bucket_count()) * sizeof(void *); + bytes += static_cast(map->size()) * + (sizeof(std::pair) + kEstimatedHashNodeOverhead); + } + return bytes; +} + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h b/kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h new file mode 100644 index 000000000..633f73d84 --- /dev/null +++ b/kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h @@ -0,0 +1,107 @@ +#pragma once + +#include +#include +#include +#include + +#include "kv_cache_manager/optimizer/liteHit/dynamic_fenwick_tree.h" + +namespace kv_cache_manager { + +// Object identity inside a LiteHit cache pool. Full blocks and Linear state +// objects may share the same prefix-chained key; the type keeps them +// from colliding. +enum class CacheObjectType { + kFull = 0, + kLinear = 1 +}; + +struct CacheObjectKey { + CacheObjectType type = CacheObjectType::kFull; + int64_t prefix_block_key = 0; +}; + +// Weighted generalization of the LiteHit LRU state: one global recency order +// over typed objects with a fixed positive byte charge per type. Range sums +// return resident BYTES instead of block counts, so residency for a byte +// capacity is decided by RequiredBytes(key) <= capacity_bytes (Mattson stack +// inclusion, byte-weighted). It is a LiteHit-internal component. +class WeightedLruPool { +public: + // Linear state charge may be zero when the pool is configured for + // Full-only replay; Linear objects are never touched in that mode. + WeightedLruPool(uint64_t full_charge_bytes, uint64_t linear_charge_bytes) + : charges_{full_charge_bytes, linear_charge_bytes} {} + + // A compaction maps every surviving old marker position to a dense new + // position. TTL watermarks/epoch boundaries use the same lower-bound + // mapping to keep their meaning after positions move. + struct PositionRemap { + bool compacted = false; + std::vector surviving_old_positions; + + std::size_t RemapBoundary(std::size_t old_position) const; + }; + + uint64_t charge_of(CacheObjectType type) const { return charges_[static_cast(type)]; } + + bool IsResident(const CacheObjectKey &key, std::size_t alive_from_position = 0) const; + + // Minimum total pool bytes that keep `key` resident right now: its own + // charge plus the bytes of every strictly newer resident object. + // Returns false when the object is not resident. + bool RequiredBytes(const CacheObjectKey &key, uint64_t &required_bytes, std::size_t alive_from_position = 0) const; + + // Moves the object to the most recent position (inserting it if absent) + // with its fixed type charge. + void Touch(const CacheObjectKey &key); + + void Reset(); + + uint64_t resident_full_count(std::size_t alive_from_position = 0) const; + uint64_t resident_linear_count(std::size_t alive_from_position = 0) const; + // Total bytes of resident objects at/above an optional liveness boundary. + uint64_t resident_bytes(std::size_t alive_from_position = 0) const; + // Number of resident Full objects whose RequiredBytes fits into + // byte_budget, i.e. the Full population of a cache bounded to that many + // bytes (Mattson inclusion boundary). + uint64_t FullObjectsWithinBytes(uint64_t byte_budget, std::size_t alive_from_position = 0) const; + + // Number of live Full markers strictly below boundary_position. Used by + // the TTL policy to count markers crossed by its watermark. + uint64_t FullObjectsBefore(std::size_t boundary_position) const; + + std::size_t position_count() const { return fenwick_.size(); } + + // Compacts dead marker positions when slack is large. Objects below + // alive_from_position are dropped together with their key-map entries. + // The caller must remap any external position boundaries with the result. + PositionRemap MaybeCompactPositions(std::size_t alive_from_position = 0); + + uint64_t memory_usage_bytes() const; + +private: + using PositionMap = std::unordered_map; + + PositionMap &positions(CacheObjectType type) { + return type == CacheObjectType::kFull ? full_positions_ : linear_positions_; + } + const PositionMap &positions(CacheObjectType type) const { + return type == CacheObjectType::kFull ? full_positions_ : linear_positions_; + } + + static std::size_t NormalizeAliveFrom(std::size_t alive_from_position) { + return alive_from_position <= 1 ? 1 : alive_from_position; + } + + uint64_t charges_[2]; + DynamicFenwickTree fenwick_; + // Parallel order-statistics view counting 1 per Full marker at the same + // positions as the byte tree; backs FullObjectsWithinBytes. + DynamicFenwickTree full_count_fenwick_; + PositionMap full_positions_; + PositionMap linear_positions_; +}; + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/manager/BUILD b/kv_cache_manager/optimizer/manager/BUILD index c8b25e82a..d738ea852 100644 --- a/kv_cache_manager/optimizer/manager/BUILD +++ b/kv_cache_manager/optimizer/manager/BUILD @@ -56,11 +56,8 @@ cc_library( "//kv_cache_manager/common", "//kv_cache_manager/optimizer/config:online_optimizer_config", "//kv_cache_manager/optimizer/config:optimizer_registry_manager", - "//kv_cache_manager/optimizer/index:cache_indexer", - "//kv_cache_manager/optimizer/index:cache_indexer_factory", - "//kv_cache_manager/optimizer/index:ttl_cache_indexer_wrapper", "//kv_cache_manager/optimizer/liteHit:hit_curve", - "//kv_cache_manager/optimizer/liteHit:lite_hit", + "//kv_cache_manager/optimizer/liteHit:lite_hit_ttl", "//kv_cache_manager/optimizer/liteHit:request_preprocess", ], ) @@ -79,7 +76,7 @@ cc_library( "//kv_cache_manager/optimizer/config:optimizer_lite_hit_config", "//kv_cache_manager/optimizer/config:optimizer_registry_manager", "//kv_cache_manager/optimizer/liteHit:facts_csv", - "//kv_cache_manager/optimizer/liteHit:lite_hit", + "//kv_cache_manager/optimizer/liteHit:lite_hit_ttl", "//kv_cache_manager/optimizer/liteHit:request_preprocess", "//kv_cache_manager/optimizer/liteHit:trace_router", "//kv_cache_manager/optimizer/trace_loader", diff --git a/kv_cache_manager/optimizer/manager/lite_hit_offline_runner.cc b/kv_cache_manager/optimizer/manager/lite_hit_offline_runner.cc index ee262a82a..2a654130b 100644 --- a/kv_cache_manager/optimizer/manager/lite_hit_offline_runner.cc +++ b/kv_cache_manager/optimizer/manager/lite_hit_offline_runner.cc @@ -16,7 +16,7 @@ #include "kv_cache_manager/common/logger.h" #include "kv_cache_manager/optimizer/config/optimizer_registry_manager.h" #include "kv_cache_manager/optimizer/liteHit/facts_csv.h" -#include "kv_cache_manager/optimizer/liteHit/lite_hit.h" +#include "kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h" #include "kv_cache_manager/optimizer/liteHit/request_preprocess.h" #include "kv_cache_manager/optimizer/liteHit/trace_router.h" #include "kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.h" @@ -30,11 +30,10 @@ namespace { // One per-instance replay lane. LiteHit state updates are serial inside a // lane; preprocessing and row formatting are parallel across the batch. -// A group with ttl_seconds > 0 layers that fixed TTL onto the lane's core. +// Every lane owns a TTL decorator around the same LiteHit core. TTL zero is a +// pass-through; Linear state scheduling remains a core concern. struct InstanceLane { - explicit InstanceLane(uint64_t ttl_ns) : core(ttl_ns) {} - - LiteHit core; + std::unique_ptr core; uint64_t block_size_tokens = 0; uint64_t block_bytes = 0; bool enable_prefix_hash = false; @@ -96,8 +95,7 @@ bool LiteHitOfflineRunner::Run() { } // Per-instance lanes. Registration through the manager validates the - // config and yields size_full_only, the per-block byte charge recorded - // into every fact row. + // config and yields the Full block charge recorded into every fact row. std::unordered_map> lanes; std::vector instance_ids; instance_ids.reserve(config_.instances().size()); @@ -108,12 +106,6 @@ bool LiteHitOfflineRunner::Run() { return false; } for (const auto &instance : config_.instances()) { - if (instance.linear_step() != 0) { - KVCM_LOG_ERROR("LiteHitOfflineRunner: instance[%s] has linear_step=%d; the facts replay is Full-only", - instance.instance_id().c_str(), - instance.linear_step()); - return false; - } if (static_cast(instance.block_size()) % config_.block_size() != 0) { KVCM_LOG_ERROR( "LiteHitOfflineRunner: instance[%s] block_size=%ld is not a multiple of the trace block_size=%lu " @@ -135,10 +127,20 @@ bool LiteHitOfflineRunner::Run() { // instance group and applies to every instance in it. RegisterInstance // already guaranteed the group exists. const OptimizerInstanceGroup &group = *groups_by_name.at(instance.instance_group_name()); - auto lane = std::make_unique(static_cast(group.ttl_seconds()) * 1000000000ULL); + auto lane = std::make_unique(); lane->block_size_tokens = static_cast(instance.block_size()); - lane->block_bytes = static_cast(register_result.size_full_only); + lane->block_bytes = static_cast(register_result.full_charge_bytes); lane->enable_prefix_hash = group.enable_prefix_hash(); + LiteHit::CacheObjectConfig object_config; + object_config.full_charge_bytes = static_cast(register_result.full_charge_bytes); + if (instance.linear_step() != 0) { + // RegisterInstance already validated linear_step as a positive + // token multiple of block_size and the non-empty Mamba spec group. + object_config.linear_charge_bytes = static_cast(register_result.linear_charge_bytes); + object_config.linear_step_blocks = static_cast(instance.linear_step() / instance.block_size()); + } + const uint64_t ttl_ns = static_cast(group.ttl_seconds()) * 1000000000ULL; + lane->core = std::make_unique(object_config, ttl_ns); if (!lanes.emplace(instance.instance_id(), std::move(lane)).second) { KVCM_LOG_ERROR("LiteHitOfflineRunner: duplicate instance_id[%s] in config", instance.instance_id().c_str()); return false; @@ -211,7 +213,14 @@ bool LiteHitOfflineRunner::Run() { // Lane commits stay in input order; only same-lane order is // semantically required, and input order trivially satisfies it. for (BatchItem &item : batch) { - item.record.fact = item.lane->core.ProcessRequest(item.normalized.block_keys, item.record.timestamp_ns); + if (item.lane->core->uses_linear()) { + item.record.fact = + item.lane->core->ProcessRequest(item.normalized.block_keys, item.record.timestamp_ns); + } else { + item.record.is_full_rle = true; + item.record.full_rle_fact = + item.lane->core->ProcessFullRequest(item.normalized.block_keys, item.record.timestamp_ns); + } } ParallelForIndex(batch.size(), worker_count, [&](std::size_t i) { diff --git a/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.cc b/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.cc index c766897f2..1e4392a50 100644 --- a/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.cc +++ b/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.cc @@ -9,7 +9,6 @@ #include "kv_cache_manager/common/logger.h" #include "kv_cache_manager/common/timestamp_util.h" #include "kv_cache_manager/optimizer/config/optimizer_registry_manager.h" -#include "kv_cache_manager/optimizer/index/online/cache_indexer_factory.h" #include "kv_cache_manager/optimizer/liteHit/hit_curve.h" #include "kv_cache_manager/optimizer/liteHit/request_preprocess.h" @@ -31,20 +30,18 @@ const LocationSpecGroup *FindLocationSpecGroup(const std::vector &capacity_gb, - int64_t block_charge_bytes, - std::vector &capacity_blocks) { - capacity_blocks.clear(); - capacity_blocks.reserve(capacity_gb.size()); +bool ConvertCapacitiesToBytes(const std::vector &capacity_gb, std::vector &capacity_bytes) { + capacity_bytes.clear(); + capacity_bytes.reserve(capacity_gb.size()); for (double capacity : capacity_gb) { if (!std::isfinite(capacity) || capacity < 0.0) { return false; } - const long double blocks = static_cast(capacity) * kBytesPerGb / block_charge_bytes; - if (blocks >= static_cast(std::numeric_limits::max())) { - capacity_blocks.push_back(std::numeric_limits::max()); + const long double bytes = static_cast(capacity) * kBytesPerGb; + if (bytes >= static_cast(std::numeric_limits::max())) { + capacity_bytes.push_back(static_cast(std::numeric_limits::max())); } else { - capacity_blocks.push_back(static_cast(blocks)); + capacity_bytes.push_back(static_cast(bytes)); } } return true; @@ -62,6 +59,34 @@ int64_t SaturatingMultiplyToInt64(uint64_t lhs, uint64_t rhs) { return static_cast(lhs * rhs); } +int64_t ClampToInt64(long double value) { + if (value >= static_cast(std::numeric_limits::max())) { + return std::numeric_limits::max(); + } + return value <= 0.0L ? 0 : static_cast(value); +} + +// Block count of one capacity tier for the register response. +// +// full-attention (step_blocks == 0): exact, and the same value doubles as the +// LiteHit projection slot because every block costs exactly full_charge_bytes. +// +// linear: an ESTIMATE only. Full blocks and Linear states have different +// charges, so no single block count describes the byte-axis cache; hits are +// always decided on the byte axis. The estimate amortizes one Linear state over +// step_blocks blocks. +int64_t CapacityBlocksForResponse(uint64_t capacity_bytes, + int64_t full_charge_bytes, + int64_t linear_charge_bytes, + int32_t step_blocks) { + if (step_blocks <= 0) { + return ClampToInt64(capacity_bytes / static_cast(full_charge_bytes)); + } + const long double bytes_per_step = + static_cast(step_blocks) * full_charge_bytes + static_cast(linear_charge_bytes); + return ClampToInt64(static_cast(capacity_bytes) * step_blocks / bytes_per_step); +} + } // namespace int64_t OnlineOptimizerManager::ComputeSizeForGroup(const std::vector &specs, @@ -101,6 +126,11 @@ ErrorCode OnlineOptimizerManager::CreateInstanceGroup(const OptimizerInstanceGro if (!registry_manager_) { return EC_ERROR; } + std::string invalid_fields; + if (!instance_group.ValidateRequiredFields(invalid_fields)) { + KVCM_LOG_ERROR("CreateInstanceGroup failed: invalid group config %s", invalid_fields.c_str()); + return EC_BADARGS; + } std::lock_guard admin_guard(admin_ops_mutex_); return registry_manager_->CreateInstanceGroup(instance_group); } @@ -109,6 +139,11 @@ ErrorCode OnlineOptimizerManager::UpdateInstanceGroup(const OptimizerInstanceGro if (!registry_manager_) { return EC_ERROR; } + std::string invalid_fields; + if (!instance_group.ValidateRequiredFields(invalid_fields)) { + KVCM_LOG_ERROR("UpdateInstanceGroup failed: invalid group config %s", invalid_fields.c_str()); + return EC_BADARGS; + } std::lock_guard admin_guard(admin_ops_mutex_); if (HasActiveInstanceInGroup(instance_group.name()) || HasPersistedInstanceInGroup(instance_group.name())) { @@ -255,11 +290,29 @@ ErrorCode OnlineOptimizerManager::RegisterInstanceInternal(const OptimizerInstan KVCM_LOG_ERROR("RegisterInstance failed: negative linear_step for instance[%s]", instance_id.c_str()); return EC_BADARGS; } - int32_t linear_step = instance_info.linear_step(); + const int32_t linear_step = instance_info.linear_step(); if (instance_info.block_size() <= 0) { KVCM_LOG_ERROR("RegisterInstance failed: non-positive token block_size for instance[%s]", instance_id.c_str()); return EC_BADARGS; } + // linear_step counts tokens; Linear states can only land on complete block + // boundaries, so it must divide evenly into whole blocks. + if (linear_step > 0 && linear_step % instance_info.block_size() != 0) { + KVCM_LOG_ERROR( + "RegisterInstance failed: linear_step=%d tokens is not a multiple of block_size=%d for instance[%s]", + linear_step, + instance_info.block_size(), + instance_id.c_str()); + return EC_BADARGS; + } + const int32_t linear_step_blocks = linear_step / instance_info.block_size(); + std::string invalid_group_fields; + if (!instance_group.ValidateRequiredFields(invalid_group_fields)) { + KVCM_LOG_ERROR("RegisterInstance failed: invalid group config %s for instance[%s]", + invalid_group_fields.c_str(), + instance_id.c_str()); + return EC_BADARGS; + } if (instance_group.eviction_policy() != "lru") { KVCM_LOG_ERROR("RegisterInstance failed: unsupported eviction_policy[%s] for instance[%s]", instance_group.eviction_policy().c_str(), @@ -285,6 +338,17 @@ ErrorCode OnlineOptimizerManager::RegisterInstanceInternal(const OptimizerInstan instance_id.c_str()); return EC_BADARGS; } + if (linear_step == 0 && !optimizer_state_info.linear_location_spec_group_name().empty()) { + KVCM_LOG_ERROR( + "RegisterInstance failed: full-attention instance[%s] must not configure a linear/Mamba spec group", + instance_id.c_str()); + return EC_BADARGS; + } + if (linear_step > 0 && optimizer_state_info.linear_location_spec_group_name().empty()) { + KVCM_LOG_ERROR("RegisterInstance failed: linear instance[%s] requires a non-empty Mamba spec group", + instance_id.c_str()); + return EC_BADARGS; + } const auto *full_group = FindLocationSpecGroup(groups, optimizer_state_info.full_location_spec_group_name()); if (!full_group) { @@ -293,16 +357,16 @@ ErrorCode OnlineOptimizerManager::RegisterInstanceInternal(const OptimizerInstan instance_id.c_str()); return EC_BADARGS; } - int64_t size_full_only = ComputeSizeForGroup(specs, *full_group); - if (size_full_only <= 0) { + int64_t full_charge_bytes = ComputeSizeForGroup(specs, *full_group); + if (full_charge_bytes <= 0) { KVCM_LOG_ERROR("RegisterInstance failed: invalid full group[%s] size[%ld] for instance[%s]", full_group->name().c_str(), - size_full_only, + full_charge_bytes, instance_id.c_str()); return EC_BADARGS; } - int64_t size_full_linear = size_full_only; + int64_t linear_charge_bytes = 0; if (!optimizer_state_info.linear_location_spec_group_name().empty()) { const auto *linear_group = FindLocationSpecGroup(groups, optimizer_state_info.linear_location_spec_group_name()); @@ -312,67 +376,48 @@ ErrorCode OnlineOptimizerManager::RegisterInstanceInternal(const OptimizerInstan instance_id.c_str()); return EC_BADARGS; } - const int64_t size_linear = ComputeSizeForGroup(specs, *linear_group); - if (size_linear <= 0) { + linear_charge_bytes = ComputeSizeForGroup(specs, *linear_group); + if (linear_charge_bytes <= 0) { KVCM_LOG_ERROR("RegisterInstance failed: invalid linear group[%s] size[%ld] for instance[%s]", linear_group->name().c_str(), - size_linear, + linear_charge_bytes, instance_id.c_str()); return EC_BADARGS; } - size_full_linear += size_linear; - } - - int64_t estimated_bytes_per_block; - if (linear_step == 0) { - estimated_bytes_per_block = size_full_only; - } else if (linear_step == 1) { - estimated_bytes_per_block = size_full_linear; - } else { - estimated_bytes_per_block = ((linear_step - 1) * size_full_only + size_full_linear) / linear_step; - } - - if (estimated_bytes_per_block <= 0) { - KVCM_LOG_ERROR("RegisterInstance failed: estimated_bytes_per_block <= 0 for instance[%s]", instance_id.c_str()); - return EC_BADARGS; } const auto &capacity_gb = instance_group.capacity_gb(); - std::vector estimated_capacity_blocks; - if (!ConvertCapacitiesToBlocks(capacity_gb, estimated_bytes_per_block, estimated_capacity_blocks)) { + std::vector capacity_bytes; + if (!ConvertCapacitiesToBytes(capacity_gb, capacity_bytes)) { KVCM_LOG_ERROR("RegisterInstance failed: invalid capacity for instance[%s]", instance_id.c_str()); return EC_BADARGS; } + std::vector estimated_capacity_blocks; + estimated_capacity_blocks.reserve(capacity_bytes.size()); + for (uint64_t bytes : capacity_bytes) { + estimated_capacity_blocks.push_back( + CapacityBlocksForResponse(bytes, full_charge_bytes, linear_charge_bytes, linear_step_blocks)); + } auto state = std::make_shared(); state->instance_info = std::make_shared(instance_info); state->instance_group = std::make_shared(instance_group); - state->size_full_only = size_full_only; - state->size_full_linear = size_full_linear; + state->full_charge_bytes = full_charge_bytes; + state->linear_charge_bytes = linear_charge_bytes; state->linear_step = linear_step; + state->linear_step_blocks = linear_step_blocks; state->total_hits_per_capacity.resize(capacity_gb.size(), 0); + state->capacity_bytes = capacity_bytes; - if (linear_step == 0) { - state->lite_hit_capacity_blocks = estimated_capacity_blocks; - // A group TTL is layered onto the LiteHit core; online time is the - // wall clock, mirroring the linear-path TtlCacheIndexerWrapper. - state->lite_hit = - std::make_unique(static_cast(instance_group.ttl_seconds()) * 1000000000ULL); - } else { - auto indexer = CacheIndexerFactory::CreateCacheIndexer(instance_group.eviction_policy(), - instance_group.enable_theoretical_max_cache(), - capacity_gb, - size_full_only, - size_full_linear, - linear_step, - instance_group.ttl_seconds()); - if (!indexer) { - KVCM_LOG_ERROR("RegisterInstance failed: initialize linear indexer for instance[%s]", instance_id.c_str()); - return EC_BADARGS; - } - state->indexer = std::move(indexer); + LiteHit::CacheObjectConfig object_config; + object_config.full_charge_bytes = static_cast(full_charge_bytes); + if (linear_step != 0) { + object_config.linear_charge_bytes = static_cast(linear_charge_bytes); + object_config.linear_step_blocks = static_cast(linear_step_blocks); } + const uint64_t ttl_ns = static_cast(instance_group.ttl_seconds()) * 1000000000ULL; + state->lite_hit = std::make_unique(object_config, ttl_ns); { std::unique_lock lock(instances_mutex_); @@ -380,15 +425,17 @@ ErrorCode OnlineOptimizerManager::RegisterInstanceInternal(const OptimizerInstan } result.estimated_capacity_blocks = estimated_capacity_blocks; - result.size_full_only = size_full_only; - result.size_full_linear = size_full_linear; - - KVCM_LOG_INFO("RegisterInstance OK: instance[%s] group[%s] linear_step=%d estimated_bytes_per_block=%ld caps=%zu", - instance_id.c_str(), - instance_info.instance_group_name().c_str(), - linear_step, - estimated_bytes_per_block, - estimated_capacity_blocks.size()); + result.full_charge_bytes = full_charge_bytes; + result.linear_charge_bytes = linear_charge_bytes; + + KVCM_LOG_INFO( + "RegisterInstance OK: instance[%s] group[%s] linear_step=%d full_charge=%ld linear_charge=%ld caps=%zu", + instance_id.c_str(), + instance_info.instance_group_name().c_str(), + linear_step, + full_charge_bytes, + linear_charge_bytes, + capacity_bytes.size()); return EC_OK; } @@ -482,8 +529,8 @@ ErrorCode OnlineOptimizerManager::TraceQuery(const std::string &instance_id, return EC_BADARGS; } - const RequestFact fact = - state->lite_hit->ProcessRequest(normalized.block_keys, TimestampUtil::GetCurrentTimeUs() * 1000); + const FullRequestFact fact = + state->lite_hit->ProcessFullRequest(normalized.block_keys, TimestampUtil::GetCurrentTimeUs() * 1000); result.input_token_len = ClampToInt64(normalized.input_token_len); const uint64_t block_size = static_cast(state->instance_info->block_size()); @@ -491,8 +538,8 @@ ErrorCode OnlineOptimizerManager::TraceQuery(const std::string &instance_id, result.hit_count_per_capacity.reserve(num_caps); result.hit_rate_per_capacity.reserve(num_caps); for (std::size_t i = 0; i < num_caps; ++i) { - const uint64_t hits = - HitCurveProjector::ProjectBlocks(fact, static_cast(state->lite_hit_capacity_blocks[i])); + const uint64_t hits = HitCurveProjector::ProjectFullBytes( + fact, state->capacity_bytes[i], static_cast(state->full_charge_bytes)); result.hit_count_per_capacity.push_back(ClampToInt64(hits)); result.hit_rate_per_capacity.push_back( normalized.input_token_len == 0 ? 0.0 : static_cast(hits * block_size) / token_denominator); @@ -500,14 +547,14 @@ ErrorCode OnlineOptimizerManager::TraceQuery(const std::string &instance_id, } const uint64_t unique_blocks = state->lite_hit->current_unique_blocks(); - result.unique_keys_per_capacity.reserve(state->lite_hit_capacity_blocks.size()); - for (int64_t capacity : state->lite_hit_capacity_blocks) { + result.unique_keys_per_capacity.reserve(state->capacity_bytes.size()); + for (uint64_t capacity_bytes : state->capacity_bytes) { result.unique_keys_per_capacity.push_back( - ClampToInt64(std::min(unique_blocks, static_cast(capacity)))); + ClampToInt64(state->lite_hit->FullObjectsWithinTotalBytes(capacity_bytes))); } if (state->instance_group->enable_theoretical_max_cache()) { - const uint64_t max_hits = HitCurveProjector::ProjectInfinite(fact); + const uint64_t max_hits = HitCurveProjector::ProjectFullInfinite(fact); result.max_hit_count = ClampToInt64(max_hits); result.max_hit_rate = normalized.input_token_len == 0 ? 0.0 : static_cast(max_hits * block_size) / token_denominator; @@ -526,44 +573,64 @@ ErrorCode OnlineOptimizerManager::TraceQuery(const std::string &instance_id, return EC_OK; } - std::vector hit_count; - int64_t max_hit_count; - if (!state->indexer) { - return EC_ERROR; - } - // Legacy analyzers keep their algorithm but share the same prefix-hash - // preprocessing switch. - if (state->instance_group->enable_prefix_hash()) { - state->indexer->ProcessKeys(ApplyPrefixHash(block_keys), hit_count, max_hit_count); - } else { - state->indexer->ProcessKeys(block_keys, hit_count, max_hit_count); - } + if (state->lite_hit) { + if (input_token_len < 0) { + return EC_BADARGS; + } - state->indexer->PostQueryMaintenance(); + NormalizedRequest normalized; + try { + normalized = NormalizeRequest(block_keys, + input_token_len, + static_cast(state->instance_info->block_size()), + state->instance_group->enable_prefix_hash()); + } catch (const std::invalid_argument &e) { + KVCM_LOG_ERROR( + "TraceQuery failed: invalid Mamba request for instance[%s]: %s", instance_id.c_str(), e.what()); + return EC_BADARGS; + } - state->total_queries++; - state->total_blocks_queried += total_blocks; - for (size_t j = 0; j < num_caps; j++) { - state->total_hits_per_capacity[j] += hit_count[j]; - } - if (max_hit_count >= 0) { - state->total_max_hits += max_hit_count; - } + const RequestFact fact = + state->lite_hit->ProcessRequest(normalized.block_keys, TimestampUtil::GetCurrentTimeUs() * 1000); + result.input_token_len = ClampToInt64(normalized.input_token_len); - hit_count.resize(num_caps); - result.hit_count_per_capacity = std::move(hit_count); - result.hit_rate_per_capacity.reserve(num_caps); - for (int64_t hits : result.hit_count_per_capacity) { - result.hit_rate_per_capacity.push_back(total_blocks > 0 ? static_cast(hits) / total_blocks : 0.0); + const uint64_t block_size = static_cast(state->instance_info->block_size()); + const double token_denominator = static_cast(normalized.input_token_len); + result.hit_count_per_capacity.reserve(num_caps); + result.hit_rate_per_capacity.reserve(num_caps); + for (std::size_t i = 0; i < num_caps; ++i) { + const uint64_t hits = HitCurveProjector::ProjectBytes(fact, state->capacity_bytes[i]); + result.hit_count_per_capacity.push_back(ClampToInt64(hits)); + result.hit_rate_per_capacity.push_back( + normalized.input_token_len == 0 ? 0.0 : static_cast(hits * block_size) / token_denominator); + state->total_hits_per_capacity[i] += static_cast(hits); + } + + result.unique_keys_per_capacity.reserve(state->capacity_bytes.size()); + for (uint64_t capacity_bytes : state->capacity_bytes) { + result.unique_keys_per_capacity.push_back( + ClampToInt64(state->lite_hit->FullObjectsWithinTotalBytes(capacity_bytes))); + } + + if (state->instance_group->enable_theoretical_max_cache()) { + const uint64_t max_hits = HitCurveProjector::ProjectInfinite(fact); + result.max_hit_count = ClampToInt64(max_hits); + result.max_hit_rate = + normalized.input_token_len == 0 ? 0.0 : static_cast(max_hits * block_size) / token_denominator; + result.theoretical_unique_keys = ClampToInt64(state->lite_hit->current_unique_blocks()); + state->total_max_hits += static_cast(max_hits); + } else { + result.max_hit_rate = -1.0; + result.theoretical_unique_keys = -1; + } + + state->total_queries++; + state->total_blocks_queried += total_blocks; + state->total_input_tokens += ClampToInt64(normalized.input_token_len); + return EC_OK; } - result.unique_keys_per_capacity = state->indexer->capacity_unique_counts(); - result.theoretical_unique_keys = max_hit_count >= 0 ? state->indexer->unique_count() : -1; - result.max_hit_count = max_hit_count; - result.max_hit_rate = total_blocks > 0 && max_hit_count >= 0 - ? static_cast(max_hit_count) / static_cast(total_blocks) - : 0.0; - return EC_OK; + return EC_ERROR; } ErrorCode OnlineOptimizerManager::ListInstances(const std::string &instance_group_filter, @@ -583,37 +650,36 @@ ErrorCode OnlineOptimizerManager::ListInstances(const std::string &instance_grou s.instance_group = state->instance_info->instance_group_name(); s.block_size = state->instance_info->block_size(); s.total_blocks_queried = state->total_blocks_queried; - s.bytes_per_block = - (state->linear_step == 0) - ? state->size_full_only - : ((state->linear_step - 1) * state->size_full_only + state->size_full_linear) / state->linear_step; + // Exact configured charge for full-attention. The Mamba branch below + // replaces it with the current resident working-set average. + s.bytes_per_block = static_cast(state->full_charge_bytes); s.linear_step = state->linear_step; + if (!state->lite_hit) { + continue; + } + // Summaries may arrive without traffic. Advance the shared TTL + // watermark for both Full-only and Linear instances so observability + // reflects the alive working set as of now. + state->lite_hit->AdvanceTime(static_cast(TimestampUtil::GetCurrentTimeUs()) * 1000); + s.total_queries = state->total_queries; + s.total_input_tokens = state->total_input_tokens; + s.ttl_eviction_count = ClampToInt64(state->lite_hit->ttl_expired_blocks()); + // LiteHit models an unbounded recency stack, so it has no capacity + // evictions; the total equals expired Full blocks in both modes. + s.eviction_count = s.ttl_eviction_count; + s.memory_usage_bytes = ClampToInt64(state->lite_hit->memory_usage_bytes()); + const auto &caps = state->instance_group->capacity_gb(); if (state->linear_step == 0) { - if (!state->lite_hit) { - continue; - } - s.total_queries = state->total_queries; - s.total_input_tokens = state->total_input_tokens; - // Summaries arrive without traffic; advance the TTL watermark so - // idle instances report the alive set as of now, not as of the - // last request. - state->lite_hit->AdvanceTime(static_cast(TimestampUtil::GetCurrentTimeUs()) * 1000); s.unique_keys = ClampToInt64(state->lite_hit->current_unique_blocks()); - s.ttl_eviction_count = ClampToInt64(state->lite_hit->ttl_expired_blocks()); - // Total-eviction contract: LiteHit has no capacity evictions, so - // the total equals the TTL expirations (the linear wrapper also - // counts harvested entries in both). - s.eviction_count = s.ttl_eviction_count; - s.memory_usage_bytes = ClampToInt64(state->lite_hit->memory_usage_bytes()); // Capacity-unbounded residency: without a TTL every distinct // block ever seen counts; with a group TTL only the alive working // set does. Finite tiers are min(U, C) of this same U and need no // separate report. s.kv_cache_usage_bytes = SaturatingMultiplyToInt64(state->lite_hit->current_unique_blocks(), - static_cast(state->size_full_only)); + static_cast(state->full_charge_bytes)); // Full-attention rates are token based: cumulative hit blocks are // converted to tokens with the fixed block size and divided by the @@ -641,45 +707,36 @@ ErrorCode OnlineOptimizerManager::ListInstances(const std::string &instance_grou s.max_hit_rate = -1.0; } } else { - if (!state->indexer) { - continue; - } - s.total_queries = state->total_queries; - s.total_max_hits = state->total_max_hits; - s.max_hit_rate = state->total_blocks_queried > 0 ? static_cast(s.total_max_hits) / - static_cast(state->total_blocks_queried) - : 0.0; - s.unique_keys = state->indexer->unique_count(); - s.eviction_count = state->indexer->eviction_count(); - s.memory_usage_bytes = state->indexer->memory_usage_bytes(); - s.kv_cache_usage_bytes = state->indexer->kv_cache_usage_bytes(); - s.ttl_eviction_count = state->indexer->ttl_eviction_count(); - - for (size_t i = 0; i < caps.size() && i < state->total_hits_per_capacity.size(); i++) { + // unique_keys intentionally counts only Full objects. + const uint64_t resident_full_blocks = state->lite_hit->current_unique_blocks(); + const uint64_t resident_bytes = state->lite_hit->resident_bytes(); + s.unique_keys = ClampToInt64(resident_full_blocks); + // Infinite-capacity working set, Full and Mamba bytes included. + s.kv_cache_usage_bytes = ClampToInt64(resident_bytes); + s.bytes_per_block = + resident_full_blocks == 0 + ? 0.0 + : static_cast(static_cast(resident_bytes) / resident_full_blocks); + + const double token_denominator = static_cast(state->total_input_tokens); + const int64_t block_size_tokens = state->instance_info->block_size(); + for (size_t i = 0; i < caps.size() && i < state->total_hits_per_capacity.size(); ++i) { PerCapacityHitRateInfo info; info.capacity_gb = caps[i]; info.total_hits = state->total_hits_per_capacity[i]; - info.hit_rate = state->total_blocks_queried > 0 ? static_cast(info.total_hits) / - static_cast(state->total_blocks_queried) - : 0.0; + info.hit_rate = state->total_input_tokens > 0 + ? static_cast(info.total_hits * block_size_tokens) / token_denominator + : 0.0; s.per_capacity_hit_rates.push_back(info); } - // Hit-age is a legacy indexer statistic and is intentionally not - // part of LiteHit's full-attention state. - auto age_buckets = state->indexer->GetHitAgeBuckets(); - int64_t bucket_total = 0; - for (const auto &bucket : age_buckets) { - bucket_total += bucket.hit_count; - } - int64_t age_denom = s.total_max_hits > 0 ? s.total_max_hits : bucket_total; - for (const auto &bucket : age_buckets) { - HitAgeBucketRatio ratio_info; - ratio_info.threshold_seconds = bucket.threshold_seconds; - ratio_info.hit_count = bucket.hit_count; - ratio_info.ratio = - age_denom > 0 ? static_cast(bucket.hit_count) / static_cast(age_denom) : 0.0; - s.hit_age_bucket_ratios.push_back(ratio_info); + if (state->instance_group->enable_theoretical_max_cache()) { + s.total_max_hits = state->total_max_hits; + s.max_hit_rate = state->total_input_tokens > 0 + ? static_cast(s.total_max_hits * block_size_tokens) / token_denominator + : 0.0; + } else { + s.max_hit_rate = -1.0; } } @@ -700,28 +757,10 @@ ErrorCode OnlineOptimizerManager::ResetStats(const std::string &instance_id) { } std::lock_guard guard(state->mutex); - if (state->linear_step == 0) { - if (!state->lite_hit) { - return EC_ERROR; - } - state->lite_hit->Reset(); - } else { - auto new_indexer = - CacheIndexerFactory::CreateCacheIndexer(state->instance_group->eviction_policy(), - state->instance_group->enable_theoretical_max_cache(), - state->instance_group->capacity_gb(), - state->size_full_only, - state->size_full_linear, - state->linear_step, - state->instance_group->ttl_seconds()); - if (!new_indexer) { - KVCM_LOG_ERROR("ResetStats failed: unsupported eviction_policy[%s] for instance[%s]", - state->instance_group->eviction_policy().c_str(), - instance_id.c_str()); - return EC_ERROR; - } - state->indexer = std::move(new_indexer); + if (!state->lite_hit) { + return EC_ERROR; } + state->lite_hit->Reset(); state->total_queries = 0; state->total_blocks_queried = 0; state->total_input_tokens = 0; diff --git a/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.h b/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.h index 6ffb3874f..cf159a17c 100644 --- a/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.h +++ b/kv_cache_manager/optimizer/manager/online_runtime/online_optimizer_manager.h @@ -12,8 +12,7 @@ #include "kv_cache_manager/common/error_code.h" #include "kv_cache_manager/optimizer/config/optimizer_instance_group.h" #include "kv_cache_manager/optimizer/config/optimizer_instance_info.h" -#include "kv_cache_manager/optimizer/index/online/cache_indexer.h" -#include "kv_cache_manager/optimizer/liteHit/lite_hit.h" +#include "kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h" namespace kv_cache_manager { @@ -23,17 +22,23 @@ struct InstanceState { std::shared_ptr instance_info; std::shared_ptr instance_group; - // Exactly one analyzer is active. Full-attention uses LiteHit directly; - // linear attention remains on the legacy CacheIndexer path. - std::unique_ptr lite_hit; - std::unique_ptr indexer; - // Static byte->block estimation of the configured capacities, kept for - // the existing response contract and used as projection slots. - std::vector lite_hit_capacity_blocks; - - int64_t size_full_only = 0; - int64_t size_full_linear = 0; + // TTL decorator around one shared LiteHit core. With ttl_seconds == 0 the + // decorator is a timestamp-free pass-through. + std::unique_ptr lite_hit; + // Configured capacity slots in bytes for both Full-only and Linear + // instances. Full RLE performs its block-size floor only at projection. + std::vector capacity_bytes; + + // Byte charges of the two object types. A Full block costs + // full_charge_bytes; a Linear state costs linear_charge_bytes + // (0 for full-attention instances). + int64_t full_charge_bytes = 0; + int64_t linear_charge_bytes = 0; + // Configured token interval (0 = full-attention only). Reported as-is. int32_t linear_step = 0; + // linear_step converted to whole blocks (linear_step / block_size); it is + // the Linear-state interval used by the Mamba policy. + int32_t linear_step_blocks = 0; std::mutex mutex; // Minimal cumulative integers per the existing contract. LiteHit itself @@ -59,9 +64,11 @@ struct TraceQueryResult { }; struct RegisterInstanceResult { + // Compatibility/display estimate derived from capacity_bytes. Runtime + // projection never reads this field. std::vector estimated_capacity_blocks; - int64_t size_full_only = 0; - int64_t size_full_linear = 0; + int64_t full_charge_bytes = 0; + int64_t linear_charge_bytes = 0; }; struct PerCapacityHitRateInfo { @@ -70,12 +77,6 @@ struct PerCapacityHitRateInfo { double hit_rate; }; -struct HitAgeBucketRatio { - int64_t threshold_seconds; // upper bound of this bucket (0 means "+inf") - int64_t hit_count; - double ratio; // hit_count / total_max_hits -}; - struct InstanceSummary { std::string instance_id; std::string instance_group; @@ -86,14 +87,15 @@ struct InstanceSummary { int64_t total_max_hits = 0; double max_hit_rate = 0.0; int64_t unique_keys = 0; - int64_t bytes_per_block = 0; + // Full-attention: exact configured Full block charge. Mamba: current + // resident Full+Mamba working-set bytes divided by resident Full blocks. + double bytes_per_block = 0.0; int32_t linear_step = 0; int64_t eviction_count = 0; int64_t memory_usage_bytes = 0; int64_t kv_cache_usage_bytes = 0; int64_t ttl_eviction_count = 0; std::vector per_capacity_hit_rates; - std::vector hit_age_bucket_ratios; }; class OnlineOptimizerManager { diff --git a/kv_cache_manager/optimizer/service/metrics/optimizer_metrics_reporter.cc b/kv_cache_manager/optimizer/service/metrics/optimizer_metrics_reporter.cc index 2de317a3d..8f7cf8e7c 100644 --- a/kv_cache_manager/optimizer/service/metrics/optimizer_metrics_reporter.cc +++ b/kv_cache_manager/optimizer/service/metrics/optimizer_metrics_reporter.cc @@ -48,7 +48,6 @@ struct OptimizerMetricsReporter::KmonContext { DECLARE_METRICS(query, capacity_efficiency); DECLARE_METRICS(trace, query_capacity_efficiency); - DECLARE_METRICS(trace, query_hit_age_bucket_ratio); struct MapHashFunc { size_t operator()(const std::map &m) const noexcept { @@ -218,8 +217,6 @@ bool OptimizerMetricsReporter::InitMetrics() { REGISTER_GAUGE_METRIC(query, max_hit_rate); REGISTER_GAUGE_METRIC(query, capacity_efficiency); - REGISTER_GAUGE_METRIC(trace, query_hit_age_bucket_ratio); - KVCM_LOG_INFO("OptimizerMetricsReporter: kmonitor initialized, prefix[%s]", prefix_.c_str()); return true; } @@ -311,14 +308,6 @@ void OptimizerMetricsReporter::ReportInterval() { achievement = cap_info.hit_rate / s.max_hit_rate; } } - - for (const auto &bucket : s.hit_age_bucket_ratios) { - std::string bucket_label = - bucket.threshold_seconds > 0 ? std::to_string(bucket.threshold_seconds) + "s" : "inf"; - MetricsTags bucket_tags = {{"instance_id", s.instance_id}, {"age_bucket", bucket_label}}; - Gauge bucket_ratio = metrics_registry_->GetGauge("trace_query_hit_age_bucket_ratio", bucket_tags); - bucket_ratio = bucket.ratio; - } } // --- Kmonitor --- @@ -353,14 +342,6 @@ void OptimizerMetricsReporter::ReportInterval() { kmon_ctx_->trace_query_capacity_efficiency_metrics->Report(&ktags, cap_info.hit_rate / s.max_hit_rate); } } - - for (const auto &bucket : s.hit_age_bucket_ratios) { - std::string bucket_label = - bucket.threshold_seconds > 0 ? std::to_string(bucket.threshold_seconds) + "s" : "inf"; - MetricsTags bucket_tags = {{"instance_id", s.instance_id}, {"age_bucket", bucket_label}}; - kmonitor::MetricsTags ktags = kmon_ctx_->GetKmonitorTags(bucket_tags); - kmon_ctx_->trace_query_hit_age_bucket_ratio_metrics->Report(&ktags, bucket.ratio); - } } } diff --git a/kv_cache_manager/optimizer/service/optimizer_service_impl.cc b/kv_cache_manager/optimizer/service/optimizer_service_impl.cc index a9ad4fa6c..c9114fd9c 100644 --- a/kv_cache_manager/optimizer/service/optimizer_service_impl.cc +++ b/kv_cache_manager/optimizer/service/optimizer_service_impl.cc @@ -232,8 +232,10 @@ void OptimizerServiceImpl::RegisterInstance(RequestContext *request_context, for (int64_t cap : result.estimated_capacity_blocks) { response->add_estimated_capacity_blocks(cap); } - response->set_size_full_only(result.size_full_only); - response->set_size_full_linear(result.size_full_linear); + // The response contract keeps the fused wording: size_full_linear is + // the byte size of one Full block plus one Linear state. + response->set_size_full_only(result.full_charge_bytes); + response->set_size_full_linear(result.full_charge_bytes + result.linear_charge_bytes); } } diff --git a/kv_cache_manager/optimizer/test/BUILD b/kv_cache_manager/optimizer/test/BUILD index 3fa3104b2..2edea698c 100644 --- a/kv_cache_manager/optimizer/test/BUILD +++ b/kv_cache_manager/optimizer/test/BUILD @@ -145,41 +145,55 @@ py_test( ) cc_test( - name = "LruCacheIndexerTest", + name = "LiteHitTest", srcs = [ - "lru_cache_indexer_test.cc", + "lite_hit_test.cc", ], copts = ["-fno-access-control"], deps = [ "//kv_cache_manager/common:unittest", - "//kv_cache_manager/optimizer/index:cache_indexer", + "//kv_cache_manager/optimizer/liteHit:hit_curve", + "//kv_cache_manager/optimizer/liteHit:lite_hit", + "//kv_cache_manager/optimizer/liteHit:request_preprocess", ], ) cc_test( - name = "LiteHitTest", + name = "LiteHitTtlTest", srcs = [ - "lite_hit_test.cc", + "lite_hit_ttl_test.cc", ], copts = ["-fno-access-control"], deps = [ "//kv_cache_manager/common:unittest", "//kv_cache_manager/optimizer/liteHit:hit_curve", - "//kv_cache_manager/optimizer/liteHit:lite_hit", - "//kv_cache_manager/optimizer/liteHit:request_preprocess", + "//kv_cache_manager/optimizer/liteHit:lite_hit_ttl", ], ) cc_test( - name = "LiteHitTtlTest", + name = "WeightedLruPoolTest", srcs = [ - "lite_hit_ttl_test.cc", + "weighted_lru_pool_test.cc", + ], + copts = ["-fno-access-control"], + deps = [ + "//kv_cache_manager/common:unittest", + "//kv_cache_manager/optimizer/liteHit:weighted_lru_pool", + ], +) + +cc_test( + name = "LiteHitLinearTest", + srcs = [ + "lite_hit_linear_test.cc", ], copts = ["-fno-access-control"], deps = [ "//kv_cache_manager/common:unittest", "//kv_cache_manager/optimizer/liteHit:hit_curve", "//kv_cache_manager/optimizer/liteHit:lite_hit", + "//kv_cache_manager/optimizer/liteHit:lite_hit_ttl", ], ) @@ -200,20 +214,6 @@ cc_test( ], ) -cc_test( - name = "TtlCacheIndexerWrapperTest", - srcs = [ - "ttl_cache_indexer_wrapper_test.cc", - ], - copts = ["-fno-access-control"], - deps = [ - "//kv_cache_manager/common:unittest", - "//kv_cache_manager/optimizer/index:cache_indexer", - "//kv_cache_manager/optimizer/index:cache_indexer_factory", - "//kv_cache_manager/optimizer/index:ttl_cache_indexer_wrapper", - ], -) - cc_test( name = "OnlineOptimizerManagerTest", srcs = [ diff --git a/kv_cache_manager/optimizer/test/lite_hit_linear_test.cc b/kv_cache_manager/optimizer/test/lite_hit_linear_test.cc new file mode 100644 index 000000000..e2405d24b --- /dev/null +++ b/kv_cache_manager/optimizer/test/lite_hit_linear_test.cc @@ -0,0 +1,289 @@ +#include +#include +#include +#include + +#include "kv_cache_manager/common/unittest.h" +#include "kv_cache_manager/optimizer/liteHit/hit_curve.h" +#include "kv_cache_manager/optimizer/liteHit/lite_hit.h" +#include "kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h" + +namespace kv_cache_manager { + +namespace { + +LiteHit::CacheObjectConfig +LinearObjectConfig(uint64_t full_charge_bytes, uint64_t linear_charge_bytes, uint64_t linear_step_blocks) { + LiteHit::CacheObjectConfig object_config; + object_config.full_charge_bytes = full_charge_bytes; + object_config.linear_charge_bytes = linear_charge_bytes; + object_config.linear_step_blocks = linear_step_blocks; + return object_config; +} + +// Naive reference of the Mamba recovery semantics: unbounded recency lists +// with per-type charges; residency for a capacity is decided by +// required-bytes thresholds (Mattson inclusion), evaluated on the +// request-start snapshot; the commit set is fixed and tail-to-head. +class NaiveMambaOracle { +public: + NaiveMambaOracle(const LiteHit::CacheObjectConfig &object_config) : object_config_(object_config) {} + + // Hits of `keys` for every capacity in `capacities`, then commit. + std::vector ProcessRequest(const std::vector &keys, const std::vector &capacities) { + std::vector hits; + hits.reserve(capacities.size()); + for (const uint64_t capacity : capacities) { + hits.push_back(EvaluateSnapshot(keys, capacity)); + } + Commit(keys); + return hits; + } + +private: + struct Entry { + CacheObjectType type; + int64_t key; + }; + using Recency = std::list; // front = most recent + + bool ShouldWriteLinear(std::size_t position, std::size_t total) const { + return (position + 1) % object_config_.linear_step_blocks == 0 || position == total - 1; + } + + uint64_t Charge(CacheObjectType type) const { + return type == CacheObjectType::kFull ? object_config_.full_charge_bytes : object_config_.linear_charge_bytes; + } + + bool RequiredBytes(const Recency &list, CacheObjectType type, int64_t key, uint64_t &required) const { + uint64_t newer = 0; + for (const Entry &entry : list) { + if (entry.type == type && entry.key == key) { + required = newer + Charge(type); + return true; + } + newer += Charge(entry.type); + } + return false; + } + + uint64_t EvaluateSnapshot(const std::vector &keys, uint64_t total_capacity) const { + std::size_t covered = 0; + for (std::size_t i = 0; i < keys.size(); ++i) { + uint64_t required = 0; + if (!RequiredBytes(recency_, CacheObjectType::kFull, keys[i], required) || required > total_capacity) { + break; + } + covered = i + 1; + } + uint64_t best = 0; + for (std::size_t p = 0; p < covered; ++p) { + uint64_t required = 0; + if (RequiredBytes(recency_, CacheObjectType::kLinear, keys[p], required) && required <= total_capacity) { + best = p + 1; + } + } + return best; + } + + void Touch(CacheObjectType type, int64_t key) { + recency_.remove_if([&](const Entry &entry) { return entry.type == type && entry.key == key; }); + recency_.push_front({type, key}); + } + + void Commit(const std::vector &keys) { + for (std::size_t i = keys.size(); i > 0; --i) { + const std::size_t position = i - 1; + if (ShouldWriteLinear(position, keys.size())) { + Touch(CacheObjectType::kLinear, keys[position]); + } + Touch(CacheObjectType::kFull, keys[position]); + } + } + + LiteHit::CacheObjectConfig object_config_; + Recency recency_; +}; + +// Generates prefix-chained request key sequences: each request is a path of +// a random trie, so equal keys imply equal prefixes. +class ChainTraceGenerator { +public: + explicit ChainTraceGenerator(uint32_t seed) : rng_(seed) {} + + std::vector NextRequest(std::size_t max_len) { + std::vector keys; + if (!history_.empty() && std::uniform_int_distribution(0, 2)(rng_) != 0) { + const auto &base = history_[std::uniform_int_distribution(0, history_.size() - 1)(rng_)]; + const std::size_t take = std::uniform_int_distribution(0, base.size())(rng_); + keys.assign(base.begin(), base.begin() + take); + } + const std::size_t fresh = std::uniform_int_distribution(0, max_len)(rng_); + for (std::size_t i = 0; i < fresh; ++i) { + keys.push_back(next_key_++); + } + if (!keys.empty()) { + history_.push_back(keys); + } + return keys; + } + +private: + std::mt19937 rng_; + int64_t next_key_ = 1; + std::vector> history_; +}; + +void RunOracleComparison(const LiteHit::CacheObjectConfig &object_config, + const std::vector &capacities, + uint32_t seed, + int requests) { + LiteHit core(object_config); + NaiveMambaOracle oracle(object_config); + ChainTraceGenerator generator(seed); + + for (int r = 0; r < requests; ++r) { + const std::vector keys = generator.NextRequest(12); + const RequestFact fact = core.ProcessRequest(keys); + const std::vector expected = oracle.ProcessRequest(keys, capacities); + + // Envelope invariants: strictly increasing in both fields. + for (std::size_t i = 1; i < fact.points.size(); ++i) { + ASSERT_LT(fact.points[i - 1].min_total_capacity_bytes, fact.points[i].min_total_capacity_bytes); + ASSERT_LT(fact.points[i - 1].hit_blocks, fact.points[i].hit_blocks); + } + + for (std::size_t c = 0; c < capacities.size(); ++c) { + ASSERT_EQ(expected[c], HitCurveProjector::ProjectBytes(fact, capacities[c])) + << "request " << r << " capacity " << capacities[c]; + } + ASSERT_EQ(HitCurveProjector::ProjectBytes(fact, std::numeric_limits::max() / 4), + HitCurveProjector::ProjectInfinite(fact)); + } +} + +} // namespace + +class LiteHitLinearTest : public TESTBASE {}; + +TEST_F(LiteHitLinearTest, EmptyRequestYieldsEmptyFactAndNoState) { + LiteHit core(LinearObjectConfig(/*full=*/8, /*linear=*/4, /*step_blocks=*/2)); + const RequestFact fact = core.ProcessRequest({}); + EXPECT_TRUE(fact.points.empty()); + EXPECT_EQ(0u, core.current_unique_blocks()); + EXPECT_EQ(0u, core.resident_bytes()); +} + +TEST_F(LiteHitLinearTest, ColdRequestMissesThenRepeatsHit) { + LiteHit core(LinearObjectConfig(/*full=*/8, /*linear=*/4, /*step_blocks=*/2)); + + // 5 blocks, step 2 -> Linear states at positions 1, 3 and forced tail 4. + const std::vector keys = {10, 20, 30, 40, 50}; + const RequestFact cold = core.ProcessRequest(keys); + EXPECT_TRUE(cold.points.empty()); + EXPECT_EQ(5u, core.current_unique_blocks()); + // 5 Full * 8 + 3 Mamba * 4. + EXPECT_EQ(52u, core.resident_bytes()); + + const RequestFact warm = core.ProcessRequest(keys); + ASSERT_FALSE(warm.points.empty()); + // Unbounded capacity recovers the forced tail Linear state: all 5 blocks. + EXPECT_EQ(5u, HitCurveProjector::ProjectInfinite(warm)); + // Zero capacity recovers nothing. + EXPECT_EQ(0u, HitCurveProjector::ProjectBytes(warm, 0)); +} + +TEST_F(LiteHitLinearTest, RecoveryRequiresBothFullCoverageAndLinearState) { + // One Linear state per block keeps the arithmetic small. + LiteHit core(LinearObjectConfig(/*full=*/10, /*linear=*/1, /*step_blocks=*/1)); + + core.ProcessRequest({1, 2, 3}); + // Recency (new->old): F1 M1 F2 M2 F3 M3; required bytes: + // F1=10, M1=11, F2=21, M2=22, F3=32, M3=33. + const RequestFact fact = core.ProcessRequest({1, 2, 3}); + // Recover position 0 needs F1 and M1 -> 11; position 1 adds F2/M2 -> 22; + // position 2 -> 33. + ASSERT_EQ(3u, fact.points.size()); + EXPECT_EQ(ByteStepPoint({11, 1}), fact.points[0]); + EXPECT_EQ(ByteStepPoint({22, 2}), fact.points[1]); + EXPECT_EQ(ByteStepPoint({33, 3}), fact.points[2]); + EXPECT_EQ(0u, HitCurveProjector::ProjectBytes(fact, 10)); + EXPECT_EQ(1u, HitCurveProjector::ProjectBytes(fact, 11)); + EXPECT_EQ(2u, HitCurveProjector::ProjectBytes(fact, 32)); + EXPECT_EQ(3u, HitCurveProjector::ProjectBytes(fact, 33)); +} + +TEST_F(LiteHitLinearTest, HistoricalForcedTailLinearStateRestoresLongerRequest) { + LiteHit core(LinearObjectConfig(/*full=*/10, /*linear=*/1, /*step_blocks=*/3)); + + // Position 1 is written only because it is this shorter request's tail. + EXPECT_TRUE(core.ProcessRequest({1, 2}).points.empty()); + + // In the longer request position 1 is neither periodic nor the tail, but + // its resident historical Linear object is still a valid restore point. + const RequestFact fact = core.ProcessRequest({1, 2, 3, 4}); + ASSERT_EQ(1u, fact.points.size()); + EXPECT_EQ(ByteStepPoint({21, 2}), fact.points[0]); + EXPECT_EQ(0u, HitCurveProjector::ProjectBytes(fact, 20)); + EXPECT_EQ(2u, HitCurveProjector::ProjectBytes(fact, 21)); + EXPECT_EQ(2u, HitCurveProjector::ProjectInfinite(fact)); +} + +TEST_F(LiteHitLinearTest, FullAndLinearStateExpireAtStrictTtlBoundary) { + TtlLiteHit core(LinearObjectConfig(/*full=*/10, /*linear=*/2, /*step_blocks=*/2), /*ttl_ns=*/100); + EXPECT_TRUE(core.ProcessRequest({1, 2}, 0).points.empty()); + + const RequestFact warm = core.ProcessRequest({1, 2}, 99); + EXPECT_EQ((std::vector{{22, 2}}), warm.points); + + core.AdvanceTime(199); + EXPECT_EQ(0u, core.current_unique_blocks()); + EXPECT_EQ(0u, core.resident_bytes()); + EXPECT_EQ(2u, core.ttl_expired_blocks()); + + // An expired access is cold, then revives both Full and Linear state. + EXPECT_TRUE(core.ProcessRequest({1, 2}, 199).points.empty()); + EXPECT_EQ(2u, core.current_unique_blocks()); + EXPECT_EQ(22u, core.resident_bytes()); +} + +TEST_F(LiteHitLinearTest, HistoricalTailLinearStateExpiresIndependently) { + TtlLiteHit core(LinearObjectConfig(/*full=*/10, /*linear=*/1, /*step_blocks=*/3), /*ttl_ns=*/100); + core.ProcessRequest({1, 2}, 0); + + // At t=50, the shorter request's tail state is still a valid restore point. + const RequestFact alive = core.ProcessRequest({1, 2, 3, 4}, 50); + ASSERT_EQ(1u, alive.points.size()); + EXPECT_EQ(ByteStepPoint({21, 2}), alive.points[0]); + + // At t=100 that old Linear state reaches its deadline. Full 1/2 were + // refreshed at t=50 and remain alive, but they cannot restore without it. + const RequestFact expired = core.ProcessRequest({1, 2}, 100); + EXPECT_TRUE(expired.points.empty()); +} + +TEST_F(LiteHitLinearTest, SharedOracleComparison) { + RunOracleComparison(LinearObjectConfig(/*full=*/48, /*linear=*/16, /*step_blocks=*/3), + {0, 16, 48, 64, 100, 160, 320, 640, 1000, 5000, 100000}, + 20260701, + 600); +} + +TEST_F(LiteHitLinearTest, SharedStepOneOracleComparison) { + RunOracleComparison(LinearObjectConfig(/*full=*/7, /*linear=*/13, /*step_blocks=*/1), + {0, 7, 13, 20, 39, 77, 200, 1000, 40000}, + 20260702, + 600); +} + +TEST_F(LiteHitLinearTest, ResetClearsEverything) { + LiteHit core(LinearObjectConfig(/*full=*/8, /*linear=*/4, /*step_blocks=*/2)); + core.ProcessRequest({1, 2, 3, 4}); + EXPECT_GT(core.resident_bytes(), 0u); + core.Reset(); + EXPECT_EQ(0u, core.current_unique_blocks()); + EXPECT_EQ(0u, core.resident_bytes()); + EXPECT_TRUE(core.ProcessRequest({1, 2}).points.empty()); +} + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/test/lite_hit_offline_runner_test.cc b/kv_cache_manager/optimizer/test/lite_hit_offline_runner_test.cc index 68f770450..11ce8aac2 100644 --- a/kv_cache_manager/optimizer/test/lite_hit_offline_runner_test.cc +++ b/kv_cache_manager/optimizer/test/lite_hit_offline_runner_test.cc @@ -39,6 +39,22 @@ class LiteHitOfflineRunnerTest : public TESTBASE { OptimizerStateInfo("full", "")); } + // Hybrid full+linear instance: full charge 16384, mamba charge 4096. + static OptimizerInstanceInfo + MakeHybridInfo(const std::string &instance_id, int32_t block_size = 4, int32_t linear_step_tokens = 12) { + return OptimizerInstanceInfo( + "g1", + instance_id, + block_size, + {LocationSpecInfo("tp0_F0", 8192), + LocationSpecInfo("tp1_F0", 8192), + LocationSpecInfo("tp0_L1", 2048), + LocationSpecInfo("tp1_L1", 2048)}, + {LocationSpecGroup("F0", {"tp0_F0", "tp1_F0"}), LocationSpecGroup("L1", {"tp0_L1", "tp1_L1"})}, + linear_step_tokens, + OptimizerStateInfo("F0", "L1")); + } + static std::string TraceLine(const std::string &instance_id, const std::string &trace_id, int64_t timestamp_ns, @@ -97,7 +113,8 @@ TEST_F(LiteHitOfflineRunnerTest, FactsCsvRowRoundTrips) { record.input_token_len = 900; record.block_size_tokens = 256; record.block_bytes = 4194304; - record.fact.hit_curve = {{1, 2}, {4, 1}}; + record.is_full_rle = true; + record.full_rle_fact.hit_curve = {{1, 2}, {4, 1}}; const std::string row = SerializeLiteHitFactRow(record); LiteHitFactRecord parsed; @@ -109,12 +126,13 @@ TEST_F(LiteHitOfflineRunnerTest, FactsCsvRowRoundTrips) { EXPECT_EQ(record.input_token_len, parsed.input_token_len); EXPECT_EQ(record.block_size_tokens, parsed.block_size_tokens); EXPECT_EQ(record.block_bytes, parsed.block_bytes); - EXPECT_EQ(record.fact.hit_curve, parsed.fact.hit_curve); + EXPECT_TRUE(parsed.is_full_rle); + EXPECT_EQ(record.full_rle_fact.hit_curve, parsed.full_rle_fact.hit_curve); - record.fact.hit_curve.clear(); + record.full_rle_fact.hit_curve.clear(); LiteHitFactRecord parsed_empty; ASSERT_TRUE(ParseLiteHitFactRow(SerializeLiteHitFactRow(record), parsed_empty, error)) << error; - EXPECT_TRUE(parsed_empty.fact.hit_curve.empty()); + EXPECT_TRUE(parsed_empty.full_rle_fact.hit_curve.empty()); LiteHitFactRecord bad; EXPECT_FALSE(ParseLiteHitFactRow("a,b,c", bad, error)); @@ -145,19 +163,19 @@ TEST_F(LiteHitOfflineRunnerTest, TtlGroupLayersTtlOntoTheHitCurve) { std::string error; LiteHitFactRecord r1; ASSERT_TRUE(ParseLiteHitFactRow(lines[1], r1, error)) << error; - EXPECT_TRUE(r1.fact.hit_curve.empty()); // cold + EXPECT_TRUE(r1.full_rle_fact.hit_curve.empty()); // cold LiteHitFactRecord r2; ASSERT_TRUE(ParseLiteHitFactRow(lines[2], r2, error)) << error; // Blocks 1,2 are 2s old (alive); 9 was never seen. - EXPECT_EQ((std::vector{{1, 2}}), r2.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 2}}), r2.full_rle_fact.hit_curve); LiteHitFactRecord r3; ASSERT_TRUE(ParseLiteHitFactRow(lines[3], r3, error)) << error; // Blocks 1,2 refreshed by r2 (1s old); block 3 is 3s old: deadline // reached, the prefix stops there for every capacity. Without the TTL // this request would produce thresholds 1,2,4. - EXPECT_EQ((std::vector{{1, 2}}), r3.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 2}}), r3.full_rle_fact.hit_curve); // Capacity stays a query-time axis on top of the fixed TTL. const double capacity_1_block_gb = 16384.0 / (1024.0 * 1024.0 * 1024.0); @@ -181,6 +199,49 @@ TEST_F(LiteHitOfflineRunnerTest, RejectsNegativeTtlGroup) { EXPECT_FALSE(LiteHitOfflineRunner(config).Run()); } +TEST_F(LiteHitOfflineRunnerTest, ByteStepAndFullRleFactsCsvRowsRoundTrip) { + LiteHitFactRecord record; + record.trace_id = "m1"; + record.instance_id = "mamba-0"; + record.timestamp_ns = 1720000000001; + record.input_token_len = 700; + record.block_size_tokens = 128; + record.block_bytes = 65536; + record.fact.points = {{4096, 1}, {131072, 5}}; + + const std::string row = SerializeLiteHitFactRow(record); + EXPECT_NE(std::string::npos, row.find("bytes:")); + + LiteHitFactRecord parsed; + std::string error; + ASSERT_TRUE(ParseLiteHitFactRow(row, parsed, error)) << error; + EXPECT_FALSE(parsed.is_full_rle); + EXPECT_EQ(record.fact.points, parsed.fact.points); + EXPECT_TRUE(parsed.full_rle_fact.hit_curve.empty()); + + record.fact.points.clear(); + LiteHitFactRecord parsed_empty; + ASSERT_TRUE(ParseLiteHitFactRow(SerializeLiteHitFactRow(record), parsed_empty, error)) << error; + EXPECT_FALSE(parsed_empty.is_full_rle); + EXPECT_TRUE(parsed_empty.fact.points.empty()); + + // New Full-only rows are explicitly tagged RLE. + LiteHitFactRecord new_full; + new_full.is_full_rle = true; + new_full.full_rle_fact.hit_curve = {{1, 2}}; + EXPECT_NE(std::string::npos, SerializeLiteHitFactRow(new_full).find("rle:")); + + // Legacy unprefixed Full RLE and "mamba:" byte-step rows stay readable. + LiteHitFactRecord full_row; + ASSERT_TRUE(ParseLiteHitFactRow("t,i,1,2,3,4,\"[[1,2]]\"", full_row, error)) << error; + EXPECT_TRUE(full_row.is_full_rle); + EXPECT_EQ((std::vector{{1, 2}}), full_row.full_rle_fact.hit_curve); + LiteHitFactRecord legacy_byte_step; + ASSERT_TRUE(ParseLiteHitFactRow("t,i,1,2,3,4,\"mamba:[[4096,1]]\"", legacy_byte_step, error)) << error; + EXPECT_FALSE(legacy_byte_step.is_full_rle); + EXPECT_EQ((std::vector{{4096, 1}}), legacy_byte_step.fact.points); +} + TEST_F(LiteHitOfflineRunnerTest, PublishesFactsAndMatchesOnlineReplay) { const std::string trace_path = WriteTrace("facts_ok.jsonl", { @@ -205,7 +266,7 @@ TEST_F(LiteHitOfflineRunnerTest, PublishesFactsAndMatchesOnlineReplay) { EXPECT_EQ(4, r3.block_size_tokens); EXPECT_EQ(16384, r3.block_bytes); // Fork [1,2,9] interleaved key 9 between 2 and 3: thresholds 1,2,4. - EXPECT_EQ((std::vector{{1, 2}, {4, 1}}), r3.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 2}, {4, 1}}), r3.full_rle_fact.hit_curve); // Reprojection of the facts must match an online replay of the same // trace with the same capacities. @@ -241,6 +302,92 @@ TEST_F(LiteHitOfflineRunnerTest, PublishesFactsAndMatchesOnlineReplay) { EXPECT_EQ(5, online_infinite_hits); } +TEST_F(LiteHitOfflineRunnerTest, PublishesMambaFactsAndMatchesOnlineReplay) { + // block_size 4, linear_step 12 tokens -> Linear state every 3 blocks plus + // the forced last block. 17 tokens -> 4 complete blocks. + const std::string trace_path = WriteTrace("facts_mamba.jsonl", + { + TraceLine("m1", "r1", 1000, {1, 2, 3, 4}, 17), + TraceLine("m1", "r2", 2000, {1, 2, 3, 4}, 17), + }); + const std::string output_dir = GetTestTempRootPath() + "/mamba"; + ASSERT_EQ(0, ::system(("mkdir -p " + output_dir).c_str())); + OptimizerLiteHitConfig config = MakeConfig(trace_path, output_dir); + config.set_instances({MakeHybridInfo("m1")}); + ASSERT_TRUE(LiteHitOfflineRunner(config).Run()); + + const std::vector lines = ReadLines(output_dir + "/" + kLiteHitFactsFileName); + ASSERT_EQ(3, lines.size()); + std::string error; + LiteHitFactRecord r1; + ASSERT_TRUE(ParseLiteHitFactRow(lines[1], r1, error)) << error; + EXPECT_FALSE(r1.is_full_rle); + EXPECT_TRUE(r1.fact.points.empty()); // cold: no recoverable Linear state + EXPECT_EQ(16384, r1.block_bytes); // per-row charge stays the Full charge + + // Warm request. Shared-pool recency after r1 (oldest->newest): + // M3,F3,M2,F2,F1,F0. Linear state p=2: max(full prefix 3x16384, Mamba + // 3x16384+4096) = 53248 -> 3 blocks; p=3: 4x16384+2x4096 = 73728 -> 4. + LiteHitFactRecord r2; + ASSERT_TRUE(ParseLiteHitFactRow(lines[2], r2, error)) << error; + EXPECT_FALSE(r2.is_full_rle); + EXPECT_EQ((std::vector{{53248, 3}, {73728, 4}}), r2.fact.points); + + const double capacity_gb = 53248.0 / (1024.0 * 1024.0 * 1024.0); + const std::string query_log = output_dir + "/query.jsonl"; + ASSERT_TRUE(RunLiteHitFactsQuery(output_dir + "/" + kLiteHitFactsFileName, {capacity_gb, -1.0}, query_log, error)) + << error; + const std::vector query_lines = ReadLines(query_log); + ASSERT_EQ(4, query_lines.size()); // 2 requests + m1 summary + overall + EXPECT_NE(std::string::npos, query_lines[3].find("\"total_hit_blocks\":[3,4]")); + EXPECT_NE(std::string::npos, query_lines[3].find("\"total_input_tokens\":34")); + + // Same trace through the online manager must agree. + auto registry = std::make_shared(""); + OnlineOptimizerManager manager(registry); + OptimizerInstanceGroup group = MakeGroup(); + group.set_capacity_gb({capacity_gb}); + group.set_enable_theoretical_max_cache(true); + ASSERT_EQ(EC_OK, registry->CreateInstanceGroup(group)); + RegisterInstanceResult reg_result; + ASSERT_EQ(EC_OK, manager.RegisterInstance(MakeHybridInfo("m1"), reg_result)); + + TraceQueryResult first; + ASSERT_EQ(EC_OK, manager.TraceQuery("m1", {1, 2, 3, 4}, 17, first)); + EXPECT_EQ(0, first.hit_count_per_capacity.at(0)); + TraceQueryResult second; + ASSERT_EQ(EC_OK, manager.TraceQuery("m1", {1, 2, 3, 4}, 17, second)); + EXPECT_EQ(3, second.hit_count_per_capacity.at(0)); + EXPECT_EQ(4, second.max_hit_count); +} + +TEST_F(LiteHitOfflineRunnerTest, LinearInstanceHonorsGroupTtl) { + const std::string trace_path = WriteTrace("facts_mamba_ttl.jsonl", + { + TraceLine("m1", "r1", 1, {1}, 4), + TraceLine("m1", "r2", 1000000001, {1}, 4), + TraceLine("m1", "r3", 1000000002, {1}, 4), + }); + const std::string output_dir = GetTestTempRootPath() + "/mamba_ttl"; + ASSERT_EQ(0, ::system(("mkdir -p " + output_dir).c_str())); + OptimizerLiteHitConfig config = MakeConfig(trace_path, output_dir); + OptimizerInstanceGroup group = MakeGroup(); + group.set_ttl_seconds(1); + config.set_instance_groups({group}); + config.set_instances({MakeHybridInfo("m1")}); + ASSERT_TRUE(LiteHitOfflineRunner(config).Run()); + + const std::vector lines = ReadLines(output_dir + "/" + kLiteHitFactsFileName); + ASSERT_EQ(4u, lines.size()); + std::string error; + LiteHitFactRecord expired; + ASSERT_TRUE(ParseLiteHitFactRow(lines[2], expired, error)) << error; + EXPECT_TRUE(expired.fact.points.empty()); // strict TTL boundary + LiteHitFactRecord revived; + ASSERT_TRUE(ParseLiteHitFactRow(lines[3], revived, error)) << error; + EXPECT_EQ((std::vector{{20480, 1}}), revived.fact.points); +} + TEST_F(LiteHitOfflineRunnerTest, ParallelPipelineMatchesSerialOutput) { std::vector lines; for (int i = 0; i < 2000; ++i) { @@ -288,7 +435,7 @@ TEST_F(LiteHitOfflineRunnerTest, AppliesOverrideInstanceIdAndPrefixHash) { std::string error; ASSERT_TRUE(ParseLiteHitFactRow(lines[2], r2, error)) << error; EXPECT_EQ("service", r2.instance_id); - EXPECT_EQ((std::vector{{1, 2}}), r2.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 2}}), r2.full_rle_fact.hit_curve); } TEST_F(LiteHitOfflineRunnerTest, OverrideInstanceIdMustMatchConfiguredInstance) { @@ -346,7 +493,7 @@ TEST_F(LiteHitOfflineRunnerTest, IgnoresWriteEvents) { std::string error; ASSERT_TRUE(ParseLiteHitFactRow(lines[2], r2, error)) << error; // The write event neither produced a fact row nor touched the LRU. - EXPECT_EQ((std::vector{{1, 1}}), r2.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 1}}), r2.full_rle_fact.hit_curve); } TEST_F(LiteHitOfflineRunnerTest, FanoutSweepsMultipleBlockSizes) { @@ -378,7 +525,7 @@ TEST_F(LiteHitOfflineRunnerTest, FanoutSweepsMultipleBlockSizes) { ++bs4_rows; EXPECT_EQ(4, record.block_size_tokens); if (record.trace_id == "r2") { - EXPECT_EQ((std::vector{{1, 3}}), record.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 3}}), record.full_rle_fact.hit_curve); } } else { ++bs8_rows; @@ -386,7 +533,7 @@ TEST_F(LiteHitOfflineRunnerTest, FanoutSweepsMultipleBlockSizes) { EXPECT_EQ(8, record.block_size_tokens); if (record.trace_id == "r2") { // 13 tokens = 1 complete 8-token block (chained key at index 1). - EXPECT_EQ((std::vector{{1, 1}}), record.fact.hit_curve); + EXPECT_EQ((std::vector{{1, 1}}), record.full_rle_fact.hit_curve); } } } diff --git a/kv_cache_manager/optimizer/test/lite_hit_test.cc b/kv_cache_manager/optimizer/test/lite_hit_test.cc index 4308bd1e3..9dcdbc074 100644 --- a/kv_cache_manager/optimizer/test/lite_hit_test.cc +++ b/kv_cache_manager/optimizer/test/lite_hit_test.cc @@ -61,127 +61,136 @@ class NaivePrefixLru { uint64_t cumulative_hits_ = 0; }; -uint64_t Project(const RequestFact &fact, int64_t capacity_blocks) { +uint64_t Project(const FullRequestFact &fact, int64_t capacity_blocks) { if (capacity_blocks == kInfinite) { - return HitCurveProjector::ProjectInfinite(fact); + return HitCurveProjector::ProjectFullInfinite(fact); } - return HitCurveProjector::ProjectBlocks(fact, static_cast(capacity_blocks)); + return HitCurveProjector::ProjectFullBlocks(fact, static_cast(capacity_blocks)); } } // namespace TEST(HitCurveProjectorTest, EmptyCurveHitsNothing) { - const RequestFact fact; - EXPECT_EQ(0, HitCurveProjector::ProjectBlocks(fact, 0)); - EXPECT_EQ(0, HitCurveProjector::ProjectBlocks(fact, 1000000)); - EXPECT_EQ(0, HitCurveProjector::ProjectInfinite(fact)); + const FullRequestFact fact; + EXPECT_EQ(0, HitCurveProjector::ProjectFullBlocks(fact, 0)); + EXPECT_EQ(0, HitCurveProjector::ProjectFullBlocks(fact, 1000000)); + EXPECT_EQ(0, HitCurveProjector::ProjectFullInfinite(fact)); } TEST(HitCurveProjectorTest, ProjectsSegmentBoundaries) { - const RequestFact fact{{HitCurveSegment{1, 2}, HitCurveSegment{4, 1}}}; - EXPECT_EQ(0, HitCurveProjector::ProjectBlocks(fact, 0)); - EXPECT_EQ(1, HitCurveProjector::ProjectBlocks(fact, 1)); - EXPECT_EQ(2, HitCurveProjector::ProjectBlocks(fact, 2)); - EXPECT_EQ(2, HitCurveProjector::ProjectBlocks(fact, 3)); - EXPECT_EQ(3, HitCurveProjector::ProjectBlocks(fact, 4)); - EXPECT_EQ(3, HitCurveProjector::ProjectBlocks(fact, 100)); - EXPECT_EQ(3, HitCurveProjector::ProjectInfinite(fact)); + const FullRequestFact fact{{HitCurveSegment{1, 2}, HitCurveSegment{4, 1}}}; + EXPECT_EQ(0, HitCurveProjector::ProjectFullBlocks(fact, 0)); + EXPECT_EQ(1, HitCurveProjector::ProjectFullBlocks(fact, 1)); + EXPECT_EQ(2, HitCurveProjector::ProjectFullBlocks(fact, 2)); + EXPECT_EQ(2, HitCurveProjector::ProjectFullBlocks(fact, 3)); + EXPECT_EQ(3, HitCurveProjector::ProjectFullBlocks(fact, 4)); + EXPECT_EQ(3, HitCurveProjector::ProjectFullBlocks(fact, 100)); + EXPECT_EQ(3, HitCurveProjector::ProjectFullInfinite(fact)); } TEST(HitCurveProjectorTest, ByteProjectionFloorsCapacity) { - const RequestFact fact{{HitCurveSegment{2, 3}}}; + const FullRequestFact fact{{HitCurveSegment{2, 3}}}; constexpr uint64_t kBlockBytes = 4096; - EXPECT_EQ(0, HitCurveProjector::ProjectBytes(fact, 2 * kBlockBytes - 1, kBlockBytes)); - EXPECT_EQ(1, HitCurveProjector::ProjectBytes(fact, 2 * kBlockBytes, kBlockBytes)); - EXPECT_EQ(2, HitCurveProjector::ProjectBytes(fact, 4 * kBlockBytes - 1, kBlockBytes)); - EXPECT_EQ(3, HitCurveProjector::ProjectBytes(fact, 4 * kBlockBytes, kBlockBytes)); + EXPECT_EQ(0, HitCurveProjector::ProjectFullBytes(fact, 2 * kBlockBytes - 1, kBlockBytes)); + EXPECT_EQ(1, HitCurveProjector::ProjectFullBytes(fact, 2 * kBlockBytes, kBlockBytes)); + EXPECT_EQ(2, HitCurveProjector::ProjectFullBytes(fact, 4 * kBlockBytes - 1, kBlockBytes)); + EXPECT_EQ(3, HitCurveProjector::ProjectFullBytes(fact, 4 * kBlockBytes, kBlockBytes)); +} + +TEST(LiteHitTest, DefaultFactIsByteStepEvenForFullOnly) { + LiteHit::CacheObjectConfig object_config; + object_config.full_charge_bytes = 8; + LiteHit core(object_config); + EXPECT_TRUE(core.ProcessRequest({1, 2, 3}).points.empty()); + const RequestFact fact = core.ProcessRequest({1, 2, 3}); + EXPECT_EQ((std::vector{{8, 1}, {16, 2}, {24, 3}}), fact.points); } TEST(LiteHitTest, EmptyRequestIsNoOp) { LiteHit lite_hit; - const RequestFact fact = lite_hit.ProcessRequest({}); + const FullRequestFact fact = lite_hit.ProcessFullRequest({}); EXPECT_TRUE(fact.hit_curve.empty()); EXPECT_EQ(0, lite_hit.current_unique_blocks()); } TEST(LiteHitTest, ColdRequestProducesEmptyCurveButCommits) { LiteHit lite_hit; - const RequestFact cold = lite_hit.ProcessRequest({1, 2, 3}); + const FullRequestFact cold = lite_hit.ProcessFullRequest({1, 2, 3}); EXPECT_TRUE(cold.hit_curve.empty()); EXPECT_EQ(3, lite_hit.current_unique_blocks()); } TEST(LiteHitTest, ContiguousChainReplayIsOneSegment) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2, 3}); + lite_hit.ProcessFullRequest({1, 2, 3}); // Reverse commit puts the chain contiguously: depths 1, 2, 3. - const RequestFact fact = lite_hit.ProcessRequest({1, 2, 3}); + const FullRequestFact fact = lite_hit.ProcessFullRequest({1, 2, 3}); EXPECT_EQ((std::vector{{1, 3}}), fact.hit_curve); } TEST(LiteHitTest, InterleavingChainBreaksSegments) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2, 3}); + lite_hit.ProcessFullRequest({1, 2, 3}); // Fork [1, 2, 9] commits 9 between 2 and 3: LRU becomes [1, 2, 9, 3]. - lite_hit.ProcessRequest({1, 2, 9}); - const RequestFact fact = lite_hit.ProcessRequest({1, 2, 3}); + lite_hit.ProcessFullRequest({1, 2, 9}); + const FullRequestFact fact = lite_hit.ProcessFullRequest({1, 2, 3}); EXPECT_EQ((std::vector{{1, 2}, {4, 1}}), fact.hit_curve); } TEST(LiteHitTest, CurveStopsAtFirstColdKey) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2, 3}); - const RequestFact fact = lite_hit.ProcessRequest({1, 7, 3}); + lite_hit.ProcessFullRequest({1, 2, 3}); + const FullRequestFact fact = lite_hit.ProcessFullRequest({1, 7, 3}); EXPECT_EQ((std::vector{{1, 1}}), fact.hit_curve); } TEST(LiteHitTest, BlocksAfterFirstMissStillUpdateGlobalLru) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2}); - const RequestFact mixed = lite_hit.ProcessRequest({1, 3, 2}); + lite_hit.ProcessFullRequest({1, 2}); + const FullRequestFact mixed = lite_hit.ProcessFullRequest({1, 3, 2}); EXPECT_EQ((std::vector{{1, 1}}), mixed.hit_curve); // Key 2 was committed even though it was after that request's first miss. // It is the oldest of the three, so it needs the full capacity of 3. - const RequestFact next = lite_hit.ProcessRequest({2}); + const FullRequestFact next = lite_hit.ProcessFullRequest({2}); EXPECT_EQ((std::vector{{3, 1}}), next.hit_curve); } TEST(LiteHitTest, ReverseCommitKeepsChainHeadMostRecent) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2, 3}); + lite_hit.ProcessFullRequest({1, 2, 3}); - const RequestFact head = lite_hit.ProcessRequest({1}); + const FullRequestFact head = lite_hit.ProcessFullRequest({1}); EXPECT_EQ((std::vector{{1, 1}}), head.hit_curve); // After the head query the LRU is [1, 2, 3] again; the leaf needs full // capacity. - const RequestFact leaf = lite_hit.ProcessRequest({3}); + const FullRequestFact leaf = lite_hit.ProcessFullRequest({3}); EXPECT_EQ((std::vector{{3, 1}}), leaf.hit_curve); } TEST(LiteHitTest, RepeatedKeysEncodeMonotonicallyAndNeverOptimistically) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2}); + lite_hit.ProcessFullRequest({1, 2}); // Snapshot depths: 1 -> 1, 2 -> 2. Request [2, 2, 1, 2] violates the // prefix-hash contract; thresholds before the guard are [2, 2, 2, 2]. // The monotonic guard encodes strictly increasing thresholds [2,3,4,5], // which is pessimistic (never optimistic) versus the naive oracle. - const RequestFact fact = lite_hit.ProcessRequest({2, 2, 1, 2}); + const FullRequestFact fact = lite_hit.ProcessFullRequest({2, 2, 1, 2}); EXPECT_EQ((std::vector{{2, 4}}), fact.hit_curve); // Sequential final state is [2, 1] with 2 most recent. - const RequestFact next = lite_hit.ProcessRequest({2, 1}); + const FullRequestFact next = lite_hit.ProcessFullRequest({2, 1}); EXPECT_EQ((std::vector{{1, 2}}), next.hit_curve); } TEST(LiteHitTest, ResetClearsLruState) { LiteHit lite_hit; - lite_hit.ProcessRequest({1, 2}); + lite_hit.ProcessFullRequest({1, 2}); ASSERT_EQ(2, lite_hit.current_unique_blocks()); lite_hit.Reset(); EXPECT_EQ(0, lite_hit.current_unique_blocks()); - EXPECT_TRUE(lite_hit.ProcessRequest({1, 2}).hit_curve.empty()); + EXPECT_TRUE(lite_hit.ProcessFullRequest({1, 2}).hit_curve.empty()); } TEST(LiteHitTest, MatchesNaiveMultiCapacityOracleOnRandomContractTraces) { @@ -206,7 +215,7 @@ TEST(LiteHitTest, MatchesNaiveMultiCapacityOracleOnRandomContractTraces) { } const std::vector block_keys = ApplyPrefixHash(raw_keys); - const RequestFact fact = lite_hit.ProcessRequest(block_keys); + const FullRequestFact fact = lite_hit.ProcessFullRequest(block_keys); for (std::size_t i = 0; i < capacities.size(); ++i) { const uint64_t expected_hits = oracles[i].Process(block_keys); const uint64_t projected = Project(fact, capacities[i]); @@ -237,7 +246,7 @@ TEST(LiteHitTest, ProjectionIsNeverOptimisticOnNonContractTraces) { block_keys.push_back(static_cast(rng() % 17)); } - const RequestFact fact = lite_hit.ProcessRequest(block_keys); + const FullRequestFact fact = lite_hit.ProcessFullRequest(block_keys); for (std::size_t i = 0; i < capacities.size(); ++i) { const uint64_t expected_hits = oracles[i].Process(block_keys); const uint64_t projected = Project(fact, capacities[i]); @@ -253,23 +262,23 @@ TEST(LiteHitTest, ProjectionIsNeverOptimisticOnNonContractTraces) { TEST(LiteHitTest, CompactsDynamicPositionsWithoutChangingResults) { LiteHit lite_hit; for (int64_t i = 0; i < 20000; ++i) { - lite_hit.ProcessRequest({i % 7}); + lite_hit.ProcessFullRequest({i % 7}); } EXPECT_EQ(7, lite_hit.current_unique_blocks()); - EXPECT_LE(lite_hit.fenwick_.size(), 2 * lite_hit.last_positions_.size() + 4096); - const RequestFact fact = lite_hit.ProcessRequest({(20000 - 7) % 7}); + EXPECT_LE(lite_hit.pool_.position_count(), 2 * lite_hit.pool_.full_positions_.size() + 4096); + const FullRequestFact fact = lite_hit.ProcessFullRequest({(20000 - 7) % 7}); EXPECT_EQ((std::vector{{7, 1}}), fact.hit_curve); } TEST(LiteHitTest, RetainsAllActiveKeysWithoutCapacityPruning) { LiteHit lite_hit; for (int64_t key = 0; key < 10000; ++key) { - lite_hit.ProcessRequest({key}); + lite_hit.ProcessFullRequest({key}); } EXPECT_EQ(10000, lite_hit.current_unique_blocks()); - const RequestFact oldest = lite_hit.ProcessRequest({0}); + const FullRequestFact oldest = lite_hit.ProcessFullRequest({0}); EXPECT_EQ((std::vector{{10000, 1}}), oldest.hit_curve); } diff --git a/kv_cache_manager/optimizer/test/lite_hit_ttl_test.cc b/kv_cache_manager/optimizer/test/lite_hit_ttl_test.cc index 15960b03d..a4a95ca0b 100644 --- a/kv_cache_manager/optimizer/test/lite_hit_ttl_test.cc +++ b/kv_cache_manager/optimizer/test/lite_hit_ttl_test.cc @@ -8,7 +8,7 @@ #include "kv_cache_manager/common/unittest.h" #include "kv_cache_manager/optimizer/liteHit/hit_curve.h" -#include "kv_cache_manager/optimizer/liteHit/lite_hit.h" +#include "kv_cache_manager/optimizer/liteHit/lite_hit_ttl.h" namespace kv_cache_manager { @@ -19,7 +19,7 @@ namespace { // blocks are all seen, alive (age strictly below TTL) and within the top C // of the recency stack on the request-start snapshot. Commits touch // tail-to-head (chain head most recent) and refresh last_access for every -// block, matching the LiteHit contract and the online TTL wrapper. +// block, matching the LiteHit contract and the TTL decorator. class NaiveLruTtlOracle { public: uint64_t Evaluate(const std::vector &keys, int64_t now_ns, uint64_t ttl_ns, uint64_t capacity) const { @@ -80,115 +80,134 @@ class NaiveLruTtlOracle { class LiteHitTtlTest : public TESTBASE {}; TEST_F(LiteHitTtlTest, StrictDeadlineBoundary) { - LiteHit core(2000); - EXPECT_TRUE(core.ProcessRequest({1, 2, 3}, 1000).hit_curve.empty()); // cold + TtlLiteHit core(2000); + EXPECT_TRUE(core.ProcessFullRequest({1, 2, 3}, 1000).hit_curve.empty()); // cold // Ages 1999 < 2000: alive, normal LRU curve. - EXPECT_EQ((std::vector{{1, 3}}), core.ProcessRequest({1, 2, 3}, 2999).hit_curve); + EXPECT_EQ((std::vector{{1, 3}}), core.ProcessFullRequest({1, 2, 3}, 2999).hit_curve); // Ages exactly 2000: deadline reached, miss (matches the online wrapper). - EXPECT_TRUE(core.ProcessRequest({1, 2, 3}, 4999).hit_curve.empty()); + EXPECT_TRUE(core.ProcessFullRequest({1, 2, 3}, 4999).hit_curve.empty()); // The expired access still refreshed last_access. - EXPECT_EQ(3, HitCurveProjector::ProjectInfinite(core.ProcessRequest({1, 2, 3}, 5000))); + EXPECT_EQ(3, HitCurveProjector::ProjectFullInfinite(core.ProcessFullRequest({1, 2, 3}, 5000))); } TEST_F(LiteHitTtlTest, ExpiredBlockStopsThePrefixLikeAColdOne) { - LiteHit core(2500); - core.ProcessRequest({10, 11}, 1000); - core.ProcessRequest({10, 11, 12}, 3000); // all refreshed at 3000 + TtlLiteHit core(2500); + core.ProcessFullRequest({10, 11}, 1000); + core.ProcessFullRequest({10, 11, 12}, 3000); // all refreshed at 3000 // At 5600 every block is 2600 old: dead despite being LRU-resident. - EXPECT_TRUE(core.ProcessRequest({10, 11, 12}, 5600).hit_curve.empty()); + EXPECT_TRUE(core.ProcessFullRequest({10, 11, 12}, 5600).hit_curve.empty()); } TEST_F(LiteHitTtlTest, TtlZeroIsPureLru) { - LiteHit core; // default: no TTL, timestamps ignored - core.ProcessRequest({1, 2, 3}, 1000); - const RequestFact fact = core.ProcessRequest({1, 2, 3}, 1000000000000); + TtlLiteHit core; // default: no TTL, timestamps ignored + core.ProcessFullRequest({1, 2, 3}, 1000); + const FullRequestFact fact = core.ProcessFullRequest({1, 2, 3}, 1000000000000); EXPECT_EQ((std::vector{{1, 3}}), fact.hit_curve); } +TEST_F(LiteHitTtlTest, DisabledDecoratorMatchesBareLinearCore) { + LiteHit::CacheObjectConfig object_config; + object_config.full_charge_bytes = 10; + object_config.linear_charge_bytes = 2; + object_config.linear_step_blocks = 2; + LiteHit bare(object_config); + TtlLiteHit decorated(object_config, /*ttl_ns=*/0); + + const std::vector> requests = {{1, 2, 3}, {1, 2}, {1, 2, 3}, {4, 5}}; + int64_t now_ns = 1000; + for (const auto &request : requests) { + EXPECT_EQ(bare.ProcessRequest(request).points, decorated.ProcessRequest(request, now_ns).points); + EXPECT_EQ(bare.current_unique_blocks(), decorated.current_unique_blocks()); + EXPECT_EQ(bare.resident_bytes(), decorated.resident_bytes()); + now_ns += 1000000; + } + EXPECT_EQ(0, decorated.ttl_expired_blocks()); +} + TEST_F(LiteHitTtlTest, ResetClearsTtlState) { - LiteHit core(1000000); - core.ProcessRequest({1, 2}, 1000); + TtlLiteHit core(1000000); + core.ProcessFullRequest({1, 2}, 1000); EXPECT_EQ(2, core.current_unique_blocks()); core.Reset(); EXPECT_EQ(0, core.current_unique_blocks()); - EXPECT_TRUE(core.ProcessRequest({1, 2}, 2000).hit_curve.empty()); + EXPECT_TRUE(core.ProcessFullRequest({1, 2}, 2000).hit_curve.empty()); } TEST_F(LiteHitTtlTest, UniqueBlocksExcludeExpired) { - LiteHit core(1000); - core.ProcessRequest({1, 2, 3}, 0); + TtlLiteHit core(1000); + core.ProcessFullRequest({1, 2, 3}, 0); EXPECT_EQ(3, core.current_unique_blocks()); // All three expire; only the new key is alive even though the expired // table entries linger until compaction. - core.ProcessRequest({100}, 5000); + core.ProcessFullRequest({100}, 5000); EXPECT_EQ(1, core.current_unique_blocks()); } TEST_F(LiteHitTtlTest, CompactionDropsExpiredEntries) { - LiteHit core(1000); + TtlLiteHit core(1000); for (int64_t r = 0; r < 600; ++r) { std::vector keys; keys.reserve(10); for (int64_t i = 0; i < 10; ++i) { keys.push_back(1000000 + r * 100 + i); } - core.ProcessRequest(keys, 0); + core.ProcessFullRequest(keys, 0); } EXPECT_EQ(6000, core.current_unique_blocks()); // Every block expires; the next request pushes the state past the // compaction threshold and the dead entries are dropped, not carried. - core.ProcessRequest({1, 2, 3}, 5000); + core.ProcessFullRequest({1, 2, 3}, 5000); EXPECT_EQ(3, core.current_unique_blocks()); - EXPECT_EQ(3u, core.last_positions_.size()); + EXPECT_EQ(3u, core.core_.pool_.full_positions_.size()); // The bucket array shrinks with the entries instead of staying at the // 6000-key high-water mark. - EXPECT_LT(core.last_positions_.bucket_count(), 1000u); + EXPECT_LT(core.core_.pool_.full_positions_.bucket_count(), 1000u); // Semantics survive the cleanup: dropped keys stay cold, alive ones hit // (the re-committed cold key occupies MRU, shifting thresholds by one). - EXPECT_TRUE(core.ProcessRequest({1000000}, 5001).hit_curve.empty()); - EXPECT_EQ((std::vector{{2, 3}}), core.ProcessRequest({1, 2, 3}, 5001).hit_curve); + EXPECT_TRUE(core.ProcessFullRequest({1000000}, 5001).hit_curve.empty()); + EXPECT_EQ((std::vector{{2, 3}}), core.ProcessFullRequest({1, 2, 3}, 5001).hit_curve); } TEST_F(LiteHitTtlTest, CompactionCollapsesEmptyEpochs) { - LiteHit core(1000000000); // TTL far beyond the replayed horizon + TtlLiteHit core(1000000000); // TTL far beyond the replayed horizon // A single hot key with a distinct timestamp per request: every commit // moves the marker and leaves the previous epoch's range empty. Without // the compaction dedupe the deque would hold one epoch per request. for (int64_t t = 1; t <= 10000; ++t) { - core.ProcessRequest({7}, t); + core.ProcessFullRequest({7}, t); } - EXPECT_LE(core.position_epochs_.size(), 4200u); + EXPECT_LE(core.ttl_state_.position_epochs_.size(), 4200u); EXPECT_EQ(1, core.current_unique_blocks()); // Exactness survives the collapse: alive within TTL, dead at the strict // deadline measured from the true last access. - EXPECT_EQ((std::vector{{1, 1}}), core.ProcessRequest({7}, 10001).hit_curve); - EXPECT_TRUE(core.ProcessRequest({7}, 10001 + 1000000000).hit_curve.empty()); + EXPECT_EQ((std::vector{{1, 1}}), core.ProcessFullRequest({7}, 10001).hit_curve); + EXPECT_TRUE(core.ProcessFullRequest({7}, 10001 + 1000000000).hit_curve.empty()); } TEST_F(LiteHitTtlTest, AdvanceTimeRefreshesUniqueCount) { - LiteHit core(1000); - core.ProcessRequest({1, 2}, 0); + TtlLiteHit core(1000); + core.ProcessFullRequest({1, 2}, 0); EXPECT_EQ(2, core.current_unique_blocks()); core.AdvanceTime(999); // ages 999 < 1000: still alive EXPECT_EQ(2, core.current_unique_blocks()); core.AdvanceTime(1000); // strict deadline: dead without any request EXPECT_EQ(0, core.current_unique_blocks()); // Re-access after the observational advance still revives normally. - core.ProcessRequest({1}, 1000); + core.ProcessFullRequest({1}, 1000); EXPECT_EQ(1, core.current_unique_blocks()); - LiteHit pure_lru; - pure_lru.ProcessRequest({1, 2}, 0); + TtlLiteHit pure_lru; + pure_lru.ProcessFullRequest({1, 2}, 0); pure_lru.AdvanceTime(1000000000); // no TTL: a pure no-op EXPECT_EQ(2, pure_lru.current_unique_blocks()); } TEST_F(LiteHitTtlTest, CountsTtlExpiredBlocks) { - LiteHit core(1000); - core.ProcessRequest({1, 2, 3}, 0); + TtlLiteHit core(1000); + core.ProcessFullRequest({1, 2, 3}, 0); EXPECT_EQ(0, core.ttl_expired_blocks()); // All three reached the deadline; key 1 revives after being counted. - core.ProcessRequest({1}, 2000); + core.ProcessFullRequest({1}, 2000); EXPECT_EQ(3, core.ttl_expired_blocks()); // Repeated advances must not double count the already-swept markers. core.AdvanceTime(2500); @@ -210,9 +229,9 @@ TEST_F(LiteHitTtlTest, RandomizedMatchesNaiveJointOracle) { // One core per fixed TTL; commits are TTL-independent, so every core and // the single oracle share the same recency and last-access evolution. - std::vector> cores; + std::vector> cores; for (uint64_t ttl : ttls) { - cores.push_back(std::make_unique(ttl)); + cores.push_back(std::make_unique(ttl)); } NaiveLruTtlOracle oracle; // Naive harvest model per TTL: a tracked key leaves the alive set and @@ -245,11 +264,11 @@ TEST_F(LiteHitTtlTest, RandomizedMatchesNaiveJointOracle) { ++it; } } - const RequestFact fact = cores[t]->ProcessRequest(keys, now); + const FullRequestFact fact = cores[t]->ProcessFullRequest(keys, now); ASSERT_EQ(naive_expired[t], cores[t]->ttl_expired_blocks()) << "step " << step << " ttl " << ttls[t]; for (uint64_t capacity : capacities) { ASSERT_EQ(oracle.Evaluate(keys, now, ttls[t], capacity), - HitCurveProjector::ProjectBlocks(fact, capacity)) + HitCurveProjector::ProjectFullBlocks(fact, capacity)) << "step " << step << " ttl " << ttls[t] << " capacity " << capacity; } naive_alive[t].insert(keys.begin(), keys.end()); diff --git a/kv_cache_manager/optimizer/test/lru_cache_indexer_test.cc b/kv_cache_manager/optimizer/test/lru_cache_indexer_test.cc deleted file mode 100644 index 2c1532cf9..000000000 --- a/kv_cache_manager/optimizer/test/lru_cache_indexer_test.cc +++ /dev/null @@ -1,588 +0,0 @@ -#include - -#include "kv_cache_manager/common/unittest.h" -#include "kv_cache_manager/optimizer/index/online/lru_cache_indexer.h" - -namespace kv_cache_manager { - -class LruCacheIndexerTest : public TESTBASE {}; - -// ==================== Basic Init Tests ==================== - -TEST_F(LruCacheIndexerTest, InitWithSingleCapacity) { - LruCacheIndexer indexer(0); - indexer.Init({1.0}, 1024, 2048, 1); - EXPECT_EQ(0, indexer.unique_count()); - EXPECT_EQ(0, indexer.eviction_count()); -} - -TEST_F(LruCacheIndexerTest, InitWithMultipleCapacities) { - LruCacheIndexer indexer(0); - indexer.Init({0.5, 1.0, 2.0}, 1024, 2048, 1); - EXPECT_EQ(0, indexer.unique_count()); - EXPECT_EQ(0, indexer.eviction_count()); - EXPECT_EQ(0, indexer.memory_usage_bytes()); - EXPECT_EQ(0, indexer.kv_cache_usage_bytes()); -} - -TEST_F(LruCacheIndexerTest, InitWithLinearStep) { - LruCacheIndexer indexer(0); - indexer.Init({1.0}, 512, 1024, 4); - EXPECT_EQ(0, indexer.unique_count()); -} - -// ==================== First Access (All Miss) Tests ==================== - -TEST_F(LruCacheIndexerTest, FirstAccessIsMiss) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({1.0, 2.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({1}, hit_count, max_hit); - - EXPECT_EQ(2u, hit_count.size()); - EXPECT_EQ(0, hit_count[0]); - EXPECT_EQ(0, hit_count[1]); -} - -TEST_F(LruCacheIndexerTest, FirstAccessMultipleKeysAllMiss) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({1.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - EXPECT_EQ(1u, hit_count.size()); - EXPECT_EQ(0, hit_count[0]); -} - -// ==================== Repeated Access (Hit) Tests ==================== - -TEST_F(LruCacheIndexerTest, RepeatedAccessIsHit) { - LruCacheIndexer indexer(0); - // Use small charge per key so cache can comfortably hold entries - indexer.Init({1.0}, 1024, 1024, 1); - - std::vector hit_count; - int64_t max_hit; - - // First access - miss - indexer.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Second access - hit - indexer.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, RepeatedAccessMultipleKeys) { - LruCacheIndexer indexer(0); - // Use small charge per key so cache can comfortably hold entries - indexer.Init({1.0}, 1024, 1024, 1); - - std::vector hit_count; - int64_t max_hit; - - // Populate cache - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Re-access same keys - all hit - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); -} - -// ==================== Multiple Capacity Tiers ==================== - -TEST_F(LruCacheIndexerTest, MultipleCapacityTiersEviction) { - LruCacheIndexer indexer(0); - // tier0 capacity ~2.5MB (holds 2 keys of 1MB + metadata), cannot hold 3 - // tier1 capacity ~10MB (easily holds 5 keys of 1MB + metadata) - // Note: kFullChargeCacheMetadata adds ~72 bytes overhead per entry - constexpr int64_t kOneMB = 1024LL * 1024; - double tier0_gb = 2.5 * kOneMB / (1024.0 * 1024.0 * 1024.0); - double tier1_gb = 10.0 * kOneMB / (1024.0 * 1024.0 * 1024.0); - indexer.Init({tier0_gb, tier1_gb}, kOneMB, kOneMB, 1); - - std::vector hit_count; - int64_t max_hit; - - // Insert keys 1..5: tier0 can hold ~2, tier1 can hold all 5 - for (int64_t i = 1; i <= 5; i++) { - indexer.ProcessKeys({i}, hit_count, max_hit); - } - - // Re-access key 1: should be evicted from tier0 but still in tier1 - indexer.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // evicted from small cache - EXPECT_EQ(1, hit_count[1]); // still in larger cache -} - -// ==================== Linear Step and Boundary Key Tests ==================== - -TEST_F(LruCacheIndexerTest, LinearStepZeroAllFullAttention) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=0: all full attention, simple prefix matching with size_full_only charge. - indexer.Init({10.0}, kOneGB / 2, kOneGB, 0); - - std::vector hit_count; - int64_t max_hit; - - // First access: all miss - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Re-access: all hit, simple prefix = 3 - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); - - // Partial hit: key 4 not cached, prefix truncates at 3 - indexer.ProcessKeys({1, 2, 3, 4, 5}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, LinearStepGrouping) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=2: step_hit at pos1,3,5,... - // Step_hit keys have charge = size_full_linear - indexer.Init({10.0}, kOneGB / 2, kOneGB, 2); - - std::vector hit_count; - int64_t max_hit; - - // First access group {1, 2}: both miss, group invalid -> hit_count = 0 - indexer.ProcessKeys({1, 2}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Re-access group {1, 2}: both should hit with correct boundary charge - indexer.ProcessKeys({1, 2}, hit_count, max_hit); - EXPECT_EQ(2, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, LinearStepColdAllMissReturnsZeroForCapacitiesAndMaxCache) { - LruCacheIndexer indexer(true); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0, 20.0}, kOneGB / 2, kOneGB, 2); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({1, 2, 3, 4}, hit_count, max_hit); - - ASSERT_EQ(2u, hit_count.size()); - EXPECT_EQ(0, hit_count[0]); - EXPECT_EQ(0, hit_count[1]); - EXPECT_EQ(0, max_hit); -} - -TEST_F(LruCacheIndexerTest, LinearStepBoundaryChargeMismatch) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=3: groups of 3, boundary key is 3rd key - indexer.Init({10.0}, kOneGB / 2, kOneGB, 3); - - std::vector hit_count; - int64_t max_hit; - - // First access as linear_step=3: group {1,2,3} - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // all miss - - // Re-access same group: all hit - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, LinearStepPartialGroup) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=3: step_hit positions are pos2,5,8,... - // Last key of each request is also stored with size_full_linear_. - indexer.Init({10.0}, kOneGB / 2, kOneGB, 3); - - std::vector hit_count; - int64_t max_hit; - - // First access {1,2}: both miss - indexer.ProcessKeys({1, 2}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Re-access {1,2}: pos1 is last key → stored with size_full_linear_ → checkpoint at pos1 → reuse = 2 - indexer.ProcessKeys({1, 2}, hit_count, max_hit); - EXPECT_EQ(2, hit_count[0]); - - // Access {1,2,3,4,5}: keys 1,2 hit (checkpoint at pos1), key 3 miss at pos2 → prefix=2 - // last_checkpoint=1 within prefix → reuse = 2 - indexer.ProcessKeys({1, 2, 3, 4, 5}, hit_count, max_hit); - EXPECT_EQ(2, hit_count[0]); - - // Re-access {1,2,3,4,5}: all hit. - // Checkpoints at pos1(key2), pos2(key3:step_hit), pos4(key5:last) - // Rightmost checkpoint at pos4 → reuse = 5 - indexer.ProcessKeys({1, 2, 3, 4, 5}, hit_count, max_hit); - EXPECT_EQ(5, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, LinearStepMultipleGroups) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=2, capacity enough for all keys - indexer.Init({10.0}, kOneGB / 2, kOneGB, 2); - - std::vector hit_count; - int64_t max_hit; - - // First access: 2 groups {1,2} {3,4} - indexer.ProcessKeys({1, 2, 3, 4}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Re-access: both groups hit - indexer.ProcessKeys({1, 2, 3, 4}, hit_count, max_hit); - EXPECT_EQ(4, hit_count[0]); -} - -// ==================== Prefix Hit Count Semantics ==================== - -TEST_F(LruCacheIndexerTest, PrefixHitTruncatesOnMiss) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - // Populate keys 1 and 3, but NOT 2 - indexer.ProcessKeys({1}, hit_count, max_hit); - indexer.ProcessKeys({3}, hit_count, max_hit); - - // Access sequence {1, 2, 3}: key 1 hits, key 2 misses -> prefix = 1 - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, PrefixHitAllHit) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - // Populate - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - // All hit -> prefix = 3 - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); -} - -// ==================== unique_count Tracking ==================== - -TEST_F(LruCacheIndexerTest, UniqueCountTracksNewKeys) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, indexer.unique_count()); - - // Re-access doesn't increase unique_count - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, indexer.unique_count()); - - // New key increases unique_count - indexer.ProcessKeys({4}, hit_count, max_hit); - EXPECT_EQ(4, indexer.unique_count()); -} - -// ==================== PostQueryMaintenance and Eviction ==================== - -TEST_F(LruCacheIndexerTest, PostQueryMaintenanceUpdatesEviction) { - constexpr int64_t kOneMB = 1024LL * 1024; - LruCacheIndexer indexer(false); - double cap_gb = 3.5 * kOneMB / (1024.0 * 1024.0 * 1024.0); - indexer.Init({cap_gb}, kOneMB, kOneMB, 1); - - std::vector hit_count; - int64_t max_hit; - - // Insert 10 keys -> largest configured capacity should evict older entries. - for (int64_t i = 1; i <= 10; i++) { - indexer.ProcessKeys({i}, hit_count, max_hit); - } - - indexer.PostQueryMaintenance(); - // After maintenance, unique_count should equal actual cache occupancy - // and eviction_count should reflect evicted keys - EXPECT_GT(indexer.eviction_count(), 0); - EXPECT_LT(indexer.unique_count(), 10); -} - -TEST_F(LruCacheIndexerTest, PostQueryMaintenanceNoEvictionWithLargeCapacity) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - indexer.ProcessKeys({1, 2, 3, 4, 5}, hit_count, max_hit); - indexer.PostQueryMaintenance(); - - EXPECT_EQ(5, indexer.unique_count()); - EXPECT_EQ(0, indexer.eviction_count()); -} - -// ==================== RemoveKey Tests ==================== - -TEST_F(LruCacheIndexerTest, RemoveKeyExisting) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, indexer.unique_count()); - - EXPECT_TRUE(indexer.RemoveKey(2)); - EXPECT_EQ(2, indexer.unique_count()); - EXPECT_EQ(1, indexer.eviction_count()); - - // Removed key should miss on re-access - indexer.ProcessKeys({2}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, RemoveKeyNonExisting) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - indexer.ProcessKeys({1}, hit_count, max_hit); - EXPECT_FALSE(indexer.RemoveKey(999)); - EXPECT_EQ(1, indexer.unique_count()); - EXPECT_EQ(0, indexer.eviction_count()); -} - -TEST_F(LruCacheIndexerTest, RemoveKeyFromMultipleTiers) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({5.0, 10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - indexer.ProcessKeys({1, 2}, hit_count, max_hit); - EXPECT_TRUE(indexer.RemoveKey(1)); - - // Key 1 should miss on both tiers - indexer.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - EXPECT_EQ(0, hit_count[1]); -} - -// ==================== Theoretical Max Cache Tests ==================== - -TEST_F(LruCacheIndexerTest, TheoreticalMaxCacheDoesNotLimitTracking) { - LruCacheIndexer indexer(true); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - for (int64_t i = 1; i <= 5; i++) { - indexer.ProcessKeys({i}, hit_count, max_hit); - } - EXPECT_EQ(5, indexer.unique_count()); - - indexer.PostQueryMaintenance(); - // Theoretical max cache is an effectively unlimited baseline, not another bounded capacity. - EXPECT_EQ(5, indexer.unique_count()); - EXPECT_EQ(0, indexer.eviction_count()); -} - -TEST_F(LruCacheIndexerTest, TheoreticalMaxCacheDisabledUsesLargestCapacity) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneKB = 1024LL; - indexer.Init({10.0}, kOneKB, kOneKB, 1); - - std::vector hit_count; - int64_t max_hit; - - for (int64_t i = 1; i <= 100; i++) { - indexer.ProcessKeys({i}, hit_count, max_hit); - } - indexer.PostQueryMaintenance(); - EXPECT_EQ(100, indexer.unique_count()); - EXPECT_EQ(0, indexer.eviction_count()); -} - -// ==================== Memory and Usage Metrics ==================== - -TEST_F(LruCacheIndexerTest, MemoryUsageBytesGrows) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - EXPECT_EQ(0, indexer.memory_usage_bytes()); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - EXPECT_GT(indexer.memory_usage_bytes(), 0); -} - -TEST_F(LruCacheIndexerTest, KvCacheUsageBytesGrows) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - EXPECT_EQ(0, indexer.kv_cache_usage_bytes()); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - EXPECT_GT(indexer.kv_cache_usage_bytes(), 0); -} - -TEST_F(LruCacheIndexerTest, KvCacheUsageBytesReflectsCharge) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // Each key uses kOneGB as charge - indexer.Init({100.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({1}, hit_count, max_hit); - - // Usage should be approximately kOneGB (for 1 key) - int64_t usage = indexer.kv_cache_usage_bytes(); - EXPECT_GE(usage, kOneGB); -} - -// ==================== max_hit_count Tests ==================== - -TEST_F(LruCacheIndexerTest, MaxHitCountWithTheoreticalCapacity) { - LruCacheIndexer indexer(true); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({1.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - // First access: miss - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(0, max_hit); - - // Re-access: all hit in max_cache - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, max_hit); -} - -// ==================== Edge Cases ==================== - -TEST_F(LruCacheIndexerTest, EmptyKeysVector) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({1.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - indexer.ProcessKeys({}, hit_count, max_hit); - - EXPECT_EQ(1u, hit_count.size()); - EXPECT_EQ(0, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, SingleKeyRepeatedAccess) { - LruCacheIndexer indexer(0); - // Use small charge per key - indexer.Init({1.0}, 1024, 1024, 1); - - std::vector hit_count; - int64_t max_hit; - - // First miss - indexer.ProcessKeys({42}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // Repeated hits - for (int i = 0; i < 10; i++) { - indexer.ProcessKeys({42}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); - } - - EXPECT_EQ(1, indexer.unique_count()); -} - -TEST_F(LruCacheIndexerTest, LinearStepOneIsDefault) { - // linear_step = 1 means every key is its own boundary - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - indexer.Init({10.0}, kOneGB, kOneGB, 1); - - std::vector hit_count; - int64_t max_hit; - - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - - // All keys should hit individually - indexer.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); -} - -TEST_F(LruCacheIndexerTest, LargeLinearStepGroupBehavior) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=4: step_hit at pos3,7,11,... - // Step_hit keys use charge=size_full_linear, others use size_full_only - indexer.Init({10.0}, kOneGB / 2, kOneGB, 4); - - std::vector hit_count; - int64_t max_hit; - - // Full group of 4 keys - indexer.ProcessKeys({10, 20, 30, 40}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // first access, all miss - - // Re-access same group - indexer.ProcessKeys({10, 20, 30, 40}, hit_count, max_hit); - EXPECT_EQ(4, hit_count[0]); // all hit with correct boundary charge -} - -TEST_F(LruCacheIndexerTest, GroupInvalidationPropagation) { - LruCacheIndexer indexer(0); - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - // linear_step=2 - indexer.Init({10.0}, kOneGB, kOneGB, 2); - - std::vector hit_count; - int64_t max_hit; - - // Populate group {1, 2} - indexer.ProcessKeys({1, 2}, hit_count, max_hit); - - // Access {1, 2, 3, 4}: group {1,2} hits, group {3,4} misses - // Prefix hit should be 2 (first group all hit) - indexer.ProcessKeys({1, 2, 3, 4}, hit_count, max_hit); - EXPECT_EQ(2, hit_count[0]); -} - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/test/online_optimizer_manager_test.cc b/kv_cache_manager/optimizer/test/online_optimizer_manager_test.cc index 8e9c6c767..e689cab20 100644 --- a/kv_cache_manager/optimizer/test/online_optimizer_manager_test.cc +++ b/kv_cache_manager/optimizer/test/online_optimizer_manager_test.cc @@ -33,7 +33,9 @@ class OnlineOptimizerManagerTest : public TESTBASE { OptimizerInstanceInfo MakeInfo(const std::string &instance_id = "i1", const std::string &group_name = "g1", int32_t block_size = 16, - int32_t linear_step = 1) { + int32_t linear_step = 0) { + // Full-only specs: linear_step > 0 would require a Mamba spec group, + // use MakeHybridInfo for linear instances. return OptimizerInstanceInfo(group_name, instance_id, block_size, @@ -46,7 +48,10 @@ class OnlineOptimizerManagerTest : public TESTBASE { OptimizerInstanceInfo MakeHybridInfo(const std::string &instance_id = "i1", const std::string &group_name = "g1", int32_t block_size = 16, - int32_t linear_step = 1) { + int32_t linear_step = -1) { + if (linear_step < 0) { + linear_step = block_size; + } return OptimizerInstanceInfo(group_name, instance_id, block_size, @@ -108,10 +113,14 @@ class OnlineOptimizerManagerTest : public TESTBASE { std::shared_ptr registry_; std::shared_ptr mgr_; - static double FullCapacityGb(int64_t capacity_blocks) { + static double CapacityGbForBytes(uint64_t capacity_bytes) { constexpr double kBytesPerGb = 1024.0 * 1024.0 * 1024.0; - constexpr double kFullBlockChargeBytes = 16384.0; - return static_cast(capacity_blocks) * kFullBlockChargeBytes / kBytesPerGb; + return static_cast(capacity_bytes) / kBytesPerGb; + } + + static double FullCapacityGb(int64_t capacity_blocks) { + constexpr uint64_t kFullBlockChargeBytes = 16384; + return CapacityGbForBytes(static_cast(capacity_blocks) * kFullBlockChargeBytes); } }; @@ -122,20 +131,26 @@ TEST_F(OnlineOptimizerManagerTest, RegisterInstanceBasic) { ErrorCode ec = RegisterInstance(info, group, result); EXPECT_EQ(EC_OK, ec); - EXPECT_EQ(16384, result.size_full_only); - EXPECT_EQ(16384, result.size_full_linear); + EXPECT_EQ(16384, result.full_charge_bytes); + EXPECT_EQ(0, result.linear_charge_bytes); EXPECT_EQ(1, result.estimated_capacity_blocks.size()); } TEST_F(OnlineOptimizerManagerTest, RegisterInstanceHybrid) { - auto info = MakeHybridInfo("i1", "g1", 16, 3); + // 48 tokens / 16 tokens-per-block = one Linear state every 3 blocks. + auto info = MakeHybridInfo("i1", "g1", 16, 48); auto group = MakeGroup("g1", {1.0}); RegisterInstanceResult result; ErrorCode ec = RegisterInstance(info, group, result); EXPECT_EQ(EC_OK, ec); - EXPECT_EQ(16384, result.size_full_only); - EXPECT_EQ(20480, result.size_full_linear); + EXPECT_EQ(16384, result.full_charge_bytes); + EXPECT_EQ(4096, result.linear_charge_bytes); + // Estimate only (hits run on the byte axis): a shared pool spends + // 3 * 16384 + 4096 = 53248 bytes per 3 blocks, so 1 GB holds about + // floor(1073741824 * 3 / 53248) blocks. + ASSERT_EQ(1, result.estimated_capacity_blocks.size()); + EXPECT_EQ(60494, result.estimated_capacity_blocks[0]); } TEST_F(OnlineOptimizerManagerTest, RegisterInstanceEmptyIdFails) { @@ -160,7 +175,7 @@ TEST_F(OnlineOptimizerManagerTest, RegisterInstanceMissingOptimizerStateInfoFail } TEST_F(OnlineOptimizerManagerTest, RegisterInstanceMissingFullGroupFails) { - OptimizerInstanceInfo info("g1", "i1", 16, MakeSpecs(), MakeGroups(), 1, OptimizerStateInfo("missing", "")); + OptimizerInstanceInfo info("g1", "i1", 16, MakeSpecs(), MakeGroups(), 16, OptimizerStateInfo("missing", "")); auto group = MakeGroup(); RegisterInstanceResult result; EXPECT_EQ(EC_BADARGS, RegisterInstance(info, group, result)); @@ -168,12 +183,27 @@ TEST_F(OnlineOptimizerManagerTest, RegisterInstanceMissingFullGroupFails) { TEST_F(OnlineOptimizerManagerTest, RegisterInstanceMissingSpecInStateGroupFails) { std::vector groups = {LocationSpecGroup("full", {"tp0", "tp_missing"})}; - OptimizerInstanceInfo info("g1", "i1", 16, MakeSpecs(), groups, 1, OptimizerStateInfo("full", "")); + OptimizerInstanceInfo info("g1", "i1", 16, MakeSpecs(), groups, 16, OptimizerStateInfo("full", "")); auto group = MakeGroup(); RegisterInstanceResult result; EXPECT_EQ(EC_BADARGS, RegisterInstance(info, group, result)); } +TEST_F(OnlineOptimizerManagerTest, RegisterInstanceLinearStepNotTokenMultipleFails) { + // linear_step counts tokens and must divide into whole blocks. + auto info = MakeHybridInfo("i1", "g1", 16, /*linear_step tokens=*/24); + auto group = MakeGroup(); + RegisterInstanceResult result; + EXPECT_EQ(EC_BADARGS, RegisterInstance(info, group, result)); + + auto ok_info = MakeHybridInfo("i1", "g1", 16, /*linear_step tokens=*/32); + EXPECT_EQ(EC_OK, RegisterInstance(ok_info, group, result)); + + // A linear instance without a Mamba spec group is rejected. + auto no_mamba_group = MakeInfo("i2", "g1", 16, /*linear_step tokens=*/32); + EXPECT_EQ(EC_BADARGS, RegisterInstance(no_mamba_group, group, result)); +} + TEST_F(OnlineOptimizerManagerTest, RegisterInstanceSharedGroupQuotaFails) { auto info = MakeInfo(); auto group = MakeGroup(); @@ -245,16 +275,35 @@ TEST_F(OnlineOptimizerManagerTest, TraceQueryMultipleCapacities) { TraceQueryResult result; mgr_->TraceQuery("i1", init_keys, result); - // This legacy (non-full-attention) path replays with the eviction-policy - // simulator: cache_hit_count uses index 0 (smallest capacity ~6 blocks), - // prefix match starts at key 0 whose stack distance (99) exceeds the small - // capacity, so prefix hit = 0. - EXPECT_EQ(0, result.hit_count_per_capacity.at(0)); - // Large capacity (index 1) should hit all 100 keys + // Full-attention LiteHit path with tail-first commit: the chain head is + // most recent, so the small capacity (~6 blocks) serves exactly its + // capacity as prefix hits. ASSERT_EQ(2, result.hit_count_per_capacity.size()); + EXPECT_EQ(reg_result.estimated_capacity_blocks[0], result.hit_count_per_capacity.at(0)); + // Large capacity (index 1) should hit all 100 keys EXPECT_EQ(100, result.hit_count_per_capacity[1]); } +TEST_F(OnlineOptimizerManagerTest, FullAttentionStoresCapacityInBytes) { + constexpr uint64_t kFullChargeBytes = 16384; + constexpr uint64_t kCapacityBytes = 2 * kFullChargeBytes - 1; + auto info = MakeInfo("i1", "g1", 4, 0); + auto group = MakeGroup("g1", {CapacityGbForBytes(kCapacityBytes)}); + RegisterInstanceResult reg_result; + ASSERT_EQ(EC_OK, RegisterInstance(info, group, reg_result)); + EXPECT_EQ((std::vector{1}), reg_result.estimated_capacity_blocks); + + ASSERT_EQ(EC_OK, mgr_->GetInstanceState("i1", [&](const InstanceState &state) { + EXPECT_EQ((std::vector{kCapacityBytes}), state.capacity_bytes); + })); + + TraceQueryResult result; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2}, 8, result)); + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2}, 8, result)); + EXPECT_EQ((std::vector{1}), result.hit_count_per_capacity); + EXPECT_EQ((std::vector{1}), result.unique_keys_per_capacity); +} + TEST_F(OnlineOptimizerManagerTest, FullAttentionUsesLiteHitTokenRates) { auto info = MakeInfo("i1", "g1", 4, 0); auto group = MakeGroup("g1", {FullCapacityGb(2), FullCapacityGb(3)}, "lru", true); @@ -279,7 +328,6 @@ TEST_F(OnlineOptimizerManagerTest, FullAttentionUsesLiteHitTokenRates) { ASSERT_EQ(EC_OK, mgr_->GetInstanceState("i1", [&](const InstanceState &state) { checked_state = true; EXPECT_NE(nullptr, state.lite_hit); - EXPECT_EQ(nullptr, state.indexer); EXPECT_EQ(2, state.total_queries); EXPECT_EQ(26, state.total_input_tokens); })); @@ -298,6 +346,101 @@ TEST_F(OnlineOptimizerManagerTest, FullAttentionUsesLiteHitTokenRates) { EXPECT_DOUBLE_EQ(12.0 / 26.0, summaries[0].max_hit_rate); } +TEST_F(OnlineOptimizerManagerTest, MambaLinearUsesSharedLiteHit) { + // block_size 16, linear_step 48 tokens -> one Linear state every 3 blocks + // plus the forced last block. Hybrid specs: full charge 16384, mamba + // charge 4096 per Linear state. + auto info = MakeHybridInfo("i1", "g1", 16, 48); + auto group = MakeGroup("g1", {1.0}, "lru", /*enable_theoretical_max_cache=*/true); + RegisterInstanceResult reg_result; + ASSERT_EQ(EC_OK, RegisterInstance(info, group, reg_result)); + + // An empty working set has no per-resident-block average yet. + { + std::vector empty_summaries; + ASSERT_EQ(EC_OK, mgr_->ListInstances("g1", empty_summaries)); + ASSERT_EQ(1u, empty_summaries.size()); + EXPECT_DOUBLE_EQ(0.0, empty_summaries[0].bytes_per_block); + } + + TraceQueryResult first; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2, 3, 4}, 70, first)); + ASSERT_EQ(1, first.hit_count_per_capacity.size()); + EXPECT_EQ(0, first.hit_count_per_capacity[0]); + EXPECT_EQ(0, first.max_hit_count); + + TraceQueryResult second; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2, 3, 4}, 70, second)); + // 1 GB covers everything: the forced tail Linear state (position 3) + // recovers all 4 complete blocks. + EXPECT_EQ(4, second.hit_count_per_capacity[0]); + EXPECT_DOUBLE_EQ(64.0 / 70.0, second.hit_rate_per_capacity[0]); + EXPECT_EQ(4, second.max_hit_count); + EXPECT_DOUBLE_EQ(64.0 / 70.0, second.max_hit_rate); + EXPECT_EQ(4, second.theoretical_unique_keys); // Full objects only + + bool checked_state = false; + ASSERT_EQ(EC_OK, mgr_->GetInstanceState("i1", [&](const InstanceState &state) { + checked_state = true; + ASSERT_NE(nullptr, state.lite_hit); + EXPECT_TRUE(state.lite_hit->uses_linear()); + EXPECT_EQ((std::vector{1ULL << 30}), state.capacity_bytes); + })); + EXPECT_TRUE(checked_state); + + std::vector summaries; + ASSERT_EQ(EC_OK, mgr_->ListInstances("g1", summaries)); + ASSERT_EQ(1, summaries.size()); + EXPECT_EQ(2, summaries[0].total_queries); + EXPECT_EQ(140, summaries[0].total_input_tokens); + EXPECT_EQ(4, summaries[0].unique_keys); + // Working set: 4 Full * 16384 + 2 Linear states (positions 2 and 3) * 4096. + EXPECT_EQ(4 * 16384 + 2 * 4096, summaries[0].kv_cache_usage_bytes); + EXPECT_DOUBLE_EQ(static_cast(4 * 16384 + 2 * 4096) / 4, summaries[0].bytes_per_block); + ASSERT_EQ(1, summaries[0].per_capacity_hit_rates.size()); + EXPECT_EQ(4, summaries[0].per_capacity_hit_rates[0].total_hits); + EXPECT_DOUBLE_EQ(64.0 / 140.0, summaries[0].per_capacity_hit_rates[0].hit_rate); + + ASSERT_EQ(EC_OK, mgr_->ResetStats("i1")); + TraceQueryResult after_reset; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2, 3, 4}, 70, after_reset)); + EXPECT_EQ(0, after_reset.hit_count_per_capacity[0]); +} + +TEST_F(OnlineOptimizerManagerTest, MambaCountsHistoricalForcedTailLinearState) { + // block_size 16, linear_step 48 -> periodic Linear states every 3 blocks. + auto info = MakeHybridInfo("i1", "g1", 16, 48); + auto group = MakeGroup("g1", {1.0}, "lru", /*enable_theoretical_max_cache=*/true); + RegisterInstanceResult reg_result; + ASSERT_EQ(EC_OK, RegisterInstance(info, group, reg_result)); + + TraceQueryResult first; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2}, 32, first)); + EXPECT_EQ(0, first.hit_count_per_capacity[0]); + EXPECT_EQ(0, first.max_hit_count); + + TraceQueryResult second; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1, 2, 3, 4}, 64, second)); + // key 2 has a Linear state solely as the first request's forced tail. It is + // not scheduled in this request, but remains a valid + // restore point and must contribute two hit blocks to both statistics. + EXPECT_EQ(2, second.hit_count_per_capacity[0]); + EXPECT_DOUBLE_EQ(0.5, second.hit_rate_per_capacity[0]); + EXPECT_EQ(2, second.max_hit_count); + EXPECT_DOUBLE_EQ(0.5, second.max_hit_rate); + + std::vector summaries; + ASSERT_EQ(EC_OK, mgr_->ListInstances("g1", summaries)); + ASSERT_EQ(1u, summaries.size()); + EXPECT_EQ(2, summaries[0].per_capacity_hit_rates[0].total_hits); + EXPECT_DOUBLE_EQ(32.0 / 96.0, summaries[0].per_capacity_hit_rates[0].hit_rate); + EXPECT_DOUBLE_EQ(32.0 / 96.0, summaries[0].max_hit_rate); + // 4 Full blocks plus current Linear states 3/4 and historical forced-tail + // Linear state 2: the real working-set average is (4F + 3M) / 4. + EXPECT_EQ(4 * 16384 + 3 * 4096, summaries[0].kv_cache_usage_bytes); + EXPECT_DOUBLE_EQ(static_cast(4 * 16384 + 3 * 4096) / 4, summaries[0].bytes_per_block); +} + TEST_F(OnlineOptimizerManagerTest, FullAttentionRequiresConsistentInputTokenLength) { auto info = MakeInfo("i1", "g1", 4, 0); auto group = MakeGroup("g1", {FullCapacityGb(2)}); @@ -344,6 +487,32 @@ TEST_F(OnlineOptimizerManagerTest, FullAttentionLayersGroupTtlOntoLiteHit) { EXPECT_EQ(EC_BADARGS, RegisterInstance(bad_info, bad_group, reg_result)); } +TEST_F(OnlineOptimizerManagerTest, LinearInstanceLayersGroupTtlOntoSharedCore) { + auto info = MakeHybridInfo("i1", "g1", 16, 48); + auto ttl_group = MakeGroup("g1", {1.0}, "lru", /*enable_theoretical_max_cache=*/false, /*ttl=*/300); + RegisterInstanceResult reg_result; + ASSERT_EQ(EC_OK, RegisterInstance(info, ttl_group, reg_result)); + + TraceQueryResult first; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1}, 16, first)); + EXPECT_EQ(0, first.hit_count_per_capacity.at(0)); + TraceQueryResult second; + ASSERT_EQ(EC_OK, mgr_->TraceQuery("i1", {1}, 16, second)); + EXPECT_EQ(1, second.hit_count_per_capacity.at(0)); + + // Drive the shared watermark past the deadline without waiting in the UT. + ASSERT_EQ(EC_OK, + mgr_->GetInstanceState("i1", [](const InstanceState &state) { state.lite_hit->AdvanceTime(LLONG_MAX); })); + std::vector summaries; + ASSERT_EQ(EC_OK, mgr_->ListInstances("g1", summaries)); + ASSERT_EQ(1u, summaries.size()); + EXPECT_EQ(0, summaries[0].unique_keys); + EXPECT_EQ(0, summaries[0].kv_cache_usage_bytes); + EXPECT_DOUBLE_EQ(0.0, summaries[0].bytes_per_block); + EXPECT_EQ(1, summaries[0].ttl_eviction_count); // Full objects only + EXPECT_EQ(summaries[0].ttl_eviction_count, summaries[0].eviction_count); +} + TEST_F(OnlineOptimizerManagerTest, ResetStatsResetsFullAttentionLiteHit) { auto info = MakeInfo("i1", "g1", 4, 0); auto group = MakeGroup("g1", {FullCapacityGb(2)}); diff --git a/kv_cache_manager/optimizer/test/optimizer_metrics_reporter_test.cc b/kv_cache_manager/optimizer/test/optimizer_metrics_reporter_test.cc index ddfabe3bd..7e734dc19 100644 --- a/kv_cache_manager/optimizer/test/optimizer_metrics_reporter_test.cc +++ b/kv_cache_manager/optimizer/test/optimizer_metrics_reporter_test.cc @@ -331,32 +331,6 @@ TEST_F(OptimizerMetricsReporterTest, ReportIntervalCapacityEfficiencySkippedWhen EXPECT_DOUBLE_EQ(0.0, hit_rate.Get()); } -TEST_F(OptimizerMetricsReporterTest, ReportIntervalHitAgeBucketRatio) { - // TTL > 0 triggers TtlCacheIndexerWrapper, enabling age-bucket tracking - ASSERT_EQ(EC_OK, - RegisterTestInstance("inst1", - {1.0}, - /*ttl_seconds=*/3600, - /*enable_theoretical_max_cache=*/false, - /*linear_step=*/1)); - - TraceQueryResult result; - manager_->TraceQuery("inst1", {1, 2, 3}, result); - manager_->TraceQuery("inst1", {1, 2, 3}, result); // all 3 keys hit - - reporter_->ReportInterval(); - - // With near-zero age, all hits should fall in the first bucket (threshold=5s) - MetricsTags bucket_tags = {{"instance_id", "inst1"}, {"age_bucket", "5s"}}; - Gauge bucket_ratio = registry_->GetGauge("trace_query_hit_age_bucket_ratio", bucket_tags); - EXPECT_GT(bucket_ratio.Get(), 0.0); - - // The "inf" bucket should have ratio = 0 (no hits that old) - MetricsTags inf_tags = {{"instance_id", "inst1"}, {"age_bucket", "inf"}}; - Gauge inf_ratio = registry_->GetGauge("trace_query_hit_age_bucket_ratio", inf_tags); - EXPECT_DOUBLE_EQ(0.0, inf_ratio.Get()); -} - TEST_F(OptimizerMetricsReporterTest, RemoveInstanceMetricsCleansUp) { ASSERT_EQ(EC_OK, RegisterTestInstance("inst1")); diff --git a/kv_cache_manager/optimizer/test/optimizer_service_impl_test.cc b/kv_cache_manager/optimizer/test/optimizer_service_impl_test.cc index a195ecb23..ed7a04884 100644 --- a/kv_cache_manager/optimizer/test/optimizer_service_impl_test.cc +++ b/kv_cache_manager/optimizer/test/optimizer_service_impl_test.cc @@ -33,8 +33,14 @@ class OptimizerServiceImplTest : public TESTBASE { proto::optimizer::OptimizerRegisterInstanceRequest MakeRegisterRequest(const std::string &group, const std::string &instance_id, int32_t block_size, - int32_t linear_step = 1, + int32_t linear_step = -1, int64_t extra_spec_size = 0) { + // linear_step counts tokens; -1 = auto: linear instances (with an + // extra Linear spec) default to one Linear state per block, plain + // instances stay full-attention. + if (linear_step < 0) { + linear_step = extra_spec_size > 0 ? block_size : 0; + } proto::optimizer::OptimizerRegisterInstanceRequest req; req.set_trace_id("test-trace"); req.set_instance_group(group); @@ -523,7 +529,7 @@ TEST_F(OptimizerServiceImplTest, ResetStatsNonExistent) { TEST_F(OptimizerServiceImplTest, RegisterWithLinearStep) { CreateTestGroup("grp_ls", 1.0); - auto req = MakeRegisterRequest("grp_ls", "inst1", 1024, 4, 256); + auto req = MakeRegisterRequest("grp_ls", "inst1", 1024, /*linear_step tokens=*/4096, 256); proto::optimizer::OptimizerRegisterInstanceResponse resp; RequestContext ctx("trace1", nullptr); @@ -536,7 +542,7 @@ TEST_F(OptimizerServiceImplTest, RegisterWithLinearStep) { TEST_F(OptimizerServiceImplTest, GetInstanceSuccess) { CreateTestGroup("grp1"); - auto reg_req = MakeRegisterRequest("grp1", "inst1", 128, 2, 64); + auto reg_req = MakeRegisterRequest("grp1", "inst1", 128, /*linear_step tokens=*/256, 64); proto::optimizer::OptimizerRegisterInstanceResponse reg_resp; RequestContext ctx1("t1", nullptr); @@ -554,7 +560,7 @@ TEST_F(OptimizerServiceImplTest, GetInstanceSuccess) { EXPECT_EQ("grp1", get_resp.instance_group()); EXPECT_EQ("inst1", get_resp.instance_id()); EXPECT_EQ(128, get_resp.block_size()); - EXPECT_EQ(2, get_resp.linear_step()); + EXPECT_EQ(256, get_resp.linear_step()); ASSERT_EQ(2, get_resp.location_spec_infos_size()); EXPECT_EQ("tp0", get_resp.location_spec_infos(0).name()); diff --git a/kv_cache_manager/optimizer/test/ttl_cache_indexer_wrapper_test.cc b/kv_cache_manager/optimizer/test/ttl_cache_indexer_wrapper_test.cc deleted file mode 100644 index 0d400262e..000000000 --- a/kv_cache_manager/optimizer/test/ttl_cache_indexer_wrapper_test.cc +++ /dev/null @@ -1,388 +0,0 @@ -#include - -#include "kv_cache_manager/common/unittest.h" -#include "kv_cache_manager/optimizer/index/online/cache_indexer_factory.h" -#include "kv_cache_manager/optimizer/index/online/ttl_cache_indexer_wrapper.h" - -namespace kv_cache_manager { - -class TtlCacheIndexerWrapperTest : public TESTBASE {}; - -static std::unique_ptr MakeInnerIndexer(const std::string &policy_type = "lru", - bool enable_theoretical_max_cache = false, - double capacity_gb = 10.0) { - constexpr int64_t kOneGB = 1024LL * 1024 * 1024; - auto indexer = CacheIndexerFactory::CreateCacheIndexer( - policy_type, enable_theoretical_max_cache, {capacity_gb}, kOneGB, kOneGB, 1); - return indexer; -} - -TEST_F(TtlCacheIndexerWrapperTest, BasicExpiration) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, wrapper.unique_count()); - EXPECT_EQ(0, wrapper.ttl_eviction_count()); - - // Re-access within TTL: should hit - now = 1005; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); // all hit (prefix match) - - // Advance past TTL (last access was at 1005, TTL=10, expires at 1015) - now = 1016; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // all expired -> all miss - EXPECT_EQ(3, wrapper.ttl_eviction_count()); -} - -TEST_F(TtlCacheIndexerWrapperTest, SlidingTtlRefreshesExpiry) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1}, hit_count, max_hit); - - // Access at t=1008, refreshes TTL to expire at 1018 - now = 1008; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); // hit - - // At t=1012, would have expired without refresh (1000+10=1010), but was refreshed - now = 1012; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); // still hit - EXPECT_EQ(0, wrapper.ttl_eviction_count()); - - // At t=1023, past refreshed expiry (1012+10=1022) - now = 1023; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // expired - EXPECT_EQ(1, wrapper.ttl_eviction_count()); -} - -TEST_F(TtlCacheIndexerWrapperTest, TtlAndCapacityEvictionInteraction) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer("lru", false, 3.0), 100, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2, 3, 4, 5}, hit_count, max_hit); - wrapper.PostQueryMaintenance(); - // Capacity eviction should have removed 2 keys - EXPECT_EQ(3, wrapper.unique_count()); - EXPECT_EQ(2, wrapper.eviction_count()); // total = capacity(2) + ttl(0) - EXPECT_EQ(0, wrapper.ttl_eviction_count()); - - // Now expire remaining keys via TTL - now = 1101; - wrapper.ProcessKeys({10}, hit_count, max_hit); - // 3 keys expired by TTL + 1 new key - EXPECT_EQ(3, wrapper.ttl_eviction_count()); - EXPECT_EQ(1, wrapper.unique_count()); -} - -TEST_F(TtlCacheIndexerWrapperTest, WorksWithFactoryLruIndexer) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2}, hit_count, max_hit); - EXPECT_EQ(2, wrapper.unique_count()); - - now = 1005; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); - - now = 1011; - wrapper.ProcessKeys({1, 2}, hit_count, max_hit); - // key 2 expired (last access 1000+10=1010), key 1 not (last access 1005+10=1015) - EXPECT_EQ(1, hit_count[0]); // key 1 hit, then key 2 miss -> prefix hit = 1 - EXPECT_EQ(1, wrapper.ttl_eviction_count()); -} - -TEST_F(TtlCacheIndexerWrapperTest, WorksWithLruIndexer) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, wrapper.unique_count()); - - now = 1005; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, hit_count[0]); // all hit - - now = 1016; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // all expired - EXPECT_EQ(3, wrapper.ttl_eviction_count()); -} - -TEST_F(TtlCacheIndexerWrapperTest, MetricsAccuracy) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, wrapper.unique_count()); - EXPECT_EQ(0, wrapper.eviction_count()); - EXPECT_EQ(0, wrapper.ttl_eviction_count()); - EXPECT_GT(wrapper.memory_usage_bytes(), 0); - EXPECT_GT(wrapper.kv_cache_usage_bytes(), 0); - - now = 1011; - wrapper.ProcessKeys({4}, hit_count, max_hit); - EXPECT_EQ(1, wrapper.unique_count()); - EXPECT_EQ(3, wrapper.ttl_eviction_count()); - EXPECT_EQ(3, wrapper.eviction_count()); // ttl evictions counted in total -} - -TEST_F(TtlCacheIndexerWrapperTest, RemoveKeyCleansTtlState) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - EXPECT_EQ(3, wrapper.unique_count()); - - wrapper.RemoveKey(2); - EXPECT_EQ(2, wrapper.unique_count()); - - // Key 2 should be a miss now - now = 1002; - wrapper.ProcessKeys({2}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); // miss (was removed, re-inserted as new) - EXPECT_EQ(3, wrapper.unique_count()); - - // Advance past TTL: key 2 (re-inserted at 1002) should survive, keys 1,3 (at 1000) should expire - now = 1011; - wrapper.ProcessKeys({2}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); // key 2 hit (inserted at 1002, expires at 1012) - EXPECT_EQ(2, wrapper.ttl_eviction_count()); // keys 1 and 3 expired -} - -TEST_F(TtlCacheIndexerWrapperTest, PartialExpiration) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1}, hit_count, max_hit); - - now = 1003; - wrapper.ProcessKeys({2}, hit_count, max_hit); - - now = 1006; - wrapper.ProcessKeys({3}, hit_count, max_hit); - - // At t=1011: key 1 (access 1000) expires, keys 2,3 survive - now = 1011; - wrapper.ProcessKeys({2, 3}, hit_count, max_hit); - EXPECT_EQ(2, hit_count[0]); - EXPECT_EQ(1, wrapper.ttl_eviction_count()); // only key 1 - - // At t=1014: key 2 (access refreshed to 1011) still alive - now = 1014; - wrapper.ProcessKeys({2}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); - - // At t=1017: key 3 (access refreshed to 1011) expires - now = 1022; - wrapper.ProcessKeys({2, 3}, hit_count, max_hit); - // key 2 refreshed at 1014, expires at 1024 -> still alive - // key 3 refreshed at 1011, expires at 1021 -> expired - EXPECT_EQ(1, hit_count[0]); // key 2 hits, key 3 misses -> prefix = 1 -} - -TEST_F(TtlCacheIndexerWrapperTest, ExactTtlBoundary) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10, clock); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1}, hit_count, max_hit); - - // One second before expiry: key should still be alive - now = 1009; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(1, hit_count[0]); - EXPECT_EQ(0, wrapper.ttl_eviction_count()); - - // Sliding TTL: access at 1009 refreshes expiry to 1019 - // At exactly last_access + ttl_seconds (1009+10=1019): key should be expired - now = 1019; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - EXPECT_EQ(1, wrapper.ttl_eviction_count()); -} - -TEST_F(TtlCacheIndexerWrapperTest, HitAgeBucketBasic) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 3600, clock); - // Use small thresholds for easy testing - wrapper.SetHitAgeBucketThresholds({10, 30, 60}); - // Buckets: [0,10) [10,30) [30,60) [60,+inf) - - std::vector hit_count; - int64_t max_hit; - - // Insert keys 1, 2, 3 - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - // Re-access key 1 at age=5 (< 10 → bucket 0) - now = 1005; - wrapper.ProcessKeys({1}, hit_count, max_hit); - - // Re-access key 2 at age=15 (>= 10, < 30 → bucket 1) - now = 1015; - wrapper.ProcessKeys({2}, hit_count, max_hit); - - // Re-access key 3 at age=45 (>= 30, < 60 → bucket 2) - now = 1045; - wrapper.ProcessKeys({3}, hit_count, max_hit); - - auto buckets = wrapper.GetHitAgeBuckets(); - ASSERT_EQ(4u, buckets.size()); - - // bucket [0,10): threshold=10, count=1 (key 1 at age 5) - EXPECT_EQ(10, buckets[0].threshold_seconds); - EXPECT_EQ(1, buckets[0].hit_count); - - // bucket [10,30): threshold=30, count=1 (key 2 at age 15) - EXPECT_EQ(30, buckets[1].threshold_seconds); - EXPECT_EQ(1, buckets[1].hit_count); - - // bucket [30,60): threshold=60, count=1 (key 3 at age 45) - EXPECT_EQ(60, buckets[2].threshold_seconds); - EXPECT_EQ(1, buckets[2].hit_count); - - // bucket [60,+inf): threshold=0, count=0 - EXPECT_EQ(0, buckets[3].threshold_seconds); - EXPECT_EQ(0, buckets[3].hit_count); -} - -TEST_F(TtlCacheIndexerWrapperTest, CapacityEvictedKeyDoesNotIncrementHitAgeBucket) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer("lru", true, 2.0), 100, clock); - wrapper.SetHitAgeBucketThresholds({10}); - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1, 2}, hit_count, max_hit); - wrapper.PostQueryMaintenance(); - - now = 1005; - wrapper.ProcessKeys({3}, hit_count, max_hit); - wrapper.PostQueryMaintenance(); - - now = 1010; - wrapper.ProcessKeys({1}, hit_count, max_hit); - EXPECT_EQ(0, hit_count[0]); - EXPECT_EQ(1, max_hit); - - auto buckets = wrapper.GetHitAgeBuckets(); - ASSERT_EQ(2u, buckets.size()); - EXPECT_EQ(0, buckets[0].hit_count); - EXPECT_EQ(0, buckets[1].hit_count); -} - -TEST_F(TtlCacheIndexerWrapperTest, HitAgeBucketInfinityBucket) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 10000, clock); - wrapper.SetHitAgeBucketThresholds({10, 100}); - // Buckets: [0,10) [10,100) [100,+inf) - - std::vector hit_count; - int64_t max_hit; - wrapper.ProcessKeys({1}, hit_count, max_hit); - - // Re-access at age=200 (>= 100 → bucket 2, the +inf bucket) - now = 1200; - wrapper.ProcessKeys({1}, hit_count, max_hit); - - auto buckets = wrapper.GetHitAgeBuckets(); - ASSERT_EQ(3u, buckets.size()); - EXPECT_EQ(0, buckets[0].hit_count); // [0,10) - EXPECT_EQ(0, buckets[1].hit_count); // [10,100) - EXPECT_EQ(1, buckets[2].hit_count); // [100,+inf) - EXPECT_EQ(0, buckets[2].threshold_seconds); -} - -TEST_F(TtlCacheIndexerWrapperTest, HitAgeBucketDefaultThresholds) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 100000, clock); - // Don't call SetHitAgeBucketThresholds — use defaults {5,30,60,120,300,600,1800,3600,7200} - - auto buckets = wrapper.GetHitAgeBuckets(); - // 9 thresholds + 1 infinity bucket = 10 - ASSERT_EQ(10u, buckets.size()); - EXPECT_EQ(5, buckets[0].threshold_seconds); - EXPECT_EQ(30, buckets[1].threshold_seconds); - EXPECT_EQ(60, buckets[2].threshold_seconds); - EXPECT_EQ(120, buckets[3].threshold_seconds); - EXPECT_EQ(300, buckets[4].threshold_seconds); - EXPECT_EQ(600, buckets[5].threshold_seconds); - EXPECT_EQ(1800, buckets[6].threshold_seconds); - EXPECT_EQ(3600, buckets[7].threshold_seconds); - EXPECT_EQ(7200, buckets[8].threshold_seconds); - EXPECT_EQ(0, buckets[9].threshold_seconds); - - // All counts should be 0 initially - for (const auto &b : buckets) { - EXPECT_EQ(0, b.hit_count); - } -} - -TEST_F(TtlCacheIndexerWrapperTest, HitAgeBucketMultipleHitsSameBucket) { - int64_t now = 1000; - auto clock = [&now]() { return now; }; - TtlCacheIndexerWrapper wrapper(MakeInnerIndexer(), 3600, clock); - wrapper.SetHitAgeBucketThresholds({10, 30}); - - std::vector hit_count; - int64_t max_hit; - - // Insert 3 keys - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - // Re-access all 3 at age=5 → all fall into bucket [0,10) - now = 1005; - wrapper.ProcessKeys({1, 2, 3}, hit_count, max_hit); - - auto buckets = wrapper.GetHitAgeBuckets(); - ASSERT_EQ(3u, buckets.size()); - EXPECT_EQ(3, buckets[0].hit_count); // [0,10): 3 hits - EXPECT_EQ(0, buckets[1].hit_count); // [10,30) - EXPECT_EQ(0, buckets[2].hit_count); // [30,+inf) -} - -TEST_F(TtlCacheIndexerWrapperTest, HitAgeBucketNoTrackingWithoutTtlWrapper) { - // Base CacheIndexer should return empty buckets - auto indexer = MakeInnerIndexer(); - auto buckets = indexer->GetHitAgeBuckets(); - EXPECT_TRUE(buckets.empty()); -} - -} // namespace kv_cache_manager diff --git a/kv_cache_manager/optimizer/test/weighted_lru_pool_test.cc b/kv_cache_manager/optimizer/test/weighted_lru_pool_test.cc new file mode 100644 index 000000000..2acc437d5 --- /dev/null +++ b/kv_cache_manager/optimizer/test/weighted_lru_pool_test.cc @@ -0,0 +1,192 @@ +#include +#include +#include +#include +#include + +#include "kv_cache_manager/common/unittest.h" +#include "kv_cache_manager/optimizer/liteHit/weighted_lru_pool.h" + +namespace kv_cache_manager { + +namespace { + +// Naive weighted LRU oracle: a recency list of typed objects with fixed +// per-type charges. RequiredBytes = own charge + bytes of all newer objects. +class NaiveWeightedLru { +public: + NaiveWeightedLru(uint64_t full_charge, uint64_t linear_charge) : charges_{full_charge, linear_charge} {} + + void Touch(const CacheObjectKey &key) { + Remove(key); + entries_.push_front(key); + } + + bool RequiredBytes(const CacheObjectKey &key, uint64_t &required) const { + uint64_t newer = 0; + for (const auto &entry : entries_) { + if (Same(entry, key)) { + required = newer + charges_[static_cast(key.type)]; + return true; + } + newer += charges_[static_cast(entry.type)]; + } + return false; + } + + uint64_t TotalBytes() const { + uint64_t total = 0; + for (const auto &entry : entries_) { + total += charges_[static_cast(entry.type)]; + } + return total; + } + + uint64_t FullObjectsWithinBytes(uint64_t budget) const { + uint64_t cumulative = 0; + uint64_t count = 0; + for (const auto &entry : entries_) { + cumulative += charges_[static_cast(entry.type)]; + if (cumulative > budget) { + break; + } + if (entry.type == CacheObjectType::kFull) { + ++count; + } + } + return count; + } + +private: + static bool Same(const CacheObjectKey &a, const CacheObjectKey &b) { + return a.type == b.type && a.prefix_block_key == b.prefix_block_key; + } + void Remove(const CacheObjectKey &key) { + entries_.remove_if([&](const CacheObjectKey &entry) { return Same(entry, key); }); + } + + uint64_t charges_[2]; + std::list entries_; +}; + +} // namespace + +class WeightedLruPoolTest : public TESTBASE {}; + +TEST_F(WeightedLruPoolTest, TypedKeysDoNotCollide) { + WeightedLruPool pool(/*full=*/100, /*linear=*/30); + const CacheObjectKey full_key{CacheObjectType::kFull, 42}; + const CacheObjectKey linear_key{CacheObjectType::kLinear, 42}; + + pool.Touch(full_key); + EXPECT_TRUE(pool.IsResident(full_key)); + EXPECT_FALSE(pool.IsResident(linear_key)); + + pool.Touch(linear_key); + EXPECT_EQ(1u, pool.resident_full_count()); + EXPECT_EQ(1u, pool.resident_linear_count()); + EXPECT_EQ(130u, pool.resident_bytes()); + + // Linear was touched later, so Full requires both charges, Linear only its own. + uint64_t required = 0; + ASSERT_TRUE(pool.RequiredBytes(linear_key, required)); + EXPECT_EQ(30u, required); + ASSERT_TRUE(pool.RequiredBytes(full_key, required)); + EXPECT_EQ(130u, required); +} + +TEST_F(WeightedLruPoolTest, TouchMovesToMostRecent) { + WeightedLruPool pool(/*full=*/10, /*linear=*/4); + const CacheObjectKey a{CacheObjectType::kFull, 1}; + const CacheObjectKey b{CacheObjectType::kFull, 2}; + const CacheObjectKey m{CacheObjectType::kLinear, 1}; + + pool.Touch(a); + pool.Touch(b); + pool.Touch(m); + // Order (old->new): a, b, m + uint64_t required = 0; + ASSERT_TRUE(pool.RequiredBytes(a, required)); + EXPECT_EQ(24u, required); // 10 + (10 + 4) + + pool.Touch(a); + // Order: b, m, a + ASSERT_TRUE(pool.RequiredBytes(a, required)); + EXPECT_EQ(10u, required); + ASSERT_TRUE(pool.RequiredBytes(b, required)); + EXPECT_EQ(24u, required); + EXPECT_EQ(24u, pool.resident_bytes()); +} + +TEST_F(WeightedLruPoolTest, NonResidentReturnsFalse) { + WeightedLruPool pool(8, 2); + uint64_t required = 0; + EXPECT_FALSE(pool.RequiredBytes({CacheObjectType::kFull, 7}, required)); + pool.Touch({CacheObjectType::kFull, 7}); + EXPECT_TRUE(pool.RequiredBytes({CacheObjectType::kFull, 7}, required)); + pool.Reset(); + EXPECT_FALSE(pool.RequiredBytes({CacheObjectType::kFull, 7}, required)); + EXPECT_EQ(0u, pool.resident_bytes()); +} + +TEST_F(WeightedLruPoolTest, RandomizedOracleComparison) { + std::mt19937 rng(20260729); + WeightedLruPool pool(/*full=*/48, /*linear=*/16); + NaiveWeightedLru oracle(48, 16); + + std::uniform_int_distribution key_dist(0, 63); + std::uniform_int_distribution type_dist(0, 4); + for (int step = 0; step < 20000; ++step) { + const CacheObjectKey key{type_dist(rng) == 0 ? CacheObjectType::kLinear : CacheObjectType::kFull, + key_dist(rng)}; + pool.Touch(key); + oracle.Touch(key); + + if (step % 7 == 0) { + const CacheObjectKey probe{type_dist(rng) == 0 ? CacheObjectType::kLinear : CacheObjectType::kFull, + key_dist(rng)}; + uint64_t expected = 0; + uint64_t actual = 0; + const bool expected_resident = oracle.RequiredBytes(probe, expected); + const bool actual_resident = pool.RequiredBytes(probe, actual); + ASSERT_EQ(expected_resident, actual_resident) << "step " << step; + if (expected_resident) { + ASSERT_EQ(expected, actual) << "step " << step; + } + } + if (step % 11 == 0) { + for (uint64_t budget : {0ull, 16ull, 48ull, 100ull, 500ull, 5000ull}) { + ASSERT_EQ(oracle.FullObjectsWithinBytes(budget), pool.FullObjectsWithinBytes(budget)) + << "step " << step << " budget " << budget; + } + } + ASSERT_EQ(oracle.TotalBytes(), pool.resident_bytes()) << "step " << step; + } +} + +TEST_F(WeightedLruPoolTest, CompactionPreservesOrderAndBytes) { + WeightedLruPool pool(/*full=*/3, /*linear=*/5); + NaiveWeightedLru oracle(3, 5); + + // Force many dead positions: repeatedly touch a small key set far beyond + // the compaction slack. + for (int round = 0; round < 3000; ++round) { + for (int64_t key = 0; key < 8; ++key) { + const CacheObjectKey object{key % 3 == 0 ? CacheObjectType::kLinear : CacheObjectType::kFull, key}; + pool.Touch(object); + oracle.Touch(object); + } + } + pool.MaybeCompactPositions(); + for (int64_t key = 0; key < 8; ++key) { + const CacheObjectKey object{key % 3 == 0 ? CacheObjectType::kLinear : CacheObjectType::kFull, key}; + uint64_t expected = 0; + uint64_t actual = 0; + ASSERT_TRUE(oracle.RequiredBytes(object, expected)); + ASSERT_TRUE(pool.RequiredBytes(object, actual)); + EXPECT_EQ(expected, actual) << "key " << key; + } + EXPECT_EQ(oracle.TotalBytes(), pool.resident_bytes()); +} + +} // namespace kv_cache_manager diff --git a/kv_cache_manager/protocol/protobuf/optimizer_service.proto b/kv_cache_manager/protocol/protobuf/optimizer_service.proto index d7688a3ea..7a66f736f 100644 --- a/kv_cache_manager/protocol/protobuf/optimizer_service.proto +++ b/kv_cache_manager/protocol/protobuf/optimizer_service.proto @@ -111,12 +111,12 @@ message OptimizerRegisterInstanceRequest { repeated LocationSpecInfo location_spec_infos = 5; // KVCM 风格 spec 定义 repeated LocationSpecGroup location_spec_groups = 6; // KVCM 风格 spec group 定义 OptimizerStateInfo optimizer_state_info = 7; // 显式声明 optimizer full/linear 状态如何映射到 KVCM group - int32 linear_step = 8; // 0 表示 full-only；>0 表示每 linear_step 个 block 产生一次 full+linear checkpoint + int32 linear_step = 8; // 0 表示 full-only；>0 表示每 linear_step 个 **token** 产生一次 Linear 状态（须为 block_size 的整数倍，Linear 状态只落在完整 block 边界） } message OptimizerRegisterInstanceResponse { CommonResponseHeader header = 1; - repeated int64 estimated_capacity_blocks = 2; // 各档容量按 average block size 折算的 block 数 + repeated int64 estimated_capacity_blocks = 2; // 兼容展示字段，不参与命中计算。full-attention：由 byte 容量精确派生；linear：仅估算（Full block 与 Linear 状态字节不等） int64 size_full_only = 3; // 仅 full spec 的单 block 字节数 int64 size_full_linear = 4; // full+linear 混合的单 block 字节数 }