--moe-cache-auto leaves no room for prefill: boots and serves decode, then the first 8k prefill OOMs and kills the worker (NVFP4 dense) #401

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. Siblings in the same resolver, none of them this: --moe-cache-auto ignores --num-tokens / --num-pages: expert fill overruns the budget and the KV pool OOMs at boot #383 is --moe-cache-auto ignoring --num-tokens/--num-pages and OOMing at boot; fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor #340 is a dummy-page off-by-one in the boot budget; AssertionError: cache budget too small #150 is cache budget too small. This one boots cleanly and serves decode, then OOMs at prefill runtime. Backend worker is gone and cannot be restarted; stopping the API server #20 is the shared symptom (Backend worker is gone).
• I am on the latest nightly wheel (0.1.2+gaf71ba432, commit af71ba432).

What happened

--moe-cache-auto sizes the expert-slot fill against free VRAM at load, with no allowance for the transient allocations that prefill makes later. The plan fits at boot, decode is fine, and then the first large prefill OOMs and takes the worker down permanently.

What I did: serve nvidia/Qwen3.6-35B-A3B-NVFP4 (modelopt MIXED_PRECISION: fp8 per-tensor dense + NVFP4 experts + NVFP4 lm_head) with --moe-backend offload --moe-cache-auto --kv-reserve-tokens 131072 on a 16 GiB card, then send 8192-token prompts.

What I expected: the auto planner leaves enough headroom that a prefill at the default --max-extend-tokens 8192 runs, since that is the configured chunk size it should be planning for.

What happened instead:

• boot is clean. --moe-cache-auto resolved moe_cache_size=4484 num_pages=131141, KV 2.50 GiB, Free memory after initialization: 1.17 GiB, health ok.
• decode works — 16/16 requests at a 1200-token prompt, 135.87 tok/s.
• the first 8192-token prefill dies in the GDN chunked-prefill workspace (chunk_fwd_o → torch.zeros_like(v)), 64 MiB short with 96 MiB free, and the supervisor then stops the API server.

The plan is not wrong about the weights — it is wrong about what prefill still needs. moe_cache_size=4484 is 7.42 GiB of expert cache (4484 × 1.6934 MiB).

Control on the same box, same flags, same --max-extend-tokens 8192: an AWQ-int4 checkpoint of the same model family and geometry (Qwen3.6-35B-A3B, 40 layers / 256 experts / H=2048 / I=512) resolves moe_cache_size=3003 — a 5.63 GiB cache — also reports 1.17 GiB free after init, and completes 8/8 prefill requests at 8192 tokens. So the failure is not the card being too small; it is the auto planner choosing a cache size that leaves no room for prefill, and how much room it leaves varies with the checkpoint's dense format.

Workaround: replace --moe-cache-auto with an explicit --moe-cache-size 3600 (5.96 GiB, ~1.3 GiB more headroom). Then prefill is 8/8 at 8192 tokens and decode is 126.9 tok/s — i.e. the fix costs ~7% decode and buys a server that does not die.

Hypothesis, offered as such rather than asserted: the extra pressure on this checkpoint is likely the dense NVFP4 prefill path, which allocates a full bf16 scratch. From kernel/triton/nvfp4_linear.py:

M > 64 (prefill): dequantize the whole weight to a bf16 scratch (memory-roof Triton kernel, N-chunked to bound the allocation) and run cuBLAS.

The AWQ/Q4_1 control makes no such allocation, which would explain why the same flags and the same reported free memory behave differently. The failing allocation in the traceback is the GDN workspace, so I am not claiming the scratch is the allocation that fails — only that it is a plausible reason this checkpoint has less room left than the control at the same nominal free-memory figure. Someone with moe_collect_stats / a memory profile could confirm quickly.

Suggested direction: have plan_cache_budget reserve a prefill headroom term derived from max_extend_tokens (and, for quantized dense formats that use a dequant scratch, from the largest dense weight it will need to materialize) before handing the remainder to expert slots. A conservative floor would be preferable to the current behaviour, because the failure mode is an unrecoverable worker death rather than a degraded plan.

How did you install FreeToken

pip / uv wheel

FreeToken version

0.1.2+gaf71ba432

OS

Fedora

OS details

Fedora Linux 44 (Workstation Edition), kernel 7.2.3-cachyos2.fc44.x86_64. Engine runs in a container (nvidia/cuda:13.3.1-devel-ubuntu24.04, venv Python 3.12.3), torch 2.11.0+cu130, flashinfer 0.6.18.

GPU and driver

RTX 4080 SUPER 16 GiB (16376 MiB, sm_89), driver 610.57.04. Single GPU, desktop attached (Free memory before loading model: 14.59 GiB).

CPU and system RAM

AMD Ryzen 9 5950X 16-Core (32 threads), 31 GiB system RAM

Checkpoint

nvidia/Qwen3.6-35B-A3B-NVFP4

Command

ft serve \
  --model /models/qwen3.6/35b-a3b/nvidia-nvfp4 \
  --served-model-name qwen3.6-35b-a3b-nvfp4 \
  --host 0.0.0.0 --port 1919 \
  --moe-backend offload \
  --moe-cache-auto \
  --kv-reserve-tokens 131072 \
  --max-seq-len-override 131072 \
  --max-running-requests 1 \
  --max-output-tokens 4096 \
  --sampling-defaults model \
  --enable-cache-report

# then, against that server, 8 sequential requests with an 8192-token prompt and 1 output
# token (evalscope perf --dataset random --min-prompt-length 8192 --max-prompt-length 8192
# --min-tokens 1 --max-tokens 1 --parallel 1 --number 8). The first one kills the worker.

Full log

NVFP4 expert backend: triton
Free memory before loading model: 14.59 GiB
expert banks: low free RAM -> serial build (avoids parallel-reader OOM; override with --expert-load parallel)
expert banks: slow path (serial build)
Loading Qwen3.5 NVFP4 experts: 100%|##########| 3/3 [00:19<00:00,  6.60s/it]
--moe-cache-auto resolved moe_cache_size=4484 num_pages=131141 (prefill_overlap=True)
Allocating 131141 tokens for KV cache, K + V = 2.50 GiB
Free memory after initialization: 1.17 GiB

... decode serves normally here (16/16 requests, 1200-token prompts) ...

... then, on the first 8192-token prefill:

    o, h = ChunkGatedDeltaRuleFunction.apply(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/ft/lib/python3.12/site-packages/torch/autograd/function.py", line 596, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/ft/lib/python3.12/site-packages/freetoken/kernel/fla/utils.py", line 178, in wrapper
    return fn(*contiguous_args, **contiguous_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/ft/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 481, in decorate_fwd
    return fwd(*args, **kwargs)  # pyrefly: ignore [not-callable]
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/ft/lib/python3.12/site-packages/freetoken/kernel/fla/chunk.py", line 112, in forward
    g, o, A, w, h, v_new = chunk_gated_delta_rule_fwd(
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/ft/lib/python3.12/site-packages/freetoken/kernel/fla/chunk.py", line 67, in chunk_gated_delta_rule_fwd
    o = chunk_fwd_o(
        ^^^^^^^^^^^^
  File "/opt/ft/lib/python3.12/site-packages/freetoken/kernel/fla/chunk_o.py", line 146, in chunk_fwd_o
    o = torch.zeros_like(v)
        ^^^^^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 64.00 MiB. GPU 0 has a total
capacity of 15.57 GiB of which 96.06 MiB is free. Including non-PyTorch memory, this process
has 14.59 GiB memory in use. Of the allocated memory 14.23 GiB is allocated by PyTorch, with
2.00 MiB allocated in private pools (e.g., CUDA Graphs), and 32.14 MiB is reserved by PyTorch
but unallocated.

Backend supervisor: backend worker freetoken-TP0-scheduler exited
Backend worker is gone and cannot be restarted; stopping the API server
INFO:     Shutting down

Anything else

The engine carries four local patches (open PRs #211, #287, #266, plus a local int4 expert repack for AWQ / compressed-tensors checkpoints). None of them touch the cache budget resolver, the GDN kernels, the NVFP4 dense path, or plan_cache_budget — and the AWQ control above runs on the identical image, so the difference between the two is the checkpoint's dense format, not the patches. expandable_segments is enabled by the engine (torch.cuda.memory._set_allocator_settings("expandable_segments:True") appears in the startup log).

Happy to re-run with any additional instrumentation, or to test a headroom-reserving patch against both checkpoints — I have the AWQ and NVFP4 checkpoints side by side on the same box, which makes the A/B cheap.