Skip to content

fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor - #340

Open
dejay2 wants to merge 1 commit into
FlashML-org:mainfrom
dejay2:pr/cache-budget-split
Open

fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor#340
dejay2 wants to merge 1 commit into
FlashML-org:mainfrom
dejay2:pr/cache-budget-split

Conversation

@dejay2

@dejay2 dejay2 commented Sep 2, 2026

Copy link
Copy Markdown

What

One-line fix to the boot-time --moe-cache-auto budget split, plus the
--kv-reserve-tokens help text that documents it:

kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + 1

Why

plan_cache_budget reasons in usable pages, but create_kv_pool(config, num_pages, ...) allocates num_pages + 1: every pool family keeps page 0 as an
unreachable dummy/sentinel.

The plan prices the KV half as num_pages * cache_per_page and then lets experts
greedily take everything the reserve does not protect (the fill is deliberately
headroom-free; see the total <= budget_bytes assert immediately below it). Because
the priced count is the usable one, the plan under-counted the KV pool by exactly one
page's bytes, and that page came out of memory already promised to expert slots. On a
model with a large cache_per_page that is not a rounding error, and the failure mode
is precisely the one --moe-cache-auto exists to prevent: an arithmetic plan that
claims to fit, followed by a CUDA OOM in the KV allocation at boot.

The same off-by-one made --kv-reserve-tokens N under-deliver. It is documented as a
KV-cache token floor, but ceil(N / page_size) usable pages minus the dummy is less
than N tokens of reachable capacity. Hence also the help-text change.

Provenance

This hunk was extracted from a larger fork commit whose message carries no recorded
symptom. The reasoning above is reconstructed from the code paths named, not from a
logged OOM; the new test shows the arithmetic changes exactly as described.

How it was tested

Windows 11, RTX 5090.

python -m pytest -q tests/engine tests/server
  • before this commit: 643 passed, 2 failed
  • after this commit: 644 passed, 2 failed (1 new test)

The 2 failures are pre-existing on this box and unrelated to this change:
tests/engine/test_cache_budget.py::test_adjust_config_resolves_num_tokens_generic and
::test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend both
require flashinfer, which is not installed here.

New test, tests/engine/test_cache_budget.py::test_resolve_auto_reserves_usable_tokens_beyond_the_dummy_page:
pins that --kv-reserve-tokens 256 at page_size=64 yields pages - 1 usable pages
covering at least 256 tokens. On origin/main this returns (size=9, pages=4); with the
fix it returns (size=8, pages=14), so the test genuinely discriminates.

What is NOT included

The originating fork commit also clamps the auto-sized KV pool to one full-length
context. That bounds aggregate KV under --max-running-requests > 1 (the default is 4)
and trades cross-request cache residency for VRAM, which is a policy choice rather than
a bug fix, so it is deliberately left out of this PR.

🤖 Generated with Claude Code

https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK

… floor

plan_cache_budget works in USABLE pages, while create_kv_pool allocates
num_pages + 1: every pool family keeps page 0 as an unreachable
dummy/sentinel. The plan prices the KV half as num_pages * cache_per_page
and hands experts everything the reserve does not protect (the fill is
greedy and deliberately leaves no headroom), so it under-counted the KV
pool by exactly one page's worth of bytes, and that page was then
allocated out of memory the split had already given to expert slots. On
a model with a large cache_per_page that page is not small, and the
failure mode is the one --moe-cache-auto exists to prevent: an
arithmetic plan that says it fits, followed by a CUDA OOM in the KV
allocation at boot.

The same off-by-one made --kv-reserve-tokens quietly under-deliver: it
is documented as a KV-cache token floor, but ceil(N / page_size) usable
pages minus the dummy is less than N tokens of reachable capacity.
Reserve the sentinel page on top of the user-visible floor and say so in
the --kv-reserve-tokens help text.

Tested (Windows, RTX 5090):

  python -m pytest -q tests/engine tests/server

  before: 643 passed, 2 failed   after: 644 passed, 2 failed

The two failures are pre-existing on this box and unrelated: both
tests/engine/test_cache_budget.py::test_adjust_config_* need flashinfer,
which is not installed here.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant