feat(moe): partial-pin serving to beat the Windows/WSL2 pinned-memory quota (--pin-exempt-layers)

[ervisthegreattt]

Problem. Windows caps cudaHostRegister-able host memory at roughly 50% of physical RAM (WDDM-managed; WSL2 inherits it and NVIDIA marks it not actionable:
see the long-standing threads on the NVIDIA forums, e.g. "Change limit of 50% for cudaHostAlloc pinned memory on Windows 10/11"). FreeToken pins every expert bank, so on any Windows/WSL2 machine the servable-model ceiling is the pin quota, not
RAM. We measured the quota at 94 GiB on a 192 GB box and confirmed it is independent of the WSL VM size (identical at 100 GB and 160 GB .wslconfig). MiniMax-M2.5's ~126 GB of banks cannot load at any setting; the paper's fallback
for unpinnable pools is the all-CPU backend.

This PR adds the middle path the code already anticipates. HostResidency has PAGEABLE, ExpertBanks.layer_residency and set_bank_sources(..., layer_residency=...) are already plumbed, and offload_cache.py raises NotImplementedError for non-pinned layers. This change implements that stub:

• New --pin-exempt-layers (same grammar as --moe-cpu-layers: ids, count evenly strided, or fraction). Exempt layers skip pinning at load (the per-layer PinPipeline sink honors an exempt set) and are auto-unioned into cpu_layer_ids, so their decode uses the existing CPU executor unchanged.
• Prefill for exempt layers uses the legacy full-layer .copy_() into the overlap double buffer; from pageable memory the driver stages it (slower, and the enqueuing thread stalls for the staging, but correct under the existing ready/release event discipline). The hit-D2D split prefetch and the fused copy plan skip exempt layers; set_bank_sources validates exempt is a subset of cpu_layer_ids and that moe_prefill_overlap is on.
• v1 scope: the nvfp4 provider path (that is where the >90 GB models live). bf16/q4_0/ds_fp4 providers, setup_offload_expert_banks models, FTW checkpoints, and --use-dummy-weight raise NotImplementedError with clear messages rather than silently pinning everything. With the flag unset, behavior is byte-identical to main.

Numbers (RTX 5090 32 GB, 192 GB DDR5, Windows 11 + WSL2, MiniMax-M2.5-NVFP4, 229B, 62 MoE layers at 1.90 GiB of banks each): --pin-exempt-layers 15 pins 47 layers (89.4 GiB, under the 94 GiB quota) and the model serves: READY in
127 s (129 GB expert read at 2.5 GB/s, parallel loader), decode 9.1-9.6 tok/s across context lengths, 28.9k-token prompt TTFT 34.8 s cold / 7.1 s on radix repeat. Without this change the same command dies at cudaHostRegister ~117 s in.

Tests. tests/moe/test_pin_policy.py (grammar, 11 cases) and tests/moe/test_residency_propagation.py (the set_bank_sources contract, 9 cases) run CPU-only. The existing CUDA suites most affected (test_offload, test_fused_copy, test_prefill_hit_d2d, test_hybrid_fetch: 29 tests) pass unchanged on all-pinned defaults.

Not in this PR, open to direction: an auto fit-to-quota policy (needs a cached quota measurement; we prototyped a probe and would follow up), a pinned staging ring to restore fully-async prefill for exempt layers, and extending the plan to the other providers. Happy to adjust naming, split the change, or add coverage wherever you want it.

[gdevenyi]

Triaged while assembling a merged deployment branch for a 2× RTX 6000 Ada / 2× Xeon Gold 6526Y Linux box serving DeepSeek-V4-Flash with offloaded experts, to benchmark the open PRs together.

Partial-pin serving for Windows/WSL2. Not merged — Linux here, with memlock raised to unlimited, so the constrained-pinning path this targets never triggers. Worth noting from the Linux side that the same symptom has a different cause there: the default RLIMIT_MEMLOCK of RAM/8 (62 GiB of 503) fails a 155 GiB registration partway through a load. If partial-pin has a Linux-visible path it might serve that case too, but I have not tested it.

Flagging only so the absence of a report from me is not read as a problem found — I merged and benchmarked #30, #48, #56, #69, #70, #71 and #81, and left this one out deliberately.

[ascorb12]

Native-Windows data point supporting this, plus what looks like a help-text/code
mismatch in v0.1.2.

Setup

Windows 11 (native, not WSL2), RTX 4090 24 GB, 63 GB RAM, driver 610.88,
FreeToken v0.1.2 wheel (Desktop-app install), google/gemma-4-26B-A4B-it bf16,
--moe-backend hybrid. Pin quota here should be ~31 GiB (50% of RAM).

Observations

run	flags	outcome
1	--moe-cache-auto only	loads; auto puts ~10 GB experts in VRAM, host-pins the rest (~30 GiB, just under quota)
2	--num-tokens 24576 --moe-cache-auto	cudaHostRegister failed for 0.9 GiB in PinPipeline at load
3	same, clean reboot-equivalent state	same failure — reproducible, not fragmentation
4	same + --moe-cpu-layers 0.5	same failure — in v0.1.2 every bank still goes through PinPipeline.pin() regardless; the OS-locked alternative is the unimplemented PAGEABLE path
5	--num-tokens 16384 --moe-cache-auto	same failure, error text now cudaHostRegister failed: out of memory — so even near-run-1 pin demand no longer fits
6	--moe-backend cpu --num-tokens 24576	same failure — the all-CPU backend also routes every bank through PinPipeline.pin() in v0.1.2, so the "fallback for unpinnable pools is the all-CPU backend" described here is not shipped behavior either

Quota sanity check between runs: a single fresh-process torch.empty(..., pin_memory=True) of 20 GiB succeeds on this box. The quota is healthy;
the failures are demand-side — with one caveat worth flagging: after the first
force-killed run, nvidia-smi permanently reports ~16 GiB of VRAM in use with
no owning process (a zombie WDDM context). If that context also retains its
host-pinned registrations, the cumulative pin ceiling drops from ~31 GiB to
roughly the ~20 GiB we can still pin — which matches the observed failures
exactly. If so, crashed FreeToken runs poison subsequent loads until reboot,
which compounds this issue: every pin failure makes the next attempt likelier
to fail.

Post-reboot data (which narrows the diagnosis)

After a reboot cleared the phantom VRAM:

• gpt-oss-20b (small banks, well under quota) loaded with zero pin
  failures — small demand is fine on a clean machine.
• Gemma with --num-tokens 16384 still failed with the same
  cudaHostRegister error — so the zombie context was not the dominant
  factor. The arithmetic is: pin demand = total banks (~40 GiB bf16 for this
  model) minus whatever expert cache fits in VRAM. The auto config leaves
  ~10 GiB of experts in VRAM (pin ≈ 30 GiB, just under the ~31 GiB quota);
  any --num-tokens shrinks that VRAM cache and pushes pin demand over.

Net: on a 63 GB Windows box, this model's usable KV is capped around 8-12k
tokens by the pin quota, not by VRAM or RAM — and there is no shipped knob
that trades against it (--moe-cpu-layers and --moe-backend cpu both still
pin every bank in v0.1.2). That makes this PR's partial-pin path the only real
fix for agent-sized contexts on Windows. Working reference config meanwhile:
--moe-cache-auto alone (8,287-token pool, 34.8 tok/s decode on a 4090).
Whether crashed runs additionally degrade the pin ceiling (our single-alloc
probe topped out at 20 GiB pre-reboot) remains suspected but unconfirmed.

Interpretation: pinning the KV pool evicts expert cache from VRAM into host
banks, which raises pin demand past the Windows quota. So on native Windows,
any attempt to fix the small-KV problem from #111 via --num-tokens trips the
pin quota instead — the two issues chain into each other.

Help-text vs code

ft serve --help in v0.1.2 says for --moe-cpu-layers: "Unset = automatic on
Windows (locks just enough head+tail layers when the banks exceed the pin
budget, none otherwise)". With it unset, no automatic engaged — every bank went
through PinPipeline.pin() and the load died. That matches the v0.1.2 source:
host_banks.py declares HostResidency.PAGEABLE but notes "their movement
paths are not implemented here". So the automatic described in the help appears
to be this PR, not shipped behavior — worth either landing this or trimming the
help text so Windows users are not debugging a mitigation that does not exist
yet.

Happy to test the PR branch on this native-Windows box if useful — it seems to
be the exact configuration it targets.

[jason-fxz]

We've implemented this in #112 — --moe-cpu-layers now keeps the selected layers' banks in OS-locked memory instead of pinned.

Closing this in favor of that.

[ascorb12]

Correction to one claim in my comment above, in the interest of accuracy.

We have now identified what actually held memory across our failed runs: an
orphaned multiprocessing.spawn python child of a crashed engine (parent
dead, command line is just python -c "from multiprocessing.spawn import spawn_main; ..." with no identifiable engine string). It survived every
name/cmdline-filtered cleanup, kept ~14 GiB of GPU memory, and — critically —
kept its host banks pinned. Killing that one pid dropped reported VRAM from
17.1 to 2.8 GiB instantly. No reboot required, and no driver-level "zombie
context" involved.

Consequence for my earlier analysis: the "Gemma at --num-tokens 16384 still
fails after a reboot" datapoint was confounded — a fresh orphan from the
preceding (successful, then force-killed) gpt-oss run was alive during that
test, holding its own pinned banks. So the clean split between "pin demand
exceeds the quota arithmetically" and "prior runs' pinned memory was never
released" was not actually established by that run. The arithmetic remains
plausible; it just is not proven by our data.

Refined practical guidance for anyone hitting cudaHostRegister failures on
Windows:

1. Before blaming the quota, look for orphaned engine children: any large
   python.exe in the GPU perf counters whose parent is dead. Kill it by pid
   with its tree.
2. The failure loop we described is real but the mechanism is prosaic: crashed
   or force-killed engines leave spawn children that hold pinned memory, and
   they are easy to miss because nothing in their name says what they belong to.
   A PID-file or job-object cleanup in the supervisor would eliminate the whole
   class.