From c41833b71740e84fc4a30aa4b4b28665e56cfa81 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 23 Aug 2026 08:15:17 +0000 Subject: [PATCH 1/3] feat(moe): per-layer host-bank residency -- split lock-CPU/pin-GPU layers, auto selection under capped pin quota --- python/freetoken/checkpoint/ftw.py | 64 ++++-- python/freetoken/engine/engine.py | 129 +++++++++++- python/freetoken/models/deepseek_v4/moe.py | 8 +- python/freetoken/moe/expert_banks.py | 76 +++++-- python/freetoken/moe/host_banks.py | 231 +++++++++++++++++---- python/freetoken/moe/offload_cache.py | 71 +++++-- python/freetoken/server/args.py | 11 +- tests/moe/test_offload.py | 164 +++++++++++++++ 8 files changed, 652 insertions(+), 102 deletions(-) diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index 4e39a4d9..9cca52d0 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/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,28 @@ def _read_layer(job): from freetoken.moe.expert_banks import ExpertBanks + unpinned = [i for i, r in enumerate(residency) 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 + 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 unpinned) + 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(unpinned)} CPU layers: {unpinned})" + ) + # 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=list(residency), + ) __all__ = [ diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985f..9ce193ae 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -512,12 +512,46 @@ 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 ftw_bank_bytes + + budget = _pin_budget_bytes() + bank_bytes = ftw_bank_bytes(config.model_path) if budget is not None else None + 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 +559,15 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # --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 +576,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 +610,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 +1117,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 ftw_bank_bytes + + bank_bytes = ftw_bank_bytes(config.model_path) + 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 +1314,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 +1450,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" ) diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 9b28922e..d7e555a9 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/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: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e456c931..0c8be6a7 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -28,6 +28,9 @@ 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,19 @@ 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 load_expert_banks( model_path: str, model_config, @@ -383,6 +398,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 +416,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 +459,35 @@ 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 + + return dataclasses.replace(banks, layer_residency=list(requested)) + 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 diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 5b50b25e..733349b5 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -17,6 +17,7 @@ from __future__ import annotations +import contextlib import ctypes import math import mmap @@ -28,18 +29,18 @@ 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 @@ class HostResidency(str, Enum): # 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()``. + + * ``"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`. - 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``.""" + The buffer is rounded up to the O_DIRECT block; ``tensor`` views exactly ``nbytes``. ``backing=None`` follows ``FREETOKEN_BANK_CUDA_ALLOC``.""" - __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned") + __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned", "_locked") - def __init__(self, shape: tuple[int, ...], dtype: torch.dtype): + 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 @@ def pin(self) -> None: 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. + + 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 + - 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") +_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,71 @@ 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") + + 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) + + def residency_for(self, layer_id: int) -> str: + self.applied = True + return self.labels[layer_id] + + +_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): + _settle( + layer_bank, + HostResidency.PINNED.value if plan is None + else plan.residency_for(layer_id), + ) 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 +289,28 @@ def __init__(self) -> None: 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 = item try: - bank.pin() + _settle(bank, residency) 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) -> None: + self._q.put((bank, residency)) 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) def _join(self) -> None: self._q.put(None) @@ -263,6 +402,8 @@ def rd(o): "PinPipeline", "alloc_banks", "alloc_layer_banks", + "born_pinned_default", "pin_banks", "read_file_into", + "requested_residency", ] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index debea282..b4d1d89f 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -162,8 +162,10 @@ def __post_init__(self) -> None: 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 +184,10 @@ def __post_init__(self) -> None: 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 +238,9 @@ def __post_init__(self) -> None: # 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 +279,8 @@ def set_bank_sources( 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 +290,22 @@ def set_bank_sources( ) 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 +357,11 @@ def _build_copy_plan(self) -> None: 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 +447,9 @@ def rebuild(self, cache_size: int) -> None: 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 +460,8 @@ def rebuild(self, cache_size: int) -> None: 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 +514,11 @@ def is_cpu_layer(self, layer_id: int) -> bool: """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 +795,7 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: 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 +814,7 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None 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 +823,7 @@ def materialize_layer(self, layer_id: int) -> None: 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 +958,18 @@ def copy_missing(self) -> None: 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 diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b75..4954c5f5 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -529,10 +529,13 @@ def _infer_reasoning_parser(model_path: str) -> str | None: 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." ), ) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 3caab91b..6f2b8bff 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -675,3 +675,167 @@ 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 From eebb3f53f357315eff8a92baba2e61ce1336c333 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 23 Aug 2026 23:06:56 +0000 Subject: [PATCH 2/3] feat(moe): estimate expert-bank bytes from the model config --- python/freetoken/engine/engine.py | 10 ++++++---- python/freetoken/moe/expert_banks.py | 21 ++++++++++++++++++++- python/freetoken/moe/offload_cache.py | 11 +++++++++++ 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 9ce193ae..22c4a6c7 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -535,10 +535,12 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: ) 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 ftw_bank_bytes + from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes budget = _pin_budget_bytes() - bank_bytes = ftw_bank_bytes(config.model_path) if budget is not None else None + 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( @@ -1159,9 +1161,9 @@ 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 ftw_bank_bytes + from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes - bank_bytes = ftw_bank_bytes(config.model_path) + 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() diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 0c8be6a7..b6fadff9 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -24,7 +24,7 @@ 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__) @@ -386,6 +386,25 @@ def ftw_bank_bytes(model_path: str) -> int | None: 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, diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index b4d1d89f..6ee76406 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -77,6 +77,17 @@ "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 From 831d38a66bc45543411078c6805de8c42e7603fa Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 23 Aug 2026 23:06:56 +0000 Subject: [PATCH 3/3] fix(moe): report the residency banks actually settle at --- python/freetoken/checkpoint/ftw.py | 25 ++++++++++++++++++++---- python/freetoken/moe/expert_banks.py | 10 +++++++++- python/freetoken/moe/host_banks.py | 27 ++++++++++++++++++-------- tests/moe/test_offload.py | 29 ++++++++++++++++++++++++++++ 4 files changed, 78 insertions(+), 13 deletions(-) diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index 9cca52d0..9d6df684 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -582,19 +582,36 @@ def _read_layer(job): from freetoken.moe.expert_banks import ExpertBanks - unpinned = [i for i, r in enumerate(residency) if r != HostResidency.PINNED.value] + # 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 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(unpinned)} CPU layers: {unpinned})" + 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 @@ -602,7 +619,7 @@ def _read_layer(job): alpha_kw = {n: alpha_hb[n].tensor for n in alpha_hb} return ExpertBanks( reader.meta("quant_format"), sources, **alpha_kw, - layer_residency=list(residency), + layer_residency=applied, ) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index b6fadff9..8b6116ba 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -500,7 +500,15 @@ def _echo_residency(banks: ExpertBanks, requested, plan) -> ExpertBanks: if plan is not None and plan.applied: import dataclasses - return dataclasses.replace(banks, layer_residency=list(requested)) + 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): diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 733349b5..5923c823 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -219,17 +219,23 @@ class _ResidencyPlan: 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") + __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 @@ -264,11 +270,13 @@ def pin_banks(banks: dict[str, HostBank | list[HostBank]]) -> None: for bank in banks.values(): if isinstance(bank, list): for layer_id, layer_bank in enumerate(bank): - _settle( - layer_bank, + residency = ( HostResidency.PINNED.value if plan is None - else plan.residency_for(layer_id), + 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() @@ -294,14 +302,17 @@ def _run(self) -> None: return if self._exc is not None: continue # drain without settling after a failure - bank, residency = item + bank, residency, plan, layer_id = item try: _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, residency: str = HostResidency.PINNED.value) -> None: - self._q.put((bank, residency)) + 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 at its ambient :func:`requested_residency` label.""" @@ -310,7 +321,7 @@ def __call__(self, layer_id: int, banks: dict[str, HostBank]) -> None: HostResidency.PINNED.value if plan is None else plan.residency_for(layer_id) ) for bank in banks.values(): - self.submit(bank, residency) + self.submit(bank, residency, plan, layer_id) def _join(self) -> None: self._q.put(None) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 6f2b8bff..422ca867 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -839,3 +839,32 @@ def test_echo_residency_stamps_honored_requests_only(): 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}