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.

@gdevenyi

gdevenyi commented Sep 4, 2026

Copy link
Copy Markdown

Tested on 2 x RTX 6000 Ada (48 GiB, sm_89) serving RadixArk/Qwen3.8-Flash-Next-NVFP4 at TP=2 (#385) with --moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --max-running-requests 16 and the implicit --moe-cache-auto. On main this flag combination OOMs at init because the auto-sized expert cache ignores --num-tokens (#383).

With this PR's engine hunk applied alone (my local equivalent clamp reverted; the test hunk not applied because it collides with my own regression test), the KV pool comes up at 262,144 tokens and the expert cache plans around it:

single-stream tok/s 8 concurrent tok/s expert residency TTFT (1k prompt)
main + local clamp 90.8 326.0 94.5% 0.81 s
main + #198 90.4 322.2 94.5% 0.85 s

Same behaviour and same numbers within noise, so this fixes #383 for that path.

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.

2 participants