--moe-cache-auto ignores --num-tokens / --num-pages: expert fill overruns the budget and the KV pool OOMs at boot

[gdevenyi]

Before you start

• I have read the FAQ and my problem is not answered there.
• I have read the Roadmap and this is not already planned there.
• I have searched existing issues and found no duplicate (fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor #340 fixes a sibling off-by-one in the same resolver, not this).
• I am on a freshly rebuilt main (86214a9).

What happened

--moe-cache-auto (the default) ignores --num-tokens / --num-pages when it plans the expert-slot fill, so a large explicit KV size boots into a CUDA OOM.

What I did: serve RadixArk/Qwen3.8-Flash-Next-NVFP4 (qwen4_exp, 512 experts x 48 layers, 262,144-token context) on a 48 GiB RTX 6000 Ada with the KV pool sized to the model's full context:

ft serve --model <RadixArk NVFP4 snapshot> --moe-backend offload --num-tokens 262144 \
  --max-running-requests 8 --cuda-graph-max-bs 8

What I expected: the auto planner reserves the 262,144 tokens of KV (+ the GDN state pool) and fills the rest of the budget with expert slots.

What happened instead: the planner reserves only kv_reserve_tokens (the small default floor), fills the budget with expert slots, logs --moe-cache-auto resolved moe_cache_size=... num_pages=... with a page count far below the request, and then the pool honours --num-tokens anyway: create_kv_pool allocates 4096 pages (6.2 GiB of QSA KV) plus the GDN state pool on top of an already-full card, and boot dies with torch.OutOfMemoryError in the KV / linear-state allocation after the experts are resident. Small --num-tokens values happen to fit inside the planner's slack, so this only shows on models with a large context and a per-token KV cost.

Root cause

Engine._resolve_auto_moe_cache_size (python/freetoken/engine/engine.py) passes kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve) to resolve_moe_cache_auto. config.num_page_override (set by --num-tokens / --num-pages) is not consulted, so the plan and the pool disagree about the KV size whenever the override exceeds the floor. The override is a requirement, not a floor: the pool will allocate exactly that many pages later.

Fix

Clamp the reserve to the override before planning; the expert fill then sees the KV the pool will actually allocate:

kv_reserve_tokens = max(config.kv_reserve_tokens, min_reserve)
if config.num_page_override is not None:
    # --num-pages / --num-tokens is a requirement, not a floor: plan the expert fill
    # against the KV the pool will actually allocate, or the two together overrun
    # the budget and the pool OOMs after the experts are already resident.
    kv_reserve_tokens = max(kv_reserve_tokens, config.num_page_override * page_tokens)
return resolve_moe_cache_auto(..., kv_reserve_tokens=kv_reserve_tokens, page_size=page_tokens, ...)

With this, the same command boots (moe_cache_size=8984 num_pages=4096, 2.6 GiB free after init) and serves the full 262k context (a 261,988-token prompt completes). A regression test added to tests/engine/test_cache_budget.py::test_engine_resolve_auto_moe_cache_size_maps_kwargs (a stub config with num_page_override = 20 against a stub without it) fails on main and passes with the clamp; tests/engine is 108 passed with it. I can open a PR with the fix + test if wanted.

Related: #340 adds the pool's dummy page to the same reserve; both apply together cleanly (measured: plan moves from 8984/4096 to 8983/4097, boot and throughput unchanged).

How did you install FreeToken

Built from source

FreeToken version

86214a9

OS

Ubuntu

OS details

Ubuntu 24.04, kernel 7.1.8 (Zabbly), no sudo on the box

GPU and driver

2x RTX 6000 Ada 48 GB (sm_89), CUDA 13.3 toolkit, torch 2.11.0+cu130

CPU and system RAM

2x Xeon Gold 6526Y, 503 GiB RAM

Checkpoint

RadixArk/Qwen3.8-Flash-Next-NVFP4 (also reproduces with Qwen/Qwen3.8-Flash-Next-FP8)

Command

ft serve --model <RadixArk NVFP4 snapshot> --moe-backend offload --num-tokens 262144 --max-running-requests 8 --cuda-graph-max-bs 8

Full log

The failing boot log was not kept (the run predates my harness); the planner line and the OOM location are as described above. Happy to reproduce and attach one on request.

[MT-z]

The clamp from the issue body was applied here and run on a different card and model.
It fixes the boot. Three notes came out of applying it, one of which is about the
regression test you describe.

--moe-cache-auto, asking for 262,144 tokens of KV	plan	KV pool	boot
--num-tokens 262144, main	9,550 / 9,559 / 9,561 slots, 8,255 / 8,255 / 8,193 pages	262,144 tokens, 5.00 GiB	torch.OutOfMemoryError ×3
--kv-reserve-tokens 262144, main	6,631 / 6,631 / 6,628 slots, 262,187 / 262,207 / 262,192 pages	plan, 5.00 GiB	serves ×3
--num-tokens 262144, + the clamp	6,623 / 6,627 / 6,628 slots, 262,222 / 262,191 / 262,147 pages	262,144 tokens, 5.00 GiB	serves ×3

Row one is the reported failure: Tried to allocate 5.00 GiB ... 3.87 GiB is free, from
mha_pool.py:50, three times out of three. The planned 8,193 / 8,255 pages are the
default kv_reserve_tokens of 8,192 at page_size 1, plus whatever the integer-division
expert fill left over -- num_pages = max(remaining // cache_per_page, kv_reserve_pages)
with the experts sized against that same reserve. So the plan is the default floor while
the pool goes on to allocate the 262,144 that was asked for.

Row two is the first of the three, and what this box can add. The same 262,144 tokens
asked for through --kv-reserve-tokens plans and allocates consistently, because that is
the value the planner reads. Under --moe-cache-auto the two flags are not
interchangeable, which is easy to miss from the help text.

On the regression test. A stub with num_page_override = 20 against a stub without it
gives the same (8, 81, True) here on main and with the clamp, so on that stub it does
not separate them. No value does, and the reason is the stub's shape rather than the
constant: its 4 experts over 2 layers all fit in the cache, so the expert fill is already
at its cap and KV holds the whole remainder -- an override at or below that changes
nothing, one above it cannot fit. Widening the stub does open a window, but a narrow one:
16 experts over 4 layers separates only at 81, and 64 over 8 only at 78-80. A test
pinned to a value in a three-wide window looked too brittle to suggest.

What separates them on the stub as it stands is the failure moving into the plan, which
is what the fix is for. This fails on main and passes with the clamp here (helpers
elided -- the stub is the one already in that test file, with num_page_override added):

def test_num_page_override_is_planned_against_not_ignored():
    natural = _engine()._resolve_auto_moe_cache_size(_config(None), _Banks())[1]
    with pytest.raises(AssertionError, match="cache budget too small"):
        _engine()._resolve_auto_moe_cache_size(_config(natural + 1), _Banks())

Two smaller things from the same run. Read as posted, config.num_page_override raises
AttributeError on the StubConfig in
tests/engine/test_cache_budget.py::test_engine_resolve_auto_moe_cache_size_maps_kwargs,
which has no such attribute -- other stubs in that file do; getattr(config, "num_page_override", None) leaves tests/engine at the 108 passed it collects on
af71ba4 here, 109 with the test above.
And the assertion the clamp now routes into reads raise memory_ratio, lower kv_reserve_tokens, or free GPU memory (cache_budget.py:89), which names a flag the
user did not set when they arrived by --num-tokens.

Measured on an RTX 4090 (24 GiB, sm_89) / i9-14900KF / 61 GiB box, driver 595.84, CUDA
13.3, torch 2.11.0+cu130, triton 3.6.0, sgl_kernel 0.4.5, freetoken 0.1.2, model
ornith-ai/Ornith-1.5-35B-A3B-NVFP4 (qwen3_5_moe, NVFP4 experts). Main af71ba4, which
has 86214a9 as an ancestor, with nothing else applied; the clamp arm is that plus the
four lines from the issue. One flag varied per run, three runs per arm:

ft serve --model-path ornith-ai/Ornith-1.5-35B-A3B-NVFP4 <FLAG> --port 1924

--moe-cache-auto is not passed in any of these; it turns itself on when no MoE cache
sizing flag is given, which is how the default path reaches it.

Assisted-by: Claude Opus 5