fix(engine): reserve explicit KV pages during MoE auto-sizing - #198
Open
pedro-moser wants to merge 1 commit into
Open
pedro-moser wants to merge 1 commit into
pedro-moser wants to merge 1 commit into
Conversation
|
Tested on 2 x RTX 6000 Ada (48 GiB, sm_89) serving RadixArk/Qwen3.8-Flash-Next-NVFP4 at TP=2 (#385) with 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:
Same behaviour and same numbers within noise, so this fixes #383 for that path. |
nomanoma121
pushed a commit
to nomanoma121/My-FreeToken
that referenced
this pull request
Sep 16, 2026
…nd pays for the dummy page Two ways the auto-sized plan could spend more VRAM than its budget, both reported upstream. An explicit --num-pages / --num-tokens is kept as the KV size -- the caller only writes the plan's pages back when no override was given -- but the plan reserved just --kv-reserve-tokens for KV and let the experts greedily take the rest. The larger explicit pool was then allocated on top of a budget the experts had already filled, and the boot OOMed in the KV allocation (upstream FlashML-org#383, fix as in FlashML-org#198). The plan now reserves the larger of the two. Every KV pool allocates one page past the usable ones for padded rows to write into (create_kv_pool and rebuild_from_config: num_pages + 1), and the plan priced only the usable pages, so a plan that filled the budget exactly allocated one page past it (upstream FlashML-org#340). Upstream's patch adds the page to the reserve; here the pools already treat num_pages as usable, so the page is charged to the budget instead -- the same accounting DSV4's solver already does. required_bytes counts it too, which moves the runtime rebuild's fit check and the budget-too-small suggestions onto the same arithmetic, so a suggested --kv-reserve-tokens still fits. The cost is one page of KV per plan (one token on a page_size 1 pool). The (1 - memory_ratio) headroom was absorbing it, which is why nothing on the 2060 or the 3060s ever tripped over it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
MT-z
added a commit
to MT-z/FreeToken
that referenced
this pull request
Sep 16, 2026
…with FlashML-org#393) Artemowka22's FlashML-org#224, open since 2026-08-26. FlashML-org#393 declines logprobs on both routes and its own compatibility table points here ("400 (see FlashML-org#224 for the sampled-token logprobs)"), so the two are complementary by design -- and both rewrite the same sampling and message plumbing, which is where the work was. Before taking it, measured what this branch actually did with a logprobs request: /v1/completions logprobs=5 400 "logprobs is not supported" declared /v1/chat/... logprobs=true, top_logprobs=5 200, no logprobs silent `ChatCompletionRequest` carries `extra="allow"` and declares neither field, so they were swallowed before any validator saw them. FlashML-org#393's table says this route returns 400; it returned 200 and dropped them. Same on pr393 alone, so it is upstream, not this stack. Merge: 9 files, 20 conflict blocks, all from FlashML-org#393's `n` sampling (many uids per request) meeting FlashML-org#224's single-uid shape. Resolved by keeping FlashML-org#393's structure and hanging the logprobs off it -- the multi-uid `choices` loop, `_completion_chunk`, `_resolve_sampling` (which also fills min_p, penalties, logit_bias, stop_token_ids) all stay. FlashML-org#224's rejections of `echo`, `suffix` and `logit_bias` were dropped: FlashML-org#393 implements all three, and those lines are older than it. Its `echo` + `logprobs` rejection is kept -- that one is real, the prompt logits are not there. Two failures had to be split apart before either could be fixed: - `tests/scheduler/test_abort_inflight_prefill.py`, 4 tests: **pr224 fails these on its own base too**, so it is the PR's regression, not the merge. `ForwardOutput` grows from 3 fields to 6, `_process_last_data` switches to reading it by name, and the upstream test (there since 3af9d90) hands it a bare 3-tuple. Fixed by having the test build a real `ForwardOutput`. - `tests/engine/test_sample_logprobs.py`, 1 test: this one is the merge. FlashML-org#393 adds the `needs_logits_processing` property to SamplingParams and FlashML-org#224's new stub is a SimpleNamespace without it. Fixed in the stub, not by loosening `_plan` -- a getattr default there would hide a real type mismatch. Third instance of this pattern today (FlashML-org#354's `k_scale`, FlashML-org#198's direct `num_page_override`). Verified on the wire, Ornith-1.5-35B-A3B-NVFP4: /v1/completions logprobs=3 ' Paris' -0.7205, top3 [' Paris' -0.72, ' a' -2.16, '\n' -2.60], text_offset [0,6,7,8]; streaming carries per-chunk logprobs /v1/chat/... with the default qwen3 reasoning parser: 400 naming the reason ("reasoning tokens are hidden from message content ... use /v1/completions"); with --reasoning-parser off: full entries with token/logprob/bytes/top_logprobs, streaming too FlashML-org#393's side survives: `n=2` returns 2 choices with correct usage, `echo` echoes, `suffix` and `logit_bias` are accepted. FlashML-org#224's own guards fire: `echo`+`logprobs` and `logprobs=9` both 400 with their reasons. Default path, no logprobs requested, 3 single-stream decodes of 399 tokens: try/all 91.7 / 95.2 / 94.1 tok/s median 94.1 try/224 92.5 / 95.3 / 94.4 tok/s median 94.4 Full suite: 11 failed, 1849 passed, 60 skipped -- the same 11 as try/all, plus the 13 tests the PR adds. Assisted-by: Claude Opus 5
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 18, 2026
…MoE auto-sizing # Conflicts: # tests/engine/test_cache_budget.py
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.
Summary
--num-pagesKV allocation in the reserve passed to the joint MoE/KV auto-sizing solver--num-tokens, which resolves tonum_page_overrideafter page size is finalizedkv_reserve_tokensand model-specific minimum-reserve floors by taking the maximum of all threeProblem
Engine._resolve_auto_moe_cache_size()runs before the KV pool is allocated. It previously passed only: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:resolve_moe_cache_autokv_reserve_tokens == 64 * 16The test failed before the fix with
0 != 1024and passes after it.Executed:
git diff --checkand Python compile checks also pass.Risk and exclusions
--num-pagesnor--num-tokensis set.Addresses part of #111.