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

python/freetoken/checkpoint/ftw.py

@@ -33,7 +33,6 @@
 from __future__ import annotations
 
 import json
-import logging
 import math
 import mmap
 import os
@@ -43,7 +42,9 @@
 
 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 _producer():
 
 
 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 _read_alpha(e):
                 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 _read_layer(job):
 
     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,19 +512,64 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
         # 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
             # expert-tensor granularity. Both pin-after-fill.
             # --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 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 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 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 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 override(attr: str, value: Any):  # this is dangerous, use with caution
     # 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 override(attr: str, value: Any):  # this is dangerous, use with caution
             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 @@ def _prefill_routed(
         # 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: