Merge pull request #112 from FlashML-org/feat/split-residency

feat(moe): per-layer host-bank residency

python/freetoken/checkpoint/ftw.py

@@ -33,7 +33,6 @@ Layout on disk::
 from __future__ import annotations
 
 import json
-import logging
 import math
 import mmap
 import os
@@ -43,7 +42,9 @@ from concurrent.futures import ThreadPoolExecutor
 
 import torch
 
-logger = logging.getLogger(__name__)
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
 
 INDEX_NAME = "freetoken_weight.json"
 FORMAT_TAG = "freetoken_weight"
@@ -402,12 +403,16 @@ def iter_ftw_weights(path: str, *, kinds=("weight",), workers: int = 8,
 
 
 def load_ftw_banks(
-    path: str, *, num_layers: int, workers: int = 8, chunk: int = _DEFAULT_CHUNK
+    path: str, *, num_layers: int, workers: int = 8, chunk: int = _DEFAULT_CHUNK,
+    layer_residency: list[str] | None = None,
 ):
     """Reconstruct the offload :class:`ExpertBanks` from the FTW's ``experts_bank``
-    entries, on the per-layer host bank contract (one pinned ``[num_experts, ...]``
+    entries, on the per-layer host bank contract (one ``[num_experts, ...]``
     HostBank per layer per bank; see ``moe.offload_cache.set_bank_sources``).
 
+    ``layer_residency`` (default: all pinned) settles each layer's banks per its ``HostResidency`` label as reads complete: PINNED -> cudaHostRegister, LOCKED -> mlock (CPU-executor resident, no pin quota spent).
+    The applied labels are echoed back on ``ExpertBanks.layer_residency``.
+
     Two on-disk row layouts, distinguished per bank name (a file never mixes them for
     the same name -- checked below):
 
@@ -430,9 +435,22 @@ def load_ftw_banks(
     vectors, unaffected by the row split (fixed GPU residency; see
     ``cache_budget.expert_bytes_per_slot``).
     """
-    from freetoken.moe.host_banks import HostBank, PinPipeline, alloc_banks
+    from freetoken.moe.host_banks import (
+        HostBank, HostResidency, PinPipeline, alloc_banks, born_pinned_default,
+    )
     from freetoken.utils.progress import byte_bar
 
+    residency = layer_residency or [HostResidency.PINNED.value] * num_layers
+    assert len(residency) == num_layers, (len(residency), num_layers)
+
+    # PINNED layers are born-pinned (cudaHostAlloc) where that wins (see born_pinned_default); LOCKED/PAGEABLE layers stay lazy mmaps
+    born = born_pinned_default()
+
+    def _backing(layer_id: int) -> str:
+        if born and residency[layer_id] == HostResidency.PINNED.value:
+            return "cuda"
+        return "mmap"
+
     reader = FTWReader(path)
     bank_entries = reader.entries("experts_bank")
     if not bank_entries:
@@ -492,10 +510,10 @@ def load_ftw_banks(
             win_off = (off // ALIGN) * ALIGN
             win_end = _align_up(off + layer_bytes)
             head_pad = off - win_off
-            bank = HostBank((win_end - win_off,), torch.uint8)
+            bank = HostBank((win_end - win_off,), torch.uint8, backing=_backing(layer_id))
             row_hb[name].append(bank)
             row_view_args[name].append((head_pad, layer_bytes, num_experts, tuple(row_shape), dtype))
-            row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes))
+            row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id))
 
     for base, by_layer in per_layer_groups.items():
         assert sorted(by_layer) == list(range(num_layers)), (
@@ -507,10 +525,10 @@ def load_ftw_banks(
         for layer_id in range(num_layers):
             e = by_layer[layer_id]
             assert e["global_off"] % ALIGN == 0, (base, layer_id, e["global_off"])  # writer invariant
-            bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]))
+            bank = HostBank(tuple(e["shape"]), _dtype_of(e["dtype"]), backing=_backing(layer_id))
             row_hb[base].append(bank)
             row_view_args[base].append(None)
-            layer_jobs.append((base, bank, e))
+            layer_jobs.append((base, bank, e, layer_id))
 
     total_bytes = sum(e["nbytes"] for e in bank_entries)
     bar = byte_bar(total_bytes, "Loading expert banks (FTW)")
@@ -528,16 +546,16 @@ def load_ftw_banks(
                 bar.update(e["nbytes"])
 
             def _read_row(job):
-                _name, bank, win_off, win_len, layer_bytes = job
+                _name, bank, win_off, win_len, layer_bytes, layer_id = job
                 reader.read_into(bank.memoryview(), {"global_off": win_off, "nbytes": win_len},
                                  workers=workers, chunk=chunk)
-                pins.submit(bank)
+                pins.submit(bank, residency[layer_id])
                 bar.update(layer_bytes)
 
             def _read_layer(job):
-                _name, bank, entry = job
+                _name, bank, entry, layer_id = job
                 reader.read_into(bank.memoryview(), entry, workers=workers, chunk=chunk)
-                pins.submit(bank)
+                pins.submit(bank, residency[layer_id])
                 bar.update(entry["nbytes"])
 
             with ThreadPoolExecutor(min(max(_BANK_CONCURRENCY, 16), max(n_jobs, 1))) as ex:
@@ -564,10 +582,45 @@ def load_ftw_banks(
 
     from freetoken.moe.expert_banks import ExpertBanks
 
+    # a failed mlock leaves a LOCKED layer pageable; the log and labels report what the banks actually settled at
+    applied = list(residency)
+    for banks in row_hb.values():
+        for layer_id, bank in enumerate(banks):
+            if (applied[layer_id] == HostResidency.LOCKED.value
+                    and bank.residency is not HostResidency.LOCKED):
+                applied[layer_id] = HostResidency.PAGEABLE.value
+    unpinned = [i for i, r in enumerate(applied) if r != HostResidency.PINNED.value]
+    if unpinned:
+        by_layer = [0] * num_layers
+        for name, banks in row_hb.items():
+            for layer_id, bank in enumerate(banks):
+                by_layer[layer_id] += bank.nbytes
+        locked = [i for i in unpinned if applied[i] == HostResidency.LOCKED.value]
+        pageable = [i for i in unpinned if i not in set(locked)]
+        pinned_b = sum(b for i, b in enumerate(by_layer) if i not in set(unpinned))
+        locked_b = sum(by_layer[i] for i in locked)
+        pageable_part = ""
+        if pageable:
+            pageable_b = sum(by_layer[i] for i in pageable)
+            pageable_part = (
+                f" + {pageable_b / 2**30:.2f} GiB pageable "
+                f"(lock failed, {len(pageable)} CPU layers: {pageable})"
+            )
+        logger.info(
+            f"MoE bank split residency: {pinned_b / 2**30:.2f} GiB pinned "
+            f"({'born-pinned cudaHostAlloc' if born else 'cudaHostRegister'}, "
+            f"{num_layers - len(unpinned)} GPU layers) + "
+            f"{locked_b / 2**30:.2f} GiB OS-locked ({len(locked)} CPU layers: {locked})"
+            f"{pageable_part}"
+        )
+
     # alphas are the small per-expert scale vectors, distinguished by their reserved names
     # (not a separate kind); everything else under experts_bank is a weight source.
     alpha_kw = {n: alpha_hb[n].tensor for n in alpha_hb}
-    return ExpertBanks(reader.meta("quant_format"), sources, **alpha_kw)
+    return ExpertBanks(
+        reader.meta("quant_format"), sources, **alpha_kw,
+        layer_residency=applied,
+    )
 
 
 __all__ = [

python/freetoken/engine/engine.py

@@ -512,12 +512,48 @@ class Engine:
         # layout; the GPU slot-cache GEMM reads those same native rows. decode_target also
         # gates the CPU executor build below.
         cpu_layer_ids = _resolve_cpu_layers(config, config.model_config.num_moe_layers)
+        if (
+            not cpu_layer_ids
+            and config.moe_cpu_layers is None
+            and config.moe_backend in ("offload", "hybrid")
+            and _pin_budget_bytes() is not None
+        ):
+            cpu_layer_ids = _auto_cpu_layers(config, config.model_config.num_moe_layers)
         if config.moe_backend == "hybrid":
             decode_target = "hybrid"
         elif cpu_layer_ids:
             decode_target = "cpu"
         else:
             decode_target = "gpu"
+        # split residency: where pinning is quota-capped (_pin_budget_bytes), pin only the GPU layers' banks and mlock the CPU layers'
+        # uncapped hosts keep every bank pinned (CPU decode reads them the same; overlap prefill stays on)
+        # not applied to plain --moe-backend cpu; all-locked under a cap = --moe-backend offload --moe-cpu-layers 1.0
+        split_residency = (
+            bool(cpu_layer_ids)
+            and config.moe_backend in ("offload", "hybrid")
+            and _pin_budget_bytes() is not None
+        )
+        if config.moe_backend == "cpu" and not split_residency:
+            # cpu mode pins every bank for the prefill double buffer; over the pin cap that dies in cudaHostRegister, so lock everything instead
+            from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes
+
+            budget = _pin_budget_bytes()
+            bank_bytes = None
+            if budget is not None:
+                bank_bytes = ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config)
+            if bank_bytes and bank_bytes > budget:
+                split_residency = True
+                logger.info_rank0(
+                    f"--moe-backend cpu: banks {bank_bytes / 2**30:.2f} GiB exceed the "
+                    f"pin budget; OS-locking all layers instead of pinning"
+                )
+        if split_residency and config.moe_prefill_overlap:
+            # locked (unregistered) layers cannot feed the async pinned H2D double buffer; their prefill is a synchronous pageable copy via materialize
+            logger.info_rank0(
+                "--moe-cpu-layers split residency: disabling MoE prefill overlap "
+                "(locked layers prefill via synchronous pageable copies)"
+            )
+            object.__setattr__(config, "moe_prefill_overlap", False)
         if cache_factory is None:
             # Fast path: an FTW checkpoint loads its repacked banks directly.
             # Slow path: load_expert_banks auto-picks parallel vs serial baseline by
@@ -525,6 +561,15 @@ class Engine:
             # --expert-load: serial/parallel force the read; auto (None) lets load_expert_banks
             # pick (parallel for scattered experts, with a low-RAM fallback to serial).
             expert_parallel = {"serial": False, "parallel": True}.get(config.expert_load, None)
+            requested_residency = None
+            if split_residency:
+                from freetoken.moe.host_banks import HostResidency
+
+                requested_residency = [
+                    HostResidency.LOCKED.value if i in cpu_layer_ids
+                    else HostResidency.PINNED.value
+                    for i in range(config.model_config.num_moe_layers)
+                ]
             banks = load_expert_banks(
                 config.model_path,
                 config.model_config,
@@ -533,6 +578,7 @@ class Engine:
                 dummy=config.use_dummy_weight,
                 parallel=expert_parallel,
                 decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"),
+                layer_residency=requested_residency,
             )
             if config.moe_cache_auto:
                 size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks)
@@ -566,15 +612,17 @@ class Engine:
                 decode_target=decode_target,
                 hybrid_max_fetch=config.moe_hybrid_max_fetch,
             )
+            # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set
+            cache.cpu_layer_ids = cpu_layer_ids
             cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency)
             cache.set_alphas(banks.gate_up_alpha, banks.down_alpha)
         else:
             cache = cache_factory(config, self.device)
             cache.decode_target = decode_target
             cache.hybrid_max_fetch = config.moe_hybrid_max_fetch
+            cache.cpu_layer_ids = cpu_layer_ids
         if decode_target == "hybrid":
             self._resolve_hybrid_fetch(config, cache)
-        cache.cpu_layer_ids = cpu_layer_ids
         # Must be set before CUDA graph capture so the (device-side) accumulation ops are
         # captured and re-run on every decode replay.
         cache.collect_stats = config.moe_collect_stats
@@ -1071,6 +1119,74 @@ def _resolve_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[
     return _parse_cpu_layers_spec(spec, num_moe_layers)
 
 
+# expert activations the CPU MoE executor supports (csrc ActKind)
+_CPU_MOE_ACTS = (
+    "silu", "swish", "gelu", "gelu_tanh", "gelu_pytorch_tanh", "swigluoai",
+)
+
+
+def _cpu_moe_executor_viable(model_config) -> bool:
+    """Whether an automatic CPU-decode decision may target the CPU MoE executor.
+
+    A default boot must degrade to GPU offload instead of crashing in CpuMoeExecutor after the whole load; explicit cpu/hybrid/--moe-cpu-layers picks still fail loudly."""
+    from freetoken.moe.cpu_executor import _WFMT_IDS, compiled_extension_supports
+
+    try:
+        from freetoken.kernel import _cpu_moe  # noqa: F401
+    except ImportError:
+        return False
+    act = getattr(model_config, "hidden_act", "silu")
+    moe_wfmt = getattr(model_config, "moe_weight_format", None)
+    if act not in _CPU_MOE_ACTS and moe_wfmt != "mxfp4":
+        return False
+    if moe_wfmt != "mxfp4" and not compiled_extension_supports(act):
+        return False
+    expert_quant = getattr(model_config, "expert_quant", "none")
+    fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16")
+    return fmt == "mxfp4" or fmt in _WFMT_IDS
+
+
+def _pin_budget_bytes() -> int | None:
+    """Bytes this process can safely cudaHostRegister, or None when the platform does not cap pinning (plain Linux).
+
+    WSL's WDDM-backed CUDA caps pinning near half of RAM, shared across processes -- budget 40%. FREETOKEN_PIN_BUDGET_GB overrides anywhere."""
+    if env := os.environ.get("FREETOKEN_PIN_BUDGET_GB"):
+        return int(float(env) * 2**30)
+    if not hasattr(os, "uname") or "microsoft" not in os.uname().release.lower():  # WSL kernel tag
+        return None
+    return int(os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") * 0.4)
+
+
+def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[int]:
+    """Pick CPU (locked) MoE layers automatically when the banks exceed the pin budget.
+
+    Locks just enough head+tail layers: per-layer decode miss rates are U-shaped, so the ends are the cheapest to move off the slot cache."""
+    from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes
+
+    bank_bytes = ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config)
+    if not bank_bytes:
+        return frozenset()
+    budget = _pin_budget_bytes()
+    if budget is None or bank_bytes <= budget:
+        return frozenset()
+    if not _cpu_moe_executor_viable(config.model_config):
+        logger.info_rank0(
+            f"--moe-cpu-layers auto: banks {bank_bytes / 2**30:.2f} GiB exceed the "
+            f"pin budget {budget / 2**30:.2f} GiB, but the CPU MoE executor cannot "
+            f"serve this model; keeping every layer pinned on the GPU offload path"
+        )
+        return frozenset()
+    n = min(num_moe_layers, math.ceil(num_moe_layers * (1 - budget / bank_bytes)))
+    head = (n + 1) // 2
+    ids = frozenset(range(head)) | frozenset(range(num_moe_layers - (n - head), num_moe_layers))
+    logger.info_rank0(
+        f"--moe-cpu-layers auto: banks {bank_bytes / 2**30:.2f} GiB > pin budget "
+        f"{budget / 2**30:.2f} GiB; locking {n} head+tail MoE layers for CPU decode "
+        f"({sorted(ids)})"
+    )
+    return ids
+
+
 # MoE-only knobs and the value each resolves to on a dense model. moe_backend is handled
 # separately (its dense value is 'fused', but 'auto' resolves there without a warning).
 _DENSE_MOE_SETTINGS = {
@@ -1200,13 +1316,10 @@ def _adjust_config(config: EngineConfig):
     # swigluoai the generic GEMV epilogue). A model with any other expert
     # activation cannot decode on the CPU: reject an explicit cpu/hybrid pick at
     # config time, and keep auto from upgrading offload -> hybrid off the profile.
-    _cpu_moe_acts = (
-        "silu", "swish", "gelu", "gelu_tanh", "gelu_pytorch_tanh", "swigluoai",
-    )
     # hidden_act (the dense activation) stands proxy for the expert activation --
     # true for every in-tree model. mxfp4 experts pass regardless: their act runs
     # inside the mxfp4 kernel, not the generic epilogue.
-    _cpu_moe_act_ok = getattr(model_config, "hidden_act", "silu") in _cpu_moe_acts or (
+    _cpu_moe_act_ok = getattr(model_config, "hidden_act", "silu") in _CPU_MOE_ACTS or (
         getattr(model_config, "moe_weight_format", None) == "mxfp4"
     )
     if (
@@ -1339,12 +1452,10 @@ def _adjust_config(config: EngineConfig):
             f"got {config.moe_backend!r}"
         )
 
-    if is_moe and config.moe_cpu_layers and config.moe_backend != "offload":
-        # The hybrid split pins a subset of *offload* layers to CPU decode; it needs the
-        # offload host banks + slot cache. 'cpu' already runs every layer on CPU; 'fused'
-        # keeps experts resident on the GPU (no host banks for the CPU executor to read).
+    if is_moe and config.moe_cpu_layers and config.moe_backend not in ("offload", "hybrid"):
+        # the layer split needs the offload host banks + slot cache; 'cpu' already runs every layer on CPU, 'fused' keeps experts resident on the GPU (no host banks)
         raise ValueError(
-            "--moe-cpu-layers requires --moe-backend offload (got "
+            "--moe-cpu-layers requires --moe-backend offload or hybrid (got "
             f"{config.moe_backend!r}); use --moe-backend cpu to run all layers on CPU"
         )
 

python/freetoken/models/deepseek_v4/moe.py

@@ -100,10 +100,14 @@ class DSV4OffloadMoELayer(OffloadMoELayer):
         # keeps short-prompt slot residency -- hence hybrid decode's GPU/CPU
         # route split -- unchanged). Mixing modes across chunks is safe: the
         # streaming buffers disown their borrowed slots on invalidation.
-        if hidden_states.shape[0] * self.top_k >= self.num_experts:
-            return super()._prefill_routed(hidden_states, topk_weights, topk_ids)
         cache = self.offload_cache
         assert cache is not None
+        # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights)
+        if (
+            hidden_states.shape[0] * self.top_k >= self.num_experts
+            or cache.is_unpinned_layer(self.layer_id)
+        ):
+            return super()._prefill_routed(hidden_states, topk_weights, topk_ids)
         cache.ensure_experts(self.layer_id, topk_ids)  # in-place expert-id -> slot
         cache.copy_missing()
         if cache.collect_stats:

python/freetoken/moe/expert_banks.py

@@ -24,10 +24,13 @@ import torch
 
 from freetoken.utils import init_logger
 
-from .offload_cache import _BANK_SCHEMAS
+from .offload_cache import _BANK_BYTES_PER_EXPERT, _BANK_SCHEMAS
 
 logger = init_logger(__name__)
 
+# the parallel expert-bank reader needs POSIX O_DIRECT + preadv; without them the serial (safetensors/mmap) build is the only option
+_PARALLEL_READER_SUPPORTED = hasattr(os, "O_DIRECT") and hasattr(os, "preadv")
+
 
 @dataclass(frozen=True)
 class ExpertBanks:
@@ -40,8 +43,7 @@ class ExpertBanks:
     # marlin/b12x per-expert global scales ([L*E]); None for formats without them
     gate_up_alpha: torch.Tensor | None = field(default=None)
     down_alpha: torch.Tensor | None = field(default=None)
-    # Per-layer HostResidency values; None -> all pinned (the only class
-    # served; policies that assign other classes are not implemented).
+    # per-layer HostResidency values actually applied by the loader; None -> all pinned (also the degrade signal when a request was not honored)
     layer_residency: list[str] | None = field(default=None)
     # True iff the ``layer_sink`` passed to the loader was actually engaged (each layer
     # streamed straight to its sink instead of staying materialized here) -- set by
@@ -371,6 +373,38 @@ def _host_ram_fits_parallel(model_path: str) -> bool:
     return avail > sum(sizes) + max(sizes)
 
 
+def ftw_bank_bytes(model_path: str) -> int | None:
+    """Total expert-bank bytes of an FTW checkpoint, from its metadata (no bank IO).
+    ``None`` when the checkpoint is not FTW -- callers that size things pre-load (auto split residency) then leave the load unchanged."""
+    import json
+
+    meta = os.path.join(model_path, "freetoken_weight.json")
+    if not os.path.isfile(meta):
+        return None
+    with open(meta, encoding="utf-8") as f:
+        tensors = json.load(f).get("tensors", [])
+    return sum(t["nbytes"] for t in tensors if t.get("kind") == "experts_bank")
+
+
+def bank_bytes_estimate(model_config) -> int | None:
+    """Estimated total expert-bank bytes of a raw checkpoint, from the model config alone.
+
+    Sizes the pin-budget decisions where FTW metadata is not available; ``None`` for unknown formats or missing dims (callers then skip the pre-load sizing).
+    nvfp4 uses the native-row formula, a slight over-estimate for the repacked backends."""
+    expert_quant = getattr(model_config, "expert_quant", "none")
+    fmt = expert_quant if expert_quant != "none" else (
+        getattr(model_config, "moe_weight_format", None) or "bf16"
+    )
+    per_expert = _BANK_BYTES_PER_EXPERT.get(fmt)
+    layers = getattr(model_config, "num_moe_layers", None)
+    experts = getattr(model_config, "num_experts", None)
+    hidden = getattr(model_config, "hidden_size", None)
+    inter = getattr(model_config, "moe_intermediate_size", None)
+    if per_expert is None or not all((layers, experts, hidden, inter)):
+        return None
+    return layers * experts * per_expert(hidden, inter)
+
+
 def load_expert_banks(
     model_path: str,
     model_config,
@@ -383,6 +417,7 @@ def load_expert_banks(
     chunk: int = _PARALLEL_CHUNK,
     decode_target: str = "gpu",
     layer_sink=None,
+    layer_residency: list[str] | None = None,
 ) -> ExpertBanks:
     """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning
     the same normalized ``ExpertBanks`` and both pinning after fill:
@@ -400,22 +435,33 @@ def load_expert_banks(
     ``layer_sink`` (the converter only): forwarded to whichever provider is picked; a
     provider only engages it (and reports ``ExpertBanks.streamed=True``) for its own
     streamable formats, so callers must check ``streamed`` rather than assume it fired.
+
+    ``layer_residency``: per-layer ``HostResidency`` labels applied at settle time -- explicitly on the FTW fast path, ambiently (``requested_residency``) in the slow-path providers.
+    Applied labels are echoed on ``ExpertBanks.layer_residency``; a loader that settles some other way leaves it ``None`` (CPU-layer decode still works on pinned banks, it just saves no pin quota).
     """
     from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks
 
     if model_path and is_ftw_checkpoint(model_path) and not dummy:
         banks = load_ftw_banks(
-            model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk
+            model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk,
+            layer_residency=layer_residency,
         )
         if banks is not None:
             logger.info_rank0(f"expert banks: FTW fast path (FTW checkpoint {model_path})")
             return banks
 
+    if parallel and not _PARALLEL_READER_SUPPORTED:
+        logger.warning_rank0(
+            "expert banks: parallel O_DIRECT reader unsupported on this platform "
+            "(no os.O_DIRECT/preadv) -> serial build"
+        )
+        parallel = False
+
     auto = parallel is None
     if auto:
         from freetoken.models.weight import experts_scattered
 
-        parallel = not dummy and experts_scattered(model_path)
+        parallel = _PARALLEL_READER_SUPPORTED and not dummy and experts_scattered(model_path)
         # Low-RAM fallback: the parallel reader holds whole-shard ANONYMOUS buffers
         # (non-reclaimable) on top of the ~bank-sized resident set, so on a memory-tight box
         # it OOMs where the serial path (reclaimable file mmap) survives. Drop to serial when
@@ -432,12 +478,43 @@ def load_expert_banks(
     # OSError on those (which would leak the banks it pre-allocated, since host banks live for
     # the process). Only NotImplementedError (quant has no parallel reader; raised before any
     # allocation) falls back to serial.
-    try:
-        return _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk,
-                                   decode_target, layer_sink)
-    except NotImplementedError as exc:
-        if not parallel:
-            raise
-        logger.warning_rank0(f"parallel reader unavailable ({exc}); falling back to serial build")
-        return _build_expert_banks(model_path, model_config, device, dtype, dummy, False, workers, chunk,
-                                   decode_target, layer_sink)
+    from freetoken.moe.host_banks import requested_residency
+
+    with requested_residency(layer_residency) as residency_plan:
+        try:
+            banks = _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk,
+                                        decode_target, layer_sink)
+        except NotImplementedError as exc:
+            if not parallel:
+                raise
+            logger.warning_rank0(f"parallel reader unavailable ({exc}); falling back to serial build")
+            banks = _build_expert_banks(model_path, model_config, device, dtype, dummy, False, workers, chunk,
+                                        decode_target, layer_sink)
+    return _echo_residency(banks, layer_residency, residency_plan)
+
+
+def _echo_residency(banks: ExpertBanks, requested, plan) -> ExpertBanks:
+    """Stamp an honored residency request onto the ExpertBanks; keep None (and warn) when no settle point consulted the plan."""
+    if requested is None or banks.layer_residency is not None:
+        return banks
+    if plan is not None and plan.applied:
+        import dataclasses
+
+        labels = [plan.actual.get(i, r) for i, r in enumerate(requested)]
+        downgraded = [i for i, r in enumerate(requested) if labels[i] != r]
+        if downgraded:
+            logger.warning_rank0(
+                f"--moe-cpu-layers: layers {downgraded} settled pageable instead of "
+                f"OS-locked (lock failed); they still decode on the CPU executor but "
+                f"may swap under memory pressure"
+            )
+        return dataclasses.replace(banks, layer_residency=labels)
+    from freetoken.moe.host_banks import HostResidency
+
+    if any(r != HostResidency.PINNED.value for r in requested):
+        logger.warning_rank0(
+            "--moe-cpu-layers: this checkpoint's bank loader settles banks without "
+            "per-layer residency (pre-pins everything); CPU-layer decode still works "
+            "but saves no pinned quota"
+        )
+    return banks

python/freetoken/moe/host_banks.py

@@ -17,6 +17,7 @@ The mmaps are held for the process lifetime (the banks live as long as the offlo
 
 from __future__ import annotations
 
+import contextlib
 import ctypes
 import math
 import mmap
@@ -28,18 +29,18 @@ from enum import Enum
 
 import torch
 
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
+
 _BLK = 4096  # O_DIRECT alignment (page size)
 
 
 class HostResidency(str, Enum):
     """Residency class of a host bank layer.
 
-    ``PINNED`` (cudaHostRegister'd) is required for anything the GPU dereferences:
-    the decode gather kernels and the DMA prefill copies. ``LOCKED`` (VirtualLock /
-    mlock: resident for the CPU executor, but with no device address) and
-    ``PAGEABLE`` are reserved for platforms where the pin quota cannot cover every
-    layer (Windows/WDDM); their movement paths are not implemented here -- layers
-    with either class must be served by the CPU executor.
+    Only PINNED (cudaHostRegister'd) memory can feed the GPU movement paths; LOCKED (mlock'd, no device address) and PAGEABLE layers must decode on the CPU executor.
+    The non-pinned classes exist for hosts that cap CUDA pin quota (WSL/WDDM: ~half of RAM).
     """
 
     PINNED = "pinned"
@@ -52,37 +53,84 @@ _DEFAULT_CHUNK = 8 << 20
 # Hold the mmaps for the process lifetime; the offload cache reads from these banks forever.
 _LIVE_BUFFERS: list[mmap.mmap] = []
 
+def _env_born_pinned() -> bool | None:
+    """``FREETOKEN_BANK_CUDA_ALLOC`` tri-state: unset -> ``None`` (default applies), else the parsed boolean."""
+    v = os.environ.get("FREETOKEN_BANK_CUDA_ALLOC", "").strip().lower()
+    if not v:
+        return None
+    return v in ("1", "true", "yes", "on")
+
+
+def born_pinned_default() -> bool:
+    """Whether PINNED serving banks use cudaHostAlloc instead of mmap + register-after-fill.
+
+    Off by default: registered mmaps already read at the PCIe roofline and lazy mmaps commit pages only on fill. ``FREETOKEN_BANK_CUDA_ALLOC`` overrides."""
+    env = _env_born_pinned()
+    if env is not None:
+        return env
+    return False
+
 
 class HostBank:
-    """A lazily-allocated, page-aligned host buffer + its torch view, registered on demand.
+    """A page-aligned host buffer + its torch view, page-locked on demand: allocate -> fill -> ``pin()``/``lock()``.
 
-    Allocate -> fill (write real data, or O_DIRECT read into it) -> ``pin()``. The mmap is
-    rounded up to the O_DIRECT block so chunked reads are always aligned; ``tensor`` views
-    exactly ``nbytes``."""
+    * ``"mmap"`` (default) -- lazy anonymous mmap; pages materialize on fill, then ``pin()`` registers or ``lock()`` OS-locks it.
+    * ``"cuda"`` -- cudaHostAlloc, born pinned+mapped; ``pin()``/``lock()``/``release()`` are no-ops and it never takes LOCKED. See :func:`born_pinned_default`.
 
-    __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned")
+    The buffer is rounded up to the O_DIRECT block; ``tensor`` views exactly ``nbytes``. ``backing=None`` follows ``FREETOKEN_BANK_CUDA_ALLOC``."""
 
-    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
+    __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned", "_locked")
+
+    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype,
+                 *, backing: str | None = None):
+        if backing is None:
+            plan = _requested_residency
+            # a plan with non-pinned labels vetoes born-pinned: cudaHostAlloc spends the pin quota the plan exists to save
+            born = _env_born_pinned() and (plan is None or not plan.has_unpinned)
+            backing = "cuda" if born else "mmap"
+        assert backing in ("mmap", "cuda"), backing
         elsize = torch.empty((), dtype=dtype).element_size()
         self.nbytes = math.prod(shape) * elsize
         asize = ((self.nbytes + _BLK - 1) // _BLK) * _BLK
-        self._buf = mmap.mmap(-1, asize)  # lazy: address space only, no resident pages yet
-        _LIVE_BUFFERS.append(self._buf)
-        self.addr = ctypes.addressof(ctypes.c_char.from_buffer(self._buf))
+        if backing == "cuda":
+            from freetoken.kernel.pinned import alloc_pinned_tensor
+
+            # direct-IO readers need page alignment, but cudaHostAlloc only guarantees ~512 in practice
+            # over-allocate one block and carve the aligned window; the numpy slice keeps the pinned storage alive via .base
+            raw = alloc_pinned_tensor(asize + _BLK, dtype=torch.uint8)  # cudaMallocHost
+            raw.zero_()  # keep the anonymous-mmap guarantee: unwritten regions stay zero
+            off = (-raw.data_ptr()) % _BLK
+            self._buf = raw.numpy()[off:off + asize]
+            self.addr = raw.data_ptr() + off
+            assert self.addr % _BLK == 0
+            self._pinned = True  # born pinned+mapped; pin() is a no-op
+        else:
+            self._buf = mmap.mmap(-1, asize)  # lazy: address space only, no resident pages yet
+            _LIVE_BUFFERS.append(self._buf)
+            self.addr = ctypes.addressof(ctypes.c_char.from_buffer(self._buf))
+            self._pinned = False
         self.tensor = torch.frombuffer(self._buf, dtype=dtype, count=self.nbytes // elsize).view(*shape)
-        self._pinned = False
+        self._locked = False
 
     @property
     def residency(self) -> HostResidency:
-        return HostResidency.PINNED if self._pinned else HostResidency.PAGEABLE
+        if self._pinned:
+            return HostResidency.PINNED
+        if self._locked:
+            return HostResidency.LOCKED
+        return HostResidency.PAGEABLE
 
     def memoryview(self) -> memoryview:
         return memoryview(self._buf)
 
     def pin(self) -> None:
-        """cudaHostRegister the (now-filled, resident) buffer -- pin-after-fill."""
+        """cudaHostRegister the (now-filled) buffer -- pin-after-fill.
+
+        ``FREETOKEN_SKIP_BANK_PIN=1`` makes this a no-op for CPU-only tooling (the FTW converter); never set it when serving, the GPU paths need registered banks."""
         if self._pinned:
             return
+        if os.environ.get("FREETOKEN_SKIP_BANK_PIN", "").strip().lower() in ("1", "true", "yes", "on"):
+            return
         from freetoken.kernel.pinned import host_register
 
         try:
@@ -94,20 +142,59 @@ class HostBank:
         self._pinned = True
 
     def release(self) -> None:
-        """Drop the buffer's resident pages (address space stays valid; contents
-        become undefined). Only for buffers that are done being read -- the
-        converter releases each layer after writing it out."""
-        assert not self._pinned, "cannot release a pinned bank"
+        """Drop the resident pages; the address space stays valid, the contents become undefined.
+
+        For buffers that are done being read (the converter). No-op for born-pinned banks: registered pages cannot be dropped."""
+        if self._pinned:
+            return
         self._buf.madvise(mmap.MADV_DONTNEED)
 
     def lock(self) -> None:
-        """Make the buffer resident without a device mapping (VirtualLock/mlock).
+        """mlock the (now-filled) buffer: resident without CUDA pin quota, but no device address -- only the CPU executor can serve a locked layer.
 
-        Platform-specific and not implemented here. A locked bank layer is
-        CPU-executor-readable but must never be handed to the GPU movement
-        paths.
-        """
-        raise NotImplementedError("HostBank.lock() is platform-specific and not implemented")
+        Lock after fill, or the lazy mmap faults+zero-fills every page. A failed lock (RLIMIT_MEMLOCK) warns once and leaves the bank PAGEABLE, which every consumer treats the same."""
+        if self._locked or self._pinned:  # cudaHostRegister already page-locks
+            return
+        global _os_lock_failed
+        if _os_lock_failed:
+            return  # the quota is exhausted for good; skip the syscall spam
+        try:
+            _os_lock(self.addr, len(self._buf))
+        except (OSError, ImportError) as exc:
+            _os_lock_failed = True
+            logger.warning(f"bank lock failed; leaving this and later banks pageable: {exc}")
+            return
+        self._locked = True
+
+
+_os_locked_total = 0  # bytes locked so far; the OS lock ceiling is a per-process quota
+_os_lock_failed = False  # sticky: once over quota, later (bigger-total) locks fail too
+
+
+def _os_lock(addr: int, nbytes: int) -> None:
+    global _os_locked_total
+    import resource
+
+    # grow the soft RLIMIT_MEMLOCK (defaults to a few MiB); the hard limit needs privilege, past it mlock fails below
+    want = _os_locked_total + nbytes + (256 << 20)
+    soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK)
+    if soft != resource.RLIM_INFINITY and soft < want:
+        new_soft = want if hard == resource.RLIM_INFINITY else min(want, hard)
+        if new_soft > soft:
+            try:
+                resource.setrlimit(resource.RLIMIT_MEMLOCK, (new_soft, hard))
+            except (OSError, ValueError):
+                pass  # keep the old limit; mlock below reports the real ceiling
+    libc = ctypes.CDLL(None, use_errno=True)
+    if libc.mlock(ctypes.c_void_p(addr), ctypes.c_size_t(nbytes)):
+        err = ctypes.get_errno()
+        raise OSError(
+            err,
+            f"mlock({nbytes / 2**30:.1f} GiB): {os.strerror(err)} "
+            f"(RLIMIT_MEMLOCK / `ulimit -l` caps OS-locked bytes; raise it or "
+            f"shrink --moe-cpu-layers)",
+        )
+    _os_locked_total += nbytes
 
 
 def alloc_banks(specs: dict[str, tuple[tuple[int, ...], torch.dtype]]) -> dict[str, HostBank]:
@@ -127,24 +214,79 @@ def alloc_layer_banks(
     }
 
 
+class _ResidencyPlan:
+    """Per-layer ``HostResidency`` labels, ambiently visible to the bank settle points.
+
+    Installed by ``load_expert_banks`` around the provider dispatch so every loader honors --moe-cpu-layers without a new parameter in each signature. ``applied`` flips once a settle point consults the plan."""
+
+    __slots__ = ("labels", "applied", "has_unpinned", "actual")
+
+    def __init__(self, labels: list[str]):
+        self.labels = list(labels)
+        self.applied = False
+        self.has_unpinned = any(r != HostResidency.PINNED.value for r in labels)
+        self.actual: dict[int, str] = {}
+
+    def residency_for(self, layer_id: int) -> str:
+        self.applied = True
+        return self.labels[layer_id]
+
+    def record(self, layer_id: int, achieved: str) -> None:
+        """One pageable bank downgrades the whole layer (a failed lock settles PAGEABLE)."""
+        if self.actual.get(layer_id) != HostResidency.PAGEABLE.value:
+            self.actual[layer_id] = achieved
+
+
+_requested_residency: _ResidencyPlan | None = None
+
+
+@contextlib.contextmanager
+def requested_residency(labels: list[str] | None):
+    """Install the ambient per-layer residency plan for the enclosed bank load (``None`` = no plan, everything pins)."""
+    global _requested_residency
+    if labels is None:
+        yield None
+        return
+    plan = _ResidencyPlan(labels)
+    prev, _requested_residency = _requested_residency, plan
+    try:
+        yield plan
+    finally:
+        _requested_residency = prev
+
+
+def _settle(bank: HostBank, residency: str) -> None:
+    """Route a filled bank to its residency class (PAGEABLE = leave the plain mmap)."""
+    if residency == HostResidency.PINNED.value:
+        bank.pin()
+    elif residency == HostResidency.LOCKED.value:
+        bank.lock()
+
+
 def pin_banks(banks: dict[str, HostBank | list[HostBank]]) -> None:
-    """Pin (register) every bank after it has been filled -- pin-after-fill."""
+    """Settle every bank after it has been filled -- pin-after-fill by default.
+    List-valued entries are per-layer and honor the ambient :func:`requested_residency` plan; scalar banks always pin."""
+    plan = _requested_residency
     for bank in banks.values():
         if isinstance(bank, list):
-            for layer_bank in bank:
-                layer_bank.pin()
+            for layer_id, layer_bank in enumerate(bank):
+                residency = (
+                    HostResidency.PINNED.value if plan is None
+                    else plan.residency_for(layer_id)
+                )
+                _settle(layer_bank, residency)
+                if plan is not None and residency == HostResidency.LOCKED.value:
+                    plan.record(layer_id, layer_bank.residency.value)
         else:
             bank.pin()
 
 
 class PinPipeline:
-    """Pin filled banks while other banks are still being read.
+    """Settle (pin or lock) filled banks while other banks are still being read.
 
-    cudaHostRegister is driver-serialized, so one background thread drains a
-    queue of completed banks; submitters never block. Total load time becomes
-    ~max(read, pin) instead of their sum. Use as a context manager: a clean exit
-    drains the queue and re-raises the first pin failure; an exceptional exit
-    still joins the thread but lets the original exception propagate.
+    cudaHostRegister is driver-serialized, so one background thread drains a queue and submitters never block: load time ~= max(read, settle).
+    LOCKED banks mlock on the same thread (the quota bookkeeping in ``_os_lock`` is not thread-safe).
+    A clean context-manager exit drains the queue and re-raises the first settle failure.
     """
 
     def __init__(self) -> None:
@@ -155,23 +297,31 @@ class PinPipeline:
 
     def _run(self) -> None:
         while True:
-            bank = self._q.get()
-            if bank is None:
+            item = self._q.get()
+            if item is None:
                 return
             if self._exc is not None:
-                continue  # drain without pinning after a failure
+                continue  # drain without settling after a failure
+            bank, residency, plan, layer_id = item
             try:
-                bank.pin()
+                _settle(bank, residency)
+                if plan is not None and residency == HostResidency.LOCKED.value:
+                    plan.record(layer_id, bank.residency.value)
             except BaseException as exc:  # surfaced by wait()/__exit__
                 self._exc = exc
 
-    def submit(self, bank: HostBank) -> None:
-        self._q.put(bank)
+    def submit(self, bank: HostBank, residency: str = HostResidency.PINNED.value,
+               plan=None, layer_id: int | None = None) -> None:
+        self._q.put((bank, residency, plan, layer_id))
 
     def __call__(self, layer_id: int, banks: dict[str, HostBank]) -> None:
-        """Layer-completion sink: queue every bank of the completed layer."""
+        """Layer-completion sink: queue every bank of the completed layer at its ambient :func:`requested_residency` label."""
+        plan = _requested_residency
+        residency = (
+            HostResidency.PINNED.value if plan is None else plan.residency_for(layer_id)
+        )
         for bank in banks.values():
-            self.submit(bank)
+            self.submit(bank, residency, plan, layer_id)
 
     def _join(self) -> None:
         self._q.put(None)
@@ -263,6 +413,8 @@ __all__ = [
     "PinPipeline",
     "alloc_banks",
     "alloc_layer_banks",
+    "born_pinned_default",
     "pin_banks",
     "read_file_into",
+    "requested_residency",
 ]

python/freetoken/moe/offload_cache.py

@@ -77,6 +77,17 @@ _BANK_SCHEMAS: dict[str, tuple[str, ...]] = {
     "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"),
 }
 
+# bytes per (expert, layer) as f(hidden, moe_intermediate), from the bank shapes above; keep in sync with _BANK_SCHEMAS
+# keyed by the config-time format tag (expert_quant / moe_weight_format), not quant_format: "mxfp4" sizes the mxfp4_triton banks, "nvfp4" also covers its repacked variants
+_BANK_BYTES_PER_EXPERT = {
+    "bf16": lambda H, I: 3 * I * H * 2,
+    "fp8_block": lambda H, I: 3 * I * H + ((2 * I // 128) * (H // 128) + (H // 128) * (I // 128)) * 2,
+    "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18,
+    "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2),
+    "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2),
+    "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32),
+}
+
 # vLLM's marlin grouped-GEMM hands the full [cache_size] slot cache as its expert
 # dimension; moe_align_block_size requires round_up(experts, 32) < 1024, i.e. <= 992.
 MARLIN_MAX_CACHE_SIZE = 992
@@ -162,8 +173,10 @@ class OffloadMoeCache:
         self.usage = torch.zeros((self.cache_size,), dtype=torch.int64, device=self.device)
         self.step = torch.zeros((), dtype=torch.int64, device=self.device)
         self.active_mask = torch.zeros((self.num_experts,), dtype=torch.int32, device=self.device)
-        self.evict_slots = torch.empty((self.num_experts,), dtype=torch.int32, device=self.device)
-        self.src_indices = torch.empty((self.num_experts,), dtype=torch.int32, device=self.device)
+        # lru_ensure validates these against plan = min(batch * top_k, cache_size), so num_experts elements would under-size them
+        plan_slots = max(self.num_experts, self.cache_size)
+        self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
+        self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
         self.num_indices = torch.zeros((1,), dtype=torch.int64, device=self.device)
         # hybrid only: full missing count BEFORE the per-step fetch cap (num_indices holds
         # the capped count that copy_missing actually fetches). The difference is what the
@@ -182,10 +195,10 @@ class OffloadMoeCache:
         self.bank_schema = _BANK_SCHEMAS[self.quant_format]
         self.bank_sources: dict[str, list[torch.Tensor]] = {}
         self.bank_caches: dict[str, torch.Tensor] = {}
-        # Per-layer host residency (HostResidency values). The GPU movement paths
-        # (fused gather, prefill DMA) require "pinned"; other residency classes
-        # are not supported here and are rejected by set_bank_sources.
+        # per-layer host residency: the GPU movement paths require "pinned"; LOCKED/PAGEABLE layers decode on the CPU executor and prefill via copy_missing's pageable branch
+        # _unpinned_layers is the derived id set the hot paths test against
         self.layer_residency: list[str] = []
+        self._unpinned_layers: frozenset = frozenset()
         # marlin/b12x per-expert global scales ([L*E], GPU resident, see set_alphas).
         self.gate_up_alpha: torch.Tensor | None = None
         self.down_alpha: torch.Tensor | None = None
@@ -236,7 +249,9 @@ class OffloadMoeCache:
         # The layer whose misses ensure_experts/materialize_layer staged last; consumed
         # by copy_missing to pick the per-layer source (part of the same pending-copy
         # state as evict_slots/src_indices/num_indices).
+        # _pending_whole_layer records WHICH staged it: the pageable branch is only sound after materialize_layer
         self._pending_src_layer: int | None = None
+        self._pending_whole_layer = False
         # Per-bank [2, num_experts, ...] double-buffer views over the slot cache's
         # first 2 * num_experts slots (set up when prefill_overlap is enabled).
         self.prefill_bank_buffers: list[torch.Tensor] = []
@@ -275,10 +290,8 @@ class OffloadMoeCache:
         repackers (see ``_BANK_SCHEMAS`` and :mod:`freetoken.moe.nvfp4_backends`)
         -- the cache machinery is layout-agnostic and just moves rows.
 
-        ``layer_residency`` labels each layer with a ``HostResidency`` value
-        (default: all pinned). Non-pinned layers have no device address, so the
-        GPU movement paths cannot serve them and they are rejected here
-        (platform-specific residency policies are not implemented).
+        ``layer_residency`` labels each layer with a ``HostResidency`` value (default: all pinned).
+        Non-pinned (LOCKED/PAGEABLE) layers have no device address: they must already be routed to the CPU executor (``cpu_layer_ids``, set BEFORE this call), the copy plan skips their rows, and their only movement is ``copy_missing``'s whole-layer pageable prefill branch -- which is why prefill overlap is incompatible with them.
         """
         from freetoken.moe.host_banks import HostResidency
 
@@ -288,11 +301,22 @@ class OffloadMoeCache:
         )
         residency = layer_residency or [HostResidency.PINNED.value] * self.num_layers
         assert len(residency) == self.num_layers, (len(residency), self.num_layers)
-        if any(r != HostResidency.PINNED.value for r in residency):
-            raise NotImplementedError(
-                "non-pinned host bank layers need platform-specific movement "
-                "paths that are not implemented; only pinned layers are served"
-            )
+        unpinned = frozenset(
+            i for i, r in enumerate(residency) if r != HostResidency.PINNED.value
+        )
+        if unpinned:
+            if not unpinned <= self.cpu_layer_ids:
+                raise ValueError(
+                    f"non-pinned layers {sorted(unpinned - self.cpu_layer_ids)} are not in "
+                    f"cpu_layer_ids: a layer without a device address can only decode on "
+                    f"the CPU executor (set cache.cpu_layer_ids before set_bank_sources)"
+                )
+            if self.prefill_overlap:
+                raise ValueError(
+                    "prefill overlap DMAs from registered banks; it must be disabled "
+                    "when any layer is LOCKED/PAGEABLE (the engine does this)"
+                )
+        self._unpinned_layers = unpinned
         self.layer_residency = list(residency)
         for name in self.bank_schema:
             per_layer = sources[name]
@@ -344,6 +368,11 @@ class OffloadMoeCache:
             if feat % 16 != 0 or cache.data_ptr() % 16 != 0:
                 return  # leave fused disabled; copy_missing uses the per-bank path
             for layer_id, source in enumerate(per_layer):
+                if layer_id in self._unpinned_layers:
+                    # unregistered layer: no device alias exists, and the row is never consumed (CPU decode; pageable prefill)
+                    # a 0 placeholder keeps the descriptor shape
+                    layer_src_ptrs[layer_id].append(0)
+                    continue
                 # The kernel dereferences these on the GPU, so store each host bank's
                 # device alias (== data_ptr() under UVA identity; differs on
                 # Windows/WDDM).
@@ -429,6 +458,9 @@ class OffloadMoeCache:
         self.slot_for_id.fill_(-1)
         self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device)
         self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device)
+        plan_slots = max(self.num_experts, cache_size)
+        self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
+        self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device)
         self.step.zero_()
         self.active_mask.zero_()
         self.num_indices.zero_()
@@ -439,6 +471,8 @@ class OffloadMoeCache:
         self.stat_calls.zero_()
         self.stat_fetched.zero_()
         self.stat_missing_layer.zero_()
+        # a rebuild is a cold start for the cache; carrying pre-rebuild hit/miss counts over would skew every post-rebuild stats report
+        self.lru_stats.zero_()
         self.stat_active_layer.zero_()
         self.stat_fetched_layer.zero_()
         self.stat_steps_layer.zero_()
@@ -491,6 +525,11 @@ class OffloadMoeCache:
         """Whether ``layer_id`` decodes on the CPU executor (vs the GPU offload path)."""
         return layer_id in self.cpu_layer_ids
 
+    def is_unpinned_layer(self, layer_id: int) -> bool:
+        """Whether ``layer_id``'s host banks have no device address (LOCKED/PAGEABLE): the GPU slot-gather paths cannot serve it.
+        ``copy_missing`` takes the whole-layer pageable branch, which presumes materialize's position == expert id (never ``ensure_experts``'s LRU slot remap)."""
+        return layer_id in self._unpinned_layers
+
     def alphas_for_slots(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | None:
         """Per-slot global scales for a decode call, or ``None`` when the format
         keeps no GPU-resident alphas (bf16 / triton-nvfp4). Slots of other layers
@@ -767,6 +806,7 @@ class OffloadMoeCache:
             ids = expert_ids.reshape(-1).long()
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
+        self._pending_whole_layer = False
         ensure_experts(self, layer_id, expert_ids)
 
     def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None:
@@ -785,6 +825,7 @@ class OffloadMoeCache:
             ids = expert_ids.reshape(-1).long()
             self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids))
         self._pending_src_layer = layer_id
+        self._pending_whole_layer = False
         ensure_experts_hybrid(
             self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction
         )
@@ -793,6 +834,7 @@ class OffloadMoeCache:
         from freetoken.moe.offload_kernels import materialize_layer
 
         self._pending_src_layer = layer_id
+        self._pending_whole_layer = True
         materialize_layer(self, layer_id)
 
     def reset(self) -> None:
@@ -927,6 +969,18 @@ class OffloadMoeCache:
         assert self.banks, "set_bank_sources must register the banks first"
         layer_id = self._pending_src_layer
         assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)"
+        if layer_id in self._unpinned_layers:
+            if not self._pending_whole_layer:
+                raise RuntimeError(
+                    f"layer {layer_id} is unpinned: its only copy is the whole-layer "
+                    f"pageable materialize (position == expert id); ensure_experts's "
+                    f"LRU slot remap cannot be honored without a device alias"
+                )
+            # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into slots [0, num_experts) with position == expert id -- a plain synchronous pageable H2D copy
+            # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor)
+            for per_layer, cache in self.banks:
+                cache[: self.num_experts].copy_(per_layer[layer_id])
+            return
         if self._copy_fused_ok:
             from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit
 

python/freetoken/server/args.py

@@ -529,10 +529,13 @@ def parse_args(
         type=str,
         default=ServerArgs.moe_cpu_layers,
         help=(
-            "Hybrid decode with --moe-backend offload: which MoE layers compute on the "
-            "CPU executor instead of the GPU offload/PCIe path. Explicit id list "
-            "('3,7,11'), a count ('8' = 8 layers evenly strided), or a fraction ('0.5'). "
-            "Unset = all layers on GPU."
+            "With --moe-backend offload/hybrid: which MoE layers compute on the "
+            "CPU executor instead of the GPU offload/PCIe path (where CUDA pinning "
+            "is quota-capped, e.g. WSL, their banks are OS-locked instead of pinned). Explicit id list ('3,7,11'), a count ('8' = 8 "
+            "layers evenly strided), or a fraction ('0.5'). Unset = automatic where "
+            "CUDA pinning is quota-capped, e.g. WSL (locks just enough head+tail "
+            "layers when the banks exceed the pin budget, none otherwise); '0' "
+            "forces all layers on GPU."
         ),
     )
 

tests/moe/test_offload.py

@@ -675,3 +675,196 @@ def test_offload_cache_validate_rebuild_enforces_marlin_cap_and_floor():
     bf16 = OffloadMoeCache(num_layers=1, num_experts=4, cache_size=6, device=torch.device("cpu"))
     with pytest.raises(ValueError, match="num_experts"):
         bf16.validate_rebuild(3)  # below the num_experts floor
+
+
+def _make_split_cache(num_layers=2, locked=(1,), prefill_overlap=False, device="cpu"):
+    """A [gate_up, down] bf16 cache with the given layers LOCKED (rest pinned)."""
+    from freetoken.moe.host_banks import HostResidency
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    dev = torch.device(device)
+    cache = OffloadMoeCache(
+        num_layers=num_layers, num_experts=4, cache_size=8,
+        device=dev, prefill_overlap=prefill_overlap,
+    )
+    cache.cpu_layer_ids = frozenset(locked)
+    src_dev = dev if dev.type == "cuda" else torch.device("cpu")
+    sources = {
+        # CUDA-resident pinned-layer sources keep _build_copy_plan's device_ptr happy in the CUDA variant; locked layers stay host tensors (never translated)
+        "gate_up": [
+            torch.randn(4, 32, 8, device=torch.device("cpu") if i in locked else src_dev)
+            for i in range(num_layers)
+        ],
+        "down": [
+            torch.randn(4, 8, 16, device=torch.device("cpu") if i in locked else src_dev)
+            for i in range(num_layers)
+        ],
+    }
+    residency = [
+        HostResidency.LOCKED.value if i in locked else HostResidency.PINNED.value
+        for i in range(num_layers)
+    ]
+    cache.set_bank_sources(sources, layer_residency=residency)
+    return cache, sources
+
+
+def test_set_bank_sources_locked_layer_requires_cpu_layer_ids():
+    # a layer without a device address can only decode on the CPU executor; labeling it LOCKED outside cpu_layer_ids is a wiring bug and must fail loudly
+    from freetoken.moe.host_banks import HostResidency
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=8, device=torch.device("cpu"),
+    )
+    sources = {
+        "gate_up": [torch.randn(4, 32, 8) for _ in range(2)],
+        "down": [torch.randn(4, 8, 16) for _ in range(2)],
+    }
+    with pytest.raises(ValueError, match="cpu_layer_ids"):
+        cache.set_bank_sources(
+            sources,
+            layer_residency=[HostResidency.PINNED.value, HostResidency.LOCKED.value],
+        )
+
+
+def test_set_bank_sources_locked_layer_rejects_prefill_overlap():
+    # prefill overlap DMAs from registered banks; a LOCKED layer cannot feed it
+    from freetoken.moe.host_banks import HostResidency
+    from freetoken.moe.offload_cache import OffloadMoeCache
+
+    _init_tp()
+    cache = OffloadMoeCache(
+        num_layers=2, num_experts=4, cache_size=8, device=torch.device("cpu"),
+        prefill_overlap=True,
+    )
+    cache.cpu_layer_ids = frozenset({1})
+    sources = {
+        "gate_up": [torch.randn(4, 32, 8) for _ in range(2)],
+        "down": [torch.randn(4, 8, 16) for _ in range(2)],
+    }
+    with pytest.raises(ValueError, match="[Pp]refill overlap"):
+        cache.set_bank_sources(
+            sources,
+            layer_residency=[HostResidency.PINNED.value, HostResidency.LOCKED.value],
+        )
+
+
+def test_locked_layer_prefill_materialize_copies_whole_layer_pageable():
+    # the only movement a LOCKED layer needs: copy_missing's pageable branch copies the whole layer into slots [0, E) with position == expert id
+    # stage the state materialize_layer would (its kernel is CUDA-only; the fixture cache lives on the CPU)
+    cache, sources = _make_split_cache(num_layers=2, locked=(1,))
+
+    cache._pending_src_layer = 1
+    cache._pending_whole_layer = True
+    cache.copy_missing()
+
+    gate_up_cache, down_cache = (c for _, c in cache.banks)
+    assert torch.equal(gate_up_cache[:4], sources["gate_up"][1])
+    assert torch.equal(down_cache[:4], sources["down"][1])
+    # (The pinned layers' staged JIT path is covered by the mocked tests above.)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+def test_copy_plan_skips_locked_layers_and_keeps_fused_path():
+    # _build_copy_plan must not resolve a device alias for a LOCKED layer; its descriptor row stays a 0 placeholder while the pinned layers keep the fused path
+    cache, _ = _make_split_cache(num_layers=2, locked=(1,), device="cuda")
+
+    assert cache._copy_fused_ok
+    assert (cache._copy_src_ptrs[1] == 0).all(), "locked layer row must stay 0"
+    assert (cache._copy_src_ptrs[0] != 0).all(), "pinned layer rows must resolve"
+
+
+def test_locked_layer_copy_missing_rejects_ensure_experts_staging():
+    # the pageable branch presumes materialize_layer's position == expert id; staging via ensure_experts (LRU slot remap) on a locked layer must fail loudly, not gather other experts' weights
+    # stage the state ensure_experts would (its kernel is CUDA-only; the fixture cache lives on the CPU)
+    cache, _ = _make_split_cache(num_layers=2, locked=(1,))
+
+    cache._pending_src_layer = 1
+    cache._pending_whole_layer = False
+    with pytest.raises(RuntimeError, match="unpinned"):
+        cache.copy_missing()
+
+
+def test_requested_residency_routes_layer_settles(monkeypatch):
+    # the ambient plan installed by load_expert_banks must route each layer's banks by label at both slow-path settle points (PinPipeline layer sink, list-valued pin_banks) and record that it was consulted
+    # without a plan everything pins
+    import freetoken.moe.host_banks as hb
+
+    settled = []
+    monkeypatch.setattr(hb.HostBank, "pin", lambda self: settled.append("pin"))
+    monkeypatch.setattr(hb.HostBank, "lock", lambda self: settled.append("lock"))
+    banks = {
+        "gate_up": [hb.HostBank((4,), torch.uint8) for _ in range(3)],
+        "down": [hb.HostBank((4,), torch.uint8) for _ in range(3)],
+    }
+    labels = [
+        hb.HostResidency.PINNED.value,
+        hb.HostResidency.LOCKED.value,
+        hb.HostResidency.PAGEABLE.value,
+    ]
+
+    with hb.requested_residency(labels) as plan:
+        with hb.PinPipeline() as pins:
+            for layer_id in range(3):
+                pins(layer_id, {name: per[layer_id] for name, per in banks.items()})
+    # the single drain thread settles FIFO: layer 0 pins, layer 1 locks, layer 2 passes
+    assert settled == ["pin", "pin", "lock", "lock"]
+    assert plan.applied
+
+    settled.clear()
+    with hb.requested_residency(labels) as plan:
+        hb.pin_banks(banks)
+    assert settled == ["pin", "lock", "pin", "lock"]  # per name: layer 0 pin, 1 lock, 2 skip
+    assert plan.applied
+
+    settled.clear()
+    hb.pin_banks(banks)  # no ambient plan -> every layer pins
+    assert settled == ["pin"] * 6
+
+
+def test_echo_residency_stamps_honored_requests_only():
+    # load_expert_banks stamps the request onto the provider's ExpertBanks only when a settle point consulted the plan; an unconsulted plan keeps None (the engine's degrade signal)
+    from freetoken.moe.expert_banks import ExpertBanks, _echo_residency
+    from freetoken.moe.host_banks import HostResidency, _ResidencyPlan
+
+    labels = [HostResidency.PINNED.value, HostResidency.LOCKED.value]
+    banks = ExpertBanks("bf16", {"gate_up": [], "down": []})
+
+    plan = _ResidencyPlan(labels)
+    plan.residency_for(1)  # a settle point consulted the plan
+    assert _echo_residency(banks, labels, plan).layer_residency == labels
+
+    stale = _ResidencyPlan(labels)  # never consulted -> keep None + warn
+    assert _echo_residency(banks, labels, stale).layer_residency is None
+    assert _echo_residency(banks, None, None) is banks
+
+
+def test_lock_failure_downgrades_echoed_residency(monkeypatch):
+    # a failed mlock leaves the bank pageable; the plan and the echoed labels must report that instead of the requested LOCKED
+    import freetoken.moe.host_banks as hb
+    from freetoken.moe.expert_banks import ExpertBanks, _echo_residency
+
+    def boom(addr, nbytes):
+        raise OSError(12, "mlock denied")
+
+    monkeypatch.setattr(hb, "_os_lock", boom)
+    monkeypatch.setattr(hb, "_os_lock_failed", False)
+    monkeypatch.setenv("FREETOKEN_SKIP_BANK_PIN", "1")  # keep the pinned layer off CUDA
+    labels = [hb.HostResidency.PINNED.value, hb.HostResidency.LOCKED.value]
+
+    banks = {"gate_up": [hb.HostBank((4,), torch.uint8) for _ in range(2)]}
+    with hb.requested_residency(labels) as plan:
+        hb.pin_banks(banks)
+    assert plan.actual == {1: hb.HostResidency.PAGEABLE.value}
+    echoed = _echo_residency(ExpertBanks("bf16", {}), labels, plan)
+    assert echoed.layer_residency == [
+        hb.HostResidency.PINNED.value, hb.HostResidency.PAGEABLE.value,
+    ]
+
+    monkeypatch.setattr(hb, "_os_lock_failed", False)
+    with hb.requested_residency(labels) as plan2:
+        with hb.PinPipeline() as pins:
+            pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)})
+    assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value}