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

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.

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
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