[optimizer] unify linear attention and TTL on one LiteHit core and retire the legacy indexer - #287
[optimizer] unify linear attention and TTL on one LiteHit core and retire the legacy indexer#287Tyndalllll wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
Review Summary
PR #287 — [optimizer] add the LiteHitMamba core for linear attention and retire the legacy indexer
This is a substantial, well-structured PR. The core design — one shared WeightedLruPool with byte-weighted Mattson stack inclusion, a capacity-independent MambaRequestFact step curve, and a clean offline/online agreement test — is sound. The removal of the legacy LruCacheIndexer/TtlCacheIndexerWrapper is clean and the proto field reservations are correctly handled. The randomized differential oracle and round-trip CSV tests provide good algorithmic confidence.
Findings
Six inline comments are attached; none are correctness bugs, but a few need attention before merge:
step_blocks == 0silent clamp (lite_hit_mamba.cc:10) — Dead code that masks upstream misconfiguration; should be an assertion.- Pareto-sweep invariant comment (
lite_hit_mamba.cc:57) — Equal-threshold resolution is correct but not documented; theMambaRequestFactinvariant comment should state how ties are handled. IsCheckpointreceivesnnotcovered_blocks(lite_hit_mamba.cc:46) — Correct behavior, but the variable name collision is confusing; a one-line comment would prevent future bugs.- Fenwick 1-based position convention (
weighted_lru_pool.cc:39) — Non-obvious; a comment stating that positions are 1-indexed andsize()== last position would guard against off-by-one regressions. - Compaction threshold formula (
weighted_lru_pool.cc:80) — The ceiling-division guard is correct but unintuitive; a comment explaining the "50% dead density" trigger would help. - Unnecessary temporary
std::stringallocation (facts_csv.cc:193) — Minor allocator churn on the CSV parsing path; usestrlenor a companionconstexprlength constant.
Noted gaps (not blocking on their own, but worth tracking)
- The PR description explicitly notes
linear_step > 0+ group TTL is rejected at registration. The rejection is tested (RejectsLinearInstanceInTtlGroup). The follow-up (grafting the LiteHit position-epoch watermark ontoWeightedLruPool) should be tracked as a known limitation in comments or a follow-up issue. - Internal-source UT is listed as unchecked in the test plan. If any internal-source tests cover the deleted
CacheIndexerpath, that coverage gap should be resolved before this lands. - The
internal_sourceUT checkbox is unchecked — confirm this is intentional (i.e. no internal tests exercise the deleted legacy indexer path) before merging.
🤖 Generated by Qoder
| // the Full prefix [0..p] both stay resident. Both live in the same pool, so | ||
| // the demand is the larger of the two. Ascending positions with a Pareto | ||
| // sweep yield the strictly monotone envelope. | ||
| for (std::size_t p = 0; p < covered_blocks; ++p) { |
There was a problem hiding this comment.
Phase 1 loop boundary: covered_blocks should use n, not the live value
The Phase-1 checkpoint loop iterates p < covered_blocks, where covered_blocks is set to i + 1 after each successfully resolved Full block. If the first Full block is non-resident, covered_blocks stays 0 and the loop body never executes — correct. But consider a request where blocks 0…k−1 are resident and block k is not. covered_blocks == k. Checkpoints at positions 0…k−1 are evaluated correctly.
However, the Mamba spec says a checkpoint fires at position n-1 (last block) regardless of step_blocks. If block n-1 is absent but n-1 < k, i.e. the final block is within the covered prefix but the total request is longer, the forced-last check position == total_blocks - 1 still uses n (the full length), so IsCheckpoint(p, n) is called with the real n. That part is fine.
The edge case I want to flag: when the request has exactly n blocks and all are resident, covered_blocks == n, so the loop processes p ∈ [0, n). The forced-last checkpoint is position n-1, which is within the loop — OK. But the IsCheckpoint call is IsCheckpoint(p, n) where n is the original total block count. This is consistent.
This looks correct; I'm flagging it for an explicit comment in the code, because the covered_blocks variable doubles as both a Fenwick walk bound and the exclusive upper limit for the Phase-1 loop, and the forced-last-checkpoint semantics are evaluated against n (request length) not covered_blocks. Adding a one-line comment clarifying that total_blocks in IsCheckpoint is always n (the full request length, not the covered prefix) would prevent future confusion.
🤖 Generated by Qoder
| const uint64_t threshold = std::max(prefix_full_required[p], mamba_required); | ||
| while (!fact.points.empty() && fact.points.back().min_total_capacity_bytes >= threshold) { | ||
| fact.points.pop_back(); | ||
| } |
There was a problem hiding this comment.
Pareto sweep direction: later checkpoints could dominate earlier ones in the wrong direction
The envelope loop removes points from the back when the new threshold is <= threshold:
while (!fact.points.empty() && fact.points.back().min_total_capacity_bytes >= threshold) {
fact.points.pop_back();
}
fact.points.push_back(MambaCurvePoint{threshold, static_cast<uint64_t>(p + 1)});Positions are visited in ascending order, so p + 1 strictly increases. The pop condition removes a previous point when the new threshold is no larger than it. This is correct for a strictly-monotone-increasing envelope: a later point (higher hit_blocks) at an equal or lower threshold dominates the earlier one, so removing the earlier one is right.
However the PR description says "Points are strictly increasing in both fields (monotone envelope)." Equality is actually popped (>=), meaning equal thresholds from two different positions only the later (higher hit-count) survives — which is semantically correct, but the invariant comment says strictly increasing. Verify that a duplicate threshold can only arise when the new point has a strictly higher hit_blocks, so the strict-increase guarantee holds after the pop. If two checkpoints at different positions happen to have the same RequiredBytes threshold (e.g. same recency rank coincidence), the pop preserves the later one (higher block count), and the resulting sequence is still strictly increasing in both fields. This appears safe, but the invariant comment on MambaRequestFact should explicitly note that equal-threshold candidates are resolved by keeping the higher hit-count (which the pop-and-replace achieves).
🤖 Generated by Qoder
| LiteHitMamba::LiteHitMamba(const Config &config) | ||
| : config_(config), pool_(config.full_charge_bytes, config.mamba_charge_bytes) { | ||
| if (config_.step_blocks == 0) { | ||
| config_.step_blocks = 1; |
There was a problem hiding this comment.
step_blocks == 0 silently clamped to 1 after construction
if (config_.step_blocks == 0) {
config_.step_blocks = 1;
}The caller in lite_hit_offline_runner.cc sets step_blocks = linear_step / block_size, and RegisterInstance already validates that linear_step is a positive whole-block multiple, so step_blocks should never be 0 at this point. The silent clamp is effectively dead code and could mask a misconfiguration upstream. Consider replacing with an assertion (assert(config_.step_blocks >= 1)) or a hard error, rather than silently correcting a value that should be impossible.
🤖 Generated by Qoder
| map.erase(previous); | ||
| } | ||
| fenwick_.AppendZero(); | ||
| full_count_fenwick_.AppendZero(); |
There was a problem hiding this comment.
fenwick_.size() used as 1-indexed position — implicit 1-based convention should be documented
fenwick_.AppendZero();
full_count_fenwick_.AppendZero();
const std::size_t current_position = fenwick_.size(); // size() after append = 1-based index of the new element
fenwick_.Add(current_position, static_cast<int64_t>(charge));This relies on DynamicFenwickTree::size() returning the count of elements after AppendZero(), which is used directly as the 1-based Fenwick position. This is a non-obvious invariant: the Fenwick tree's positions are 1-indexed (a standard Fenwick convention), and size() happens to equal the index of the most recently appended element. A comment here that positions are 1-based and size() == last element's index would prevent a future contributor from accidentally treating this as 0-based and introducing an off-by-one in RequiredBytes.
🤖 Generated by Qoder
| if (fenwick_.size() <= kCompactionSlackPositions) { | ||
| return; | ||
| } | ||
| const std::size_t positions_over_slack = fenwick_.size() - kCompactionSlackPositions; |
There was a problem hiding this comment.
Compaction threshold: active_positions >= (positions_over_slack + 1) / 2 uses integer division ceiling but the intent is unclear
const std::size_t positions_over_slack = fenwick_.size() - kCompactionSlackPositions;
if (active_positions >= (positions_over_slack + 1) / 2) {
return;
}The check skips compaction when active positions are at least half the dead positions. (positions_over_slack + 1) / 2 computes ceiling division of positions_over_slack / 2. The intent appears to be: don't compact unless at least half the Fenwick tree entries are dead (i.e. live density < 50%). But the condition mixes the "over slack" count with the live count in a way that isn't immediately obvious — a comment explaining the trigger condition ("compact when dead entries exceed live entries by at least a factor of 2") would help. Also worth verifying: when positions_over_slack == 0 (size exactly equals slack), the check 0 >= 0 is true and returns early — that's correct (no compaction needed), but the ceiling formula gives 0, so any non-zero active_positions would also satisfy it. This edge case is harmless but another reason a comment is warranted.
🤖 Generated by Qoder
| std::vector<std::pair<uint64_t, uint64_t>> pairs; | ||
| if (field.rfind(kMambaCurvePrefix, 0) == 0) { | ||
| record.is_mamba = true; | ||
| if (!ParsePairArray(field.substr(std::string(kMambaCurvePrefix).size()), pairs, error)) { |
There was a problem hiding this comment.
ParseHitCurveField creates a temporary std::string from kMambaCurvePrefix to get its size()
if (!ParsePairArray(field.substr(std::string(kMambaCurvePrefix).size()), pairs, error)) {kMambaCurvePrefix is a constexpr const char *. std::string(kMambaCurvePrefix).size() allocates a temporary std::string just to call .size(), when strlen(kMambaCurvePrefix) or a constexpr std::size_t kMambaCurvePrefixLen = 6 would be allocation-free. Minor, but worth cleaning up since this is on the CSV parsing hot path.
🤖 Generated by Qoder
| } else { | ||
| lo = mid + 1; | ||
| } | ||
| } |
There was a problem hiding this comment.
FullObjectsWithinBytes: binary search returns 0 when lo > n, but the invariant should be lo == n + 1
if (lo > n) {
return 0;
}
return full_count_fenwick_.PrefixSum(n) - full_count_fenwick_.PrefixSum(lo - 1);After the binary search, lo is in [1, n+1]. The condition lo > n covers only lo == n + 1 (since lo is bounded to n + 1), meaning nothing in the pool fits the budget. The comment "n + 1 means nothing fits" documents this. The check is correct, but writing lo == n + 1 instead of lo > n would more precisely express the invariant and prevent readers from wondering about larger values.
🤖 Generated by Qoder
0059f51 to
49b9dcf
Compare
A linear-attention request can only continue from a checkpoint, and a checkpoint is usable only if its Mamba state and the entire Full prefix in front of it are both cached. The two object types have different byte sizes, so residency cannot be decided on a block-count axis. WeightedLruPool generalizes the LiteHit stack: one recency order over typed objects (kFull / kMamba) with a fixed byte charge per type, where range sums return resident BYTES, so residency for a byte capacity is RequiredBytes(key) <= capacity_bytes - byte-weighted Mattson inclusion, which keeps one replay exact for every capacity. A parallel order-statistics tree counts Full markers at the same positions, which is what the resident-unique reporting needs once the main tree stores bytes. LiteHitMamba evaluates every checkpoint candidate on the request-start snapshot and emits a capacity-independent MambaRequestFact: an explicit step function on the total byte axis. Unlike the full-attention arithmetic runs it cannot be run-length encoded, because the interleaved Mamba markers break the equal-step property of consecutive blocks. Both object types share one pool and one budget, so a request's own checkpoints compete with its Full blocks for the same bytes. The facts pipeline gains a "mamba:" row encoding so the offline post-query can project the step curve for any byte capacity through the same projector as the online path.
…y indexer Linear-attention instances were simulated by the legacy LruCacheIndexer, whose model fuses a Full block and its Mamba checkpoint into one cache entry whose charge doubles as the checkpoint tag. The two could therefore never be evicted independently - exactly the case that matters, since a Full block stays warm through prefix reuse while its checkpoint goes cold. It also built one real LRU per capacity tier, had no facts pipeline, and inferred a checkpoint by comparing a stored charge against a constant. InstanceState now holds exactly one analyzer: LiteHit for full attention, LiteHitMamba for linear attention. The legacy indexer, its factory, its TTL wrapper and its hit-age histogram are deleted (the histogram had no proto field and no client or dashboard consumer). Linear attention with a group TTL is rejected at registration: the Mamba core carries no time axis yet, and the offline runner already refused that combination. Two byte quantities replace the fused wording internally: a Full block always costs full_charge_bytes and a checkpoint additionally stores mamba_charge_bytes, instead of adding them into size_full_linear and subtracting them back out at every use. The register response is unchanged; the fused value is composed once at the proto boundary. The block count in that response is now computed explicitly instead of dividing by an "average bytes per block" that does not exist when two charges differ. It stays exact for full attention, where it doubles as the projection slot, and is labelled an estimate for linear attention, where hits are decided on the byte axis and nothing consumes the estimate.
…etire LiteHitMamba LiteHitMamba was a second core that re-implemented snapshot evaluation and commit ordering next to LiteHit, on the same WeightedLruPool. The two copies had already diverged on TTL: the core carried a position-epoch time axis that the Mamba copy lacked, which is why linear attention + TTL had to be rejected at registration. LiteHit is now the single core for both attention modes. Linear attention becomes LiteHitLinearPolicy, a stateless helper that reuses the core pool: RequiredLinearBytes answers thresholds through the same Fenwick prefix sums, and CommitLinearIfNeeded writes Linear states at periodic positions plus the request tail. EvaluateRecoveryCurve merges the Full prefix threshold with the Linear-state threshold via max() and keeps a monotone envelope, so a restore point is usable exactly when its state and its entire Full prefix are resident. TTL moves out of the core into the TtlLiteHit decorator. The core no longer understands time, only an alive_from_position boundary; the decorator owns the epoch queue and advances the watermark before every evaluation. Because the boundary lives on the shared position axis, TTL now covers Linear states for free, removing the linear + TTL restriction at its root. WeightedLruPool's existing position compaction now takes the alive boundary into account and returns a PositionRemap, so the decorator outside the pool can relocate its epoch boundaries after a compaction instead of the pool reaching back into TTL state it no longer owns.
With one core serving both attention modes, the facts CSV no longer needs a Mamba-specific wording. The default hit_curve encoding is the explicit byte-step curve tagged "bytes:"; Full-only rows keep the smaller arithmetic-run RLE on the block axis, now tagged "rle:" instead of being the untagged default. The record field is_mamba becomes is_full_rle accordingly: the row describes its encoding, not the instance kind. The parser still accepts legacy "mamba:" byte steps and unprefixed Full RLE rows, so facts files written before this change stay projectable, and the projection itself keeps going through the same HitCurveProjector as the online path.
InstanceState now holds exactly one analyzer for every instance: a TtlLiteHit decorator around the shared LiteHit core, a timestamp-free pass-through when the group has no TTL. The full/linear split disappears from the manager, the offline runner and the service: both modes register the same way, feed the same ProcessRequest/ProcessFullRequest entry points and project through the same HitCurveProjector. Linear attention with a group TTL is accepted now that the TTL watermark lives on the shared position axis; the registration-time rejection and the offline runner refusal are gone. Capacity slots are kept in bytes for both modes; the Full RLE path performs its block floor only at projection. estimated_capacity_blocks in the register response becomes a compatibility/display estimate that nothing at runtime consumes, and the per-instance bytes_per_block summary turns into a double: exact Full charge for full attention, resident working-set average for linear. Proto field numbers and types are unchanged apart from that comment wording.
49b9dcf to
0ae911b
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0ae911bb41
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| } | ||
| // 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) { |
There was a problem hiding this comment.
Preserve the existing block units for linear_step
Existing clients and registry records store linear_step in blocks, as defined by the previous protobuf contract, but this validation now interprets the same unversioned field as tokens. After an upgrade, common persisted values such as linear_step=3 with block_size=16 fail Recover() with EC_BADARGS; values that happen to be divisible silently use a different checkpoint interval. Keep the wire/persistence unit in blocks or add an explicit versioned migration before applying token-based validation.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Reviewed the reworked PR (unified TtlLiteHit decorator + LiteHit core for both Full and Linear attention, retiring the legacy CacheIndexer path). The implementation is well-structured: the byte-weighted LRU pool with Fenwick order-statistics, the stateless Linear policy, and the epoch-based TTL watermark all check out under careful analysis. No new correctness issues found beyond the 6 patterns already called out in the previous review comments (step_blocks clamp, Pareto sweep, Fenwick convention, compaction threshold, std::string allocation, covered_blocks semantics) — all of which remain present on the restructured files/lines. Good test coverage with randomized differential testing and oracle comparisons.
🤖 Generated by Qoder
背景
optimizer 里只有 full attention 有精确的回放核心(LiteHit),线性注意力(Mamba)实例一直挂在老的
LruCacheIndexer上模拟。这两种注意力复用缓存的方式不一样。full attention 是逐块前缀命中;线性注意力必须从某个 Linear state 接着往下算,而一个 state 能不能用要同时满足两个条件:它自己在缓存里,它前面那一整段 Full block 也都在。
老 indexer 表达不了这件事。它把一个 block 和这个 block 上的 state 塞进同一个 cache entry,还拿 entry 的 charge 兼做"这是不是 state"的标记,所以两者只能一起进、一起出。线上实际情况往往相反:Full block 被前缀复用反复刷新,Linear state 却很久没人用。它还给每个容量档位各建一个真的 LRU,档位一多,时间和内存都成倍涨。
方案
最终形态是一个 LiteHit 核心,不再有独立的第二套分析器:
RequiredBytes <= capacity_bytes——字节加权的 Mattson 栈包含性,一次回放同时回答所有容量档。LiteHitLinearPolicy不持有任何自己的 LRU 状态,完全复用核心 pool;评估时用max(Full 前缀门槛, Linear state 门槛)合并成单调 envelope,提交时在周期位置和请求尾写入 Linear state。alive_from_position存活边界;TtlLiteHit装饰器持有请求时间 epoch 队列(O(Q)),每次评估前推进水位线。TTL 边界在共享位置轴上,天然同时覆盖 Full 与 Linear——原先"linear + TTL 注册即拒绝"的限制从根上取消。TTL 为 0 时装饰器透明透传。每个请求产出一条容量无关曲线:默认是总字节轴的 byte-step
RequestFact;Full-only 等 charge 场景无损压缩成 block 轴等差段 RLE(FullRequestFact)。facts CSV 的 hit_curve 用bytes:/rle:前缀区分两种编码,读取端兼容旧的mamba:与无前缀行,历史 facts 文件可继续投影;事后查任意容量与在线使用同一个HitCurveProjector。口径变化
capacity_bytes,Full RLE 只在最终投影时做一次 floor。estimated_capacity_blocks由字节容量派生,运行时没有任何逻辑消费它;full attention 精确、linear 明确是估算。bytes_per_block汇总改为 double。 full attention 是精确的 Full charge;linear 是当前 resident working set 字节数除以 resident Full block 数。兼容性
mamba:byte-step 与无前缀 Full RLE 行测试
--config=debug --config=asan)19/19,rebase 到最新 main 后验证kv_cache_manager/共享代码)LiteHitLinearTest)LiteHitTtlTest)TraceQuery,逐容量比对