Skip to content

fix(engine): reserve explicit KV pages during MoE auto-sizing - #198

Open
pedro-moser wants to merge 1 commit into
FlashML-org:mainfrom
pedro-moser:fix/moe-auto-compat
Open

fix(engine): reserve explicit KV pages during MoE auto-sizing#198
pedro-moser wants to merge 1 commit into
FlashML-org:mainfrom
pedro-moser:fix/moe-auto-compat

Conversation

@pedro-moser

Copy link
Copy Markdown

Summary

  • include an explicit --num-pages KV allocation in the reserve passed to the joint MoE/KV auto-sizing solver
  • also covers --num-tokens, which resolves to num_page_override after page size is finalized
  • preserve the existing kv_reserve_tokens and model-specific minimum-reserve floors by taking the maximum of all three

Problem

Engine._resolve_auto_moe_cache_size() runs before the KV pool is allocated. It previously passed only:

max(config.kv_reserve_tokens, min_reserve)

When a user explicitly requested more KV pages/tokens, that allocation did not participate in the MoE cache budget. The solver could therefore assign those bytes to expert slots and leave startup to fail later during KV allocation.

This is the explicit-context half of the tradeoff described in #111. It does not change the default 8,192-token reserve or admission behavior for over-capacity requests.

Test plan

Regression test on current main:

  • configure 64 pages with a resolved page size of 16
  • intercept the arguments passed to resolve_moe_cache_auto
  • verify kv_reserve_tokens == 64 * 16

The test failed before the fix with 0 != 1024 and passes after it.

Executed:

PYTHONPATH=$PWD/python pytest -q tests/engine/test_cache_budget.py
24 passed in 1.50s

git diff --check and Python compile checks also pass.

Risk and exclusions

  • No behavior changes when neither --num-pages nor --num-tokens is set.
  • An explicit context request can now make the budget solver reject startup earlier instead of allowing a later CUDA OOM; that is intentional.
  • This PR does not change the default MoE/KV split, scheduler admission, GGUF cache layout, or parser behavior.

Addresses part of #111.

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