fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor - #340
Open
dejay2 wants to merge 1 commit into
Open
fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor#340dejay2 wants to merge 1 commit into
dejay2 wants to merge 1 commit into
Conversation
… 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
One-line fix to the boot-time
--moe-cache-autobudget split, plus the--kv-reserve-tokenshelp text that documents it:Why
plan_cache_budgetreasons in usable pages, butcreate_kv_pool(config, num_pages, ...)allocatesnum_pages + 1: every pool family keeps page 0 as anunreachable dummy/sentinel.
The plan prices the KV half as
num_pages * cache_per_pageand then lets expertsgreedily take everything the reserve does not protect (the fill is deliberately
headroom-free; see the
total <= budget_bytesassert immediately below it). Becausethe 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_pagethat is not a rounding error, and the failure modeis precisely the one
--moe-cache-autoexists to prevent: an arithmetic plan thatclaims to fit, followed by a CUDA OOM in the KV allocation at boot.
The same off-by-one made
--kv-reserve-tokens Nunder-deliver. It is documented as aKV-cache token floor, but
ceil(N / page_size)usable pages minus the dummy is lessthan
Ntokens 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.
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_genericand::test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backendbothrequire 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 256atpage_size=64yieldspages - 1usable pagescovering at least 256 tokens. On
origin/mainthis returns(size=9, pages=4); with thefix 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