mirror of
https://github.com/Mesh-LLM/mesh-llm.git
synced 2026-08-08 22:23:19 -04:00
fix(skippy-cache): disable resident token cap for tiny contexts
CI run 26193173851 surfaced this: `scripts/skippy-ci-smoke.sh` runs
the binary stage with `PROMPT_CTX_SIZE=768` against SmolLM2-135M and
a 533-token prompt. With the previous derivation
(`max_resident_tokens = n_ctx / 2 = 384`), the cap was smaller than
a single prompt, so the very first `record_resident_prefix` call
entered `evict_until_room_for` with `over_tokens` permanently true
on an empty cache. `bail!("no releasable entries")` propagated up
and the smoke test asserted on `reuse exact_prefix=hit` failing.
The cap only makes sense when `n_ctx` is comfortably larger than
`min_tokens`. Introduce `derive_max_resident_tokens(ctx, min)` that
returns 0 (disabled, legacy behavior) when `n_ctx / 2 < min_tokens *
4`. Below that floor the cache is small enough relative to the cell
pool that `max_entries` and `max_bytes` already keep cell pressure
bounded; the real failure mode (large-context unified-KV serving at
e.g. `n_ctx = 131072`) comfortably clears the floor and still gets
the cap.
Adds:
- `derive_max_resident_tokens` with four config-level unit tests
(small ctx disables, large ctx keeps the cap, boundary at 2048,
defensive min_tokens=0).
- `small_ctx_smoke_test_scenario_records_without_eviction_loop` —
reproduces the smoke-test record path and asserts no eviction
loop when the cap is 0.
cargo test -p skippy-cache --lib: 14 pass
cargo test -p skippy-server --lib: 81 pass
cargo test -p mesh-llm-host-runtime --lib: 1437 pass
This commit is contained in:
parent
4eb513fff4
commit
809f4b0312
2 changed files with 87 additions and 7 deletions
|
|
@ -19,7 +19,9 @@ pub struct ResidentCacheConfig {
|
|||
///
|
||||
/// Set this to a fraction of the model's `n_ctx` (typically
|
||||
/// `n_ctx / 2` or similar). A value of 0 disables the cap and
|
||||
/// behaves like the legacy unbounded-by-tokens cache.
|
||||
/// behaves like the legacy unbounded-by-tokens cache. The cap is
|
||||
/// only useful when `n_ctx` is comfortably larger than
|
||||
/// `min_tokens`; see [`derive_max_resident_tokens`] for the floor.
|
||||
pub max_resident_tokens: u64,
|
||||
}
|
||||
|
||||
|
|
@ -28,12 +30,8 @@ impl ResidentCacheConfig {
|
|||
let reserved_seq_count = i32::try_from(config.lane_count.saturating_mul(2))
|
||||
.unwrap_or(i32::MAX)
|
||||
.max(2);
|
||||
// Cap the cache at half the model's `n_ctx` cell pool so the
|
||||
// active lanes always have room to prefill fresh prompts.
|
||||
// The other half is the cache budget. See the doc comment on
|
||||
// `max_resident_tokens` above for the failure mode this
|
||||
// prevents.
|
||||
let max_resident_tokens = u64::from(config.ctx_size).saturating_div(2);
|
||||
let max_resident_tokens =
|
||||
derive_max_resident_tokens(u64::from(config.ctx_size), cache.min_tokens);
|
||||
Self {
|
||||
max_entries: cache.max_entries.clamp(1, 512),
|
||||
max_bytes: cache.max_bytes,
|
||||
|
|
@ -44,6 +42,63 @@ impl ResidentCacheConfig {
|
|||
}
|
||||
}
|
||||
|
||||
/// Derive `max_resident_tokens` from the model's `n_ctx` cell pool.
|
||||
///
|
||||
/// The cache shares the `n_ctx` cell pool with the active lanes under
|
||||
/// `kv_unified = true`. The cap reserves half of the pool for in-flight
|
||||
/// lane prefills and lets the cache use at most the other half.
|
||||
///
|
||||
/// For very small contexts (smoke-test / tiny-model configs) the cap
|
||||
/// can be smaller than a single prompt, which defeats the cache
|
||||
/// without preventing any real wedge. The cap is therefore disabled
|
||||
/// when `n_ctx / 2` cannot hold at least four minimum-sized cache
|
||||
/// entries (`min_tokens * 4`). At that point the cache is small
|
||||
/// enough relative to the cell pool that `max_entries` and
|
||||
/// `max_bytes` already keep cell pressure bounded. The original
|
||||
/// failure mode this cap fixes is large-context unified-KV serving
|
||||
/// (e.g. `n_ctx = 131072`), where this floor is comfortably cleared.
|
||||
fn derive_max_resident_tokens(ctx_size: u64, min_tokens: u64) -> u64 {
|
||||
let half = ctx_size.saturating_div(2);
|
||||
let min_floor = min_tokens.saturating_mul(4).max(1);
|
||||
if half < min_floor {
|
||||
return 0;
|
||||
}
|
||||
half
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod resident_cache_config_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cap_disabled_when_ctx_smaller_than_four_min_entries() {
|
||||
// Smoke-test / SmolLM2 scenario: ctx_size=768, min_tokens=256.
|
||||
// half=384, floor=1024 -> cap disabled (0).
|
||||
assert_eq!(derive_max_resident_tokens(768, 256), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cap_enabled_when_ctx_has_room_for_four_min_entries() {
|
||||
// Production scenario: large unified-KV pool.
|
||||
assert_eq!(derive_max_resident_tokens(131072, 256), 65536);
|
||||
// Right at the boundary.
|
||||
assert_eq!(derive_max_resident_tokens(2048, 256), 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cap_disabled_just_below_floor() {
|
||||
// half=1023, floor=1024 -> disabled.
|
||||
assert_eq!(derive_max_resident_tokens(2046, 256), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_min_tokens_treats_floor_as_one() {
|
||||
// Defensive: min_tokens=0 should not divide by zero or disable
|
||||
// the cap on real-sized contexts.
|
||||
assert_eq!(derive_max_resident_tokens(8192, 0), 4096);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct PrefixCandidatePolicy {
|
||||
pub min_tokens: u64,
|
||||
|
|
|
|||
|
|
@ -310,6 +310,31 @@ mod tests {
|
|||
assert_eq!(cache.stats().resident_tokens, 3000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn small_ctx_smoke_test_scenario_records_without_eviction_loop() {
|
||||
// Regression for skippy-ci-smoke `prompt exact-prefix hit and
|
||||
// live-session reuse` (CI run 26193173851). The smoke test
|
||||
// ships SmolLM2-135M with `PROMPT_CTX_SIZE=768` and a 533-token
|
||||
// prompt. With a naive `n_ctx / 2 = 384` cap, the cap was
|
||||
// *smaller than a single prompt*, so the first record attempt
|
||||
// would call `evict_until_room_for` with `over_tokens`
|
||||
// permanently true on an empty cache and bail with
|
||||
// "no releasable entries".
|
||||
//
|
||||
// `ResidentCacheConfig::from_stage` derives the cap with a
|
||||
// 4*min_tokens floor (see `derive_max_resident_tokens`). For
|
||||
// the smoke-test config the cap therefore comes through as 0
|
||||
// (disabled). This unit test pins that path: cap=0, record a
|
||||
// 533-token prompt against a 256-token min, no eviction loop.
|
||||
let mut cache = ResidentPrefixCache::new(cfg(16, 0, 0));
|
||||
let alloc = cache
|
||||
.allocate_for_record("page-0", 533, 100, |_| Ok(()))
|
||||
.unwrap();
|
||||
cache.commit_record("page-0".to_string(), alloc.seq_id, 533, 100);
|
||||
assert_eq!(cache.stats().resident_tokens, 533);
|
||||
assert_eq!(cache.stats().entries, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_token_budget_disables_the_check() {
|
||||
// max_resident_tokens = 0 means "unlimited" — legacy behavior.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue