From cd5f27f1b689f282a69669129176d63587266615 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:39:17 +0000 Subject: [PATCH 01/38] feat(moe): disk tier v0 -- NVMe-backed MoE experts (VRAM <- RAM <- NVMe) Lets the offload backend serve models whose experts don't fit in pinned RAM: with --moe-disk-tier on --expert-ram-experts K, only the first K experts per layer are pinned; the rest keep allocated bank rows whose pages are released (MADV_DONTNEED) after load and are fetched from the ORIGINAL safetensors checkpoint on slot-cache miss. * moe/disk_tier.py: Nvfp4DiskIndex (index json + shard headers -> per (bank, layer, expert) byte ranges) and DiskTier (O_DIRECT preadv into a small pinned staging buffer, H2D into the LRU-assigned slot, miss-list rewrite so the existing PCIe copy path only moves RAM-resident misses). * host_banks: HostBank.pin_prefix / release_range; PinPipeline(prefix_rows). * nvfp4 loaders: disk_tier param -> partial pin + tail release (serial and parallel paths). * offload_cache: attach_disk_tier + copy_missing hook; materialize_layer takes the routed ids when disk-tiered. * offload_kernels: materialize kernel gains materialize_count (disk-tier prefill streams only the RAM prefix; routed disk-resident experts are fetched into their identity slots, so the prefill GEMM is unchanged). * engine: --moe-disk-tier / --expert-ram-experts / --disk-fetch-workers, guards (native NVFP4, gpu decode, no prefill overlap, no cuda graphs). v0 scope: native NVFP4 (triton) layout, synchronous fetch, no FTW, no converter path. CPU unit tests: tests/moe/test_disk_tier.py. --- python/freetoken/engine/config.py | 7 + python/freetoken/engine/engine.py | 25 ++ python/freetoken/layers/moe.py | 7 +- python/freetoken/models/nvfp4_banks.py | 20 +- python/freetoken/models/qwen3_5_moe/weight.py | 32 +- python/freetoken/models/weight.py | 9 +- python/freetoken/moe/disk_tier.py | 287 ++++++++++++++++++ python/freetoken/moe/expert_banks.py | 31 +- python/freetoken/moe/host_banks.py | 49 ++- python/freetoken/moe/offload_cache.py | 34 ++- python/freetoken/moe/offload_kernels.py | 11 +- python/freetoken/server/args.py | 27 ++ tests/moe/test_disk_tier.py | 246 +++++++++++++++ 13 files changed, 756 insertions(+), 29 deletions(-) create mode 100644 python/freetoken/moe/disk_tier.py create mode 100644 tests/moe/test_disk_tier.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index bcbe6bcf2..b5588a81d 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -41,6 +41,13 @@ class EngineConfig: # (cudaMemcpyBatchAsync); no-op unless moe_cache_size > 2 * num_experts. moe_prefill_hit_d2d: bool = False moe_collect_stats: bool = False # capture decode miss-rate counters into the cuda graph + # Disk tier (--moe-disk-tier, see moe/disk_tier.py): "off" = classic behavior. + # When "on", experts [0, expert_ram_experts) per layer stay pinned in RAM and the + # rest are fetched from the original checkpoint on slot-cache miss. Requires the + # native NVFP4 layout, gpu decode target, no prefill overlap, no cuda graphs. + moe_disk_tier: str = "off" + expert_ram_experts: int = 0 + disk_fetch_workers: int = 8 # CPU MoE backend (--moe-backend cpu): number of CPU worker threads computing # the decode experts. 0 = auto (physical cores). Ignored by other backends. moe_cpu_threads: int = 0 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b5a6fa3b0..266022beb 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -565,6 +565,22 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: "(locked layers prefill via synchronous pageable copies)" ) object.__setattr__(config, "moe_prefill_overlap", False) + disk_tier = None + if config.moe_disk_tier == "on": + from freetoken.moe.disk_tier import DiskTierSpec + + E = config.model_config.num_experts + if not 0 < config.expert_ram_experts < E: + raise ValueError( + f"--expert-ram-experts must be in (0, {E}) with --moe-disk-tier on") + if decode_target != "gpu": + raise ValueError( + "--moe-disk-tier v0 requires the gpu decode path (--moe-backend offload)") + if config.moe_prefill_overlap: + raise ValueError("--moe-disk-tier v0 requires --disable-moe-prefill-overlap") + if config.cuda_graph_max_bs is not None: + raise ValueError("--moe-disk-tier v0 requires cuda graphs disabled") + disk_tier = DiskTierSpec(ram_experts=config.expert_ram_experts) 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 @@ -590,6 +606,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: parallel=expert_parallel, decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"), layer_residency=requested_residency, + disk_tier=disk_tier, ) if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks) @@ -626,6 +643,14 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # 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) + if banks.disk_index is not None: + cache.attach_disk_tier( + banks.disk_index, banks.disk_ram_experts, + workers=config.disk_fetch_workers) + logger.info_rank0( + f"disk tier: {banks.disk_ram_experts}/{config.model_config.num_experts} " + f"experts/layer pinned in RAM; the rest fetched from " + f"{config.model_path} on slot-cache miss") cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) else: cache = cache_factory(config, self.device) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded5..0293a82d3 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -397,7 +397,12 @@ def _prefill_routed( ) cache.release_prefill_layer(self.layer_id) return out - cache.materialize_layer(self.layer_id) + if cache.disk_tier_enabled: + # Disk tier: stream only the RAM-resident prefix + fetch the routed + # disk-resident experts (identity slots, so topk_ids pass through). + cache.materialize_layer(self.layer_id, topk_ids) + else: + cache.materialize_layer(self.layer_id) cache.copy_missing() return self._expert_gemm( cache, diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 6b933ff1d..68a5cfbf0 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -86,6 +86,7 @@ def load_nvfp4_expert_source_banks( drop_page_cache: DropPageCache, primary: bool, layer_sink=None, + disk_tier=None, ) -> dict[str, list[torch.Tensor]]: """Build the 6 native NVFP4 source banks by streaming checkpoint shards (serial per-shard read). @@ -150,6 +151,9 @@ def load_nvfp4_expert_source_banks( drop_page_cache(path) _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill + K = disk_tier.ram_experts if disk_tier is not None else None + if disk_tier is not None and layer_sink is not None: + raise NotImplementedError("disk tier: the converter (layer_sink) path is not supported yet") gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] gate_up_global = [b.tensor for b in _hb["gate_up_global"]] @@ -202,8 +206,12 @@ def _load(sink) -> int: if layer_sink is not None: placed = _load(layer_sink) else: - with PinPipeline() as pins: + with PinPipeline(prefix_rows=K) as pins: placed = _load(pins) + if K is not None: + from freetoken.moe.disk_tier import release_bank_tails + + release_bank_tails(_hb, E, K) expected = num_layers * E * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" @@ -227,6 +235,7 @@ def load_nvfp4_expert_source_banks_parallel( workers: int = 8, chunk: int = 8 << 20, layer_sink=None, + disk_tier=None, ) -> dict[str, list[torch.Tensor]]: """parallel counterpart of :func:`load_nvfp4_expert_source_banks`, byte-for-byte same placement. bulk weight/weight_scale read via chunked multi-threaded O_DIRECT reader @@ -274,6 +283,9 @@ def load_nvfp4_expert_source_banks_parallel( drop_page_cache(path) _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill + K = disk_tier.ram_experts if disk_tier is not None else None + if disk_tier is not None and layer_sink is not None: + raise NotImplementedError("disk tier: the converter (layer_sink) path is not supported yet") gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] gate_up_global = [b.tensor for b in _hb["gate_up_global"]] @@ -321,8 +333,12 @@ def _load(sink) -> int: if layer_sink is not None: placed = _load(layer_sink) else: - with PinPipeline() as pins: + with PinPipeline(prefix_rows=K) as pins: placed = _load(pins) + if K is not None: + from freetoken.moe.disk_tier import release_bank_tails + + release_bank_tails(_hb, E, K) expected = num_layers * E * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..3818e9aac 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -834,7 +834,7 @@ def close(self) -> None: def setup_offload_expert_banks( model_path: str, model_config, *, device: torch.device, dtype: torch.dtype, dummy: bool = False, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20, - decode_target: str = "gpu", layer_sink=None, + decode_target: str = "gpu", layer_sink=None, disk_tier=None, ): """Build the routed-expert offload banks. The qwen3_5_moe module always exports this hook, so it intercepts *every* qwen3_5_moe offload load -- defer non-block-fp8 checkpoints (plain @@ -849,14 +849,29 @@ def setup_offload_expert_banks( providers for non-block-fp8 checkpoints. ``decode_target`` is forwarded so the cpu backend gets CPU-readable (native, non- - GPU-tiled) bank layouts -- e.g. native ``nvfp4`` rows rather than marlin/b12x.""" + GPU-tiled) bank layouts -- e.g. native ``nvfp4`` rows rather than marlin/b12x. + + ``disk_tier`` (a ``moe.disk_tier.DiskTierSpec``): pin only the first + ``ram_experts`` experts per layer, release the rest, and attach a + :class:`~freetoken.moe.disk_tier.Nvfp4DiskIndex` so the offload cache can + fetch the disk-resident experts on miss.""" eq = getattr(model_config, "expert_quant", "none") if eq != "fp8_block": from freetoken.moe.expert_banks import _PROVIDERS # nvfp4 -> _nvfp4_banks, none -> _bf16_banks - return _PROVIDERS[eq](model_path, model_config, device, dtype, dummy, - parallel=parallel, workers=workers, chunk=chunk, - decode_target=decode_target, layer_sink=layer_sink) + banks = _PROVIDERS[eq](model_path, model_config, device, dtype, dummy, + parallel=parallel, workers=workers, chunk=chunk, + decode_target=decode_target, layer_sink=layer_sink, + disk_tier=disk_tier) + if disk_tier is not None: + if eq != "nvfp4": + raise NotImplementedError( + f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") + from freetoken.moe.disk_tier import Nvfp4DiskIndex + + banks.disk_index = Nvfp4DiskIndex(model_path, model_config, _NVFP4_SOURCE_SPEC) + banks.disk_ram_experts = disk_tier.ram_experts + return banks if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only") from freetoken.moe.expert_banks import ExpertBanks @@ -1061,7 +1076,7 @@ def _load(sink) -> None: def load_nvfp4_expert_sources( - model_path: str, config, *, layer_sink=None + model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, torch.Tensor]: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" @@ -1072,11 +1087,13 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, + disk_tier=None, ): """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -1090,6 +1107,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9b..5c10e92a3 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -302,12 +302,14 @@ def load_nvfp4_moe_expert_sources( workers: int = 8, chunk: int = 8 << 20, layer_sink=None, + disk_tier=None, ) -> dict: """Load (or fabricate, with ``dummy=True``) packed NVFP4 expert source banks. ``parallel=True`` uses the model's ``load_nvfp4_expert_sources_parallel`` (common chunked multi-threaded O_DIRECT reader). ``layer_sink``: see ``models.nvfp4_banks.load_nvfp4_expert_source_banks``; forwarded to the per-model - loader, which forwards it on.""" + loader, which forwards it on. ``disk_tier`` (a ``moe.disk_tier.DiskTierSpec``) + pins only the first ``ram_experts`` experts per layer and releases the rest.""" _config, spec = _spec_for_model_path(model_path) if dummy: builder = ( @@ -319,9 +321,10 @@ def load_nvfp4_moe_expert_sources( if loader is None: # no parallel reader -> let the caller fall back to serial raise NotImplementedError( f"{spec.module} provides no load_nvfp4_expert_sources_parallel") - return loader(model_path, model_config, workers=workers, chunk=chunk, layer_sink=layer_sink) + return loader(model_path, model_config, workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier) loader = _load_attr(spec.module, "load_nvfp4_expert_sources") - return loader(model_path, model_config, layer_sink=layer_sink) + return loader(model_path, model_config, layer_sink=layer_sink, disk_tier=disk_tier) def load_q4_0_moe_expert_sources( diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py new file mode 100644 index 000000000..3aa315a9c --- /dev/null +++ b/python/freetoken/moe/disk_tier.py @@ -0,0 +1,287 @@ +"""Disk tier: NVMe-backed MoE experts (VRAM <- RAM <- NVMe). + +Lets the offload backend serve experts that do NOT fit in pinned RAM: the RAM +bank holds only the first ``ram_experts`` experts per layer (pinned), the rest +stay on disk in the original checkpoint. When the GPU slot cache misses a +disk-resident expert, :class:`DiskTier` fetches its rows with O_DIRECT preadv +into a small pinned staging buffer and H2D-copies them into the slot the LRU +kernel already assigned, then shrinks the miss list so the existing PCIe +``copy_missing`` path only moves the RAM-resident misses. + +v0 scope (prototype): +* native NVFP4 layout only (the "triton" backend banks -- what sm_120 picks); +* ``decode_target == "gpu"`` (offload) only -- the CPU executor reads banks + directly and would read released pages; +* synchronous fetch (the layer waits for its disk misses); no CUDA-graph + capture (the miss-list D2H/H2D round trip is host-side and variable); +* prefill_overlap off (the double-buffer prefill path bypasses the slot cache). + +The bank rows are read from the ORIGINAL safetensors shards: every expert +tensor is a contiguous per-expert tensor, so a bank row is one (or two, for +the gate|up-fused banks) aligned super-block preads. No FTW conversion needed. +""" + +from __future__ import annotations + +import ctypes +import json +import math +import os +import struct +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import torch + +from freetoken.moe.host_banks import HostBank + +_ALIGN = 4096 + + +@dataclass(frozen=True) +class DiskTierSpec: + """Engine -> loader: how many experts per layer stay pinned in RAM. + + Experts ``[0, ram_experts)`` are pinned as usual; ``[ram_experts, E)`` keep + their bank rows allocated but their pages are released after load and are + served from disk by :class:`DiskTier`.""" + + ram_experts: int + + +def release_bank_tails(banks_by_name: dict[str, list[HostBank]], num_experts: int, + ram_experts: int) -> None: + """MADV_DONTNEED the unpinned tail rows of every bank layer (post-load).""" + for layer_banks in banks_by_name.values(): + for bank in layer_banks: + row_bytes = bank.nbytes // num_experts + bank.release_range(ram_experts * row_bytes, bank.nbytes - ram_experts * row_bytes) + +# Native NVFP4 bank order (== _BANK_SCHEMAS["nvfp4"]) and, per bank, the +# checkpoint segments that make up one expert row: (proj, kind, dst_row_start, +# dst_row_end). The gate|up-fused banks splice gate rows then up rows on the +# output-row axis; down banks are a single segment. Row ends are None = rest. +_NVP4_BANK_SEGS = ( + (("gate_proj", "weight", 0, None), ("up_proj", "weight", None, None)), + (("gate_proj", "weight_scale", 0, None), ("up_proj", "weight_scale", None, None)), + (("gate_proj", "weight_scale_2", 0, None), ("up_proj", "weight_scale_2", None, None)), + (("down_proj", "weight", 0, None),), + (("down_proj", "weight_scale", 0, None),), + (("down_proj", "weight_scale_2", 0, None),), +) + + +def _read_safetensors_offsets(path: str) -> dict[str, tuple[int, int]]: + """{tensor_name: (start, end)} from a shard's safetensors header.""" + with open(path, "rb") as f: + (hlen,) = struct.unpack(" per-segment (shard_idx, offset, nbytes) locations. + + Built from the original checkpoint: the HF index json (name -> shard) plus + each referenced shard's safetensors header (name -> byte range). Expert + tensors are per-expert and contiguous, so a row is exactly one byte range + per segment. + """ + + def __init__(self, model_dir: str, config, spec) -> None: + from freetoken.models.nvfp4_banks import _num_moe_layers + from freetoken.utils.hf import download_hf_weight + + model_dir = download_hf_weight(model_dir) # hub id -> local cache dir; no-op if local + index_path = os.path.join(model_dir, "model.safetensors.index.json") + with open(index_path, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + + num_layers = _num_moe_layers(config) + # (bank_layer, expert, proj, kind) -> (tensor_name, shard) + loc: dict[tuple[int, int, str, str], tuple[str, str]] = {} + for name, shard in weight_map.items(): + m = spec.key_pattern.match(name) + if m is None: + continue + bank_layer = spec.layer_to_bank(int(m.group("layer")), config) + if bank_layer is None: + continue + loc[(bank_layer, int(m.group("expert")), m.group("proj"), m.group("kind"))] = ( + name, shard) + + shards = sorted(set(shard for _, shard in loc.values())) + self.shard_paths = [os.path.join(model_dir, s) for s in shards] + offsets = {s: _read_safetensors_offsets(os.path.join(model_dir, s)) for s in shards} + shard_idx = {s: i for i, s in enumerate(shards)} + + E = config.num_experts + seg_size = struct.calcsize(" packed segments per expert + for bank_idx in range(len(_NVP4_BANK_SEGS)): + per_layer = [] + for layer in range(num_layers): + rows = bytearray() + for e in range(E): + for proj, kind, _, _ in _NVP4_BANK_SEGS[bank_idx]: + key = (layer, e, proj, kind) + entry = loc.get(key) + if entry is None: + raise KeyError( + f"disk tier: no {proj}.{kind} tensor for layer {layer} expert {e} " + f"(bank {bank_idx}) in {index_path}" + ) + name, shard = entry + start, end = offsets[shard][name] + rows += struct.pack(" list[tuple[int, int, int]]: + """[(shard_idx, offset, nbytes)] for one expert row, in segment order.""" + base = expert * self._seg_size * len(_NVP4_BANK_SEGS[bank_idx]) + raw = self.entries[bank_idx][layer][base:base + self._seg_size * len(_NVP4_BANK_SEGS[bank_idx])] + return [ + struct.unpack_from(" staging -> GPU slot.""" + + def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int = 8) -> None: + self._index = index + self._ram = ram_experts + self._banks = list(cache.banks) # [(per_layer_host, gpu_cache)] in schema order + self._row_bytes = [ + math.prod(b[0][0][0].shape[1:]) * b[0][0][0].element_size() for b in self._banks + ] + # Per-bank destination row slices (gate|up split at the row midpoint). + self._dst_slices: list[list[tuple[int, int]]] = [] + for bank_idx, (host_layer, _gpu) in enumerate(self._banks): + row = host_layer[0][0] + if len(_NVP4_BANK_SEGS[bank_idx]) == 2: + mid = row.shape[0] // 2 + self._dst_slices.append([(0, mid), (mid, row.shape[0])]) + else: + self._dst_slices.append([(0, row.shape[0])]) + max_row = max(self._row_bytes) + self._staging_size = ((max_row + 2 * _ALIGN) // _ALIGN) * _ALIGN + self._staging = threading.local() + self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="disk-tier") + self._fd_lock = threading.Lock() + self._fds: dict[int, tuple[int, bool]] = {} + self._fetches = 0 + self._fetch_bytes = 0 + + # ------------------------------------------------------------------ fds + def _fd(self, shard_idx: int) -> tuple[int, bool]: + """(fd, o_direct) for a shard; O_DIRECT falls back to plain preadv where the + filesystem refuses it (tmpfs/overlayfs -- tests).""" + ent = self._fds.get(shard_idx) + if ent is None: + with self._fd_lock: + ent = self._fds.get(shard_idx) + if ent is None: + path = self._index.shard_paths[shard_idx] + try: + fd = os.open(path, os.O_RDONLY | os.O_DIRECT) + direct = True + except OSError: + fd = os.open(path, os.O_RDONLY) + direct = False + ent = (fd, direct) + self._fds[shard_idx] = ent + return ent + + # --------------------------------------------------------------- staging + def _staging_buf(self) -> HostBank: + buf = getattr(self._staging, "buf", None) + if buf is None: + buf = HostBank((self._staging_size,), torch.uint8) + buf.pin() # small; pin once per worker thread + self._staging.buf = buf + return buf + + # ---------------------------------------------------------------- fetch + def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: + staging = self._staging_buf() + for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): + row = gpu_cache[slot] + segs = self._index.row_segments(bank_idx, layer, expert) + for (d0, d1), (shard_idx, off, nbytes) in zip(self._dst_slices[bank_idx], segs): + fd, direct = self._fd(shard_idx) + if direct: + a0 = off & ~(_ALIGN - 1) + slen = (off + nbytes - a0 + _ALIGN - 1) & ~(_ALIGN - 1) + else: + a0, slen = off, nbytes + mv = (ctypes.c_char * slen).from_address(staging.addr) + os.preadv(fd, [mv], a0) + row_off = off - a0 + src = staging.tensor[row_off:row_off + nbytes] + dst = row[d0:d1] + dst.copy_(src.view(dst.shape), non_blocking=True) + self._fetches += 1 + self._fetch_bytes += sum(self._row_bytes) + + def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: + """Disk-tier prefill: materialize the RAM-resident prefix into identity slots + (the normal kernel restricted to K experts; the following ``copy_missing`` + streams it over PCIe), then fetch the routed disk-resident experts into + THEIR identity slots. The identity mapping (position == expert id) is + preserved, so the prefill GEMM is unchanged.""" + from freetoken.moe.offload_kernels import _materialize_layer_gpu + + _materialize_layer_gpu(cache, layer_id, materialize_count=self._ram) + routed = expert_ids.reshape(-1) + disk = torch.unique(routed[routed >= self._ram]) + if disk.numel() == 0: + return + step = int(cache.step.item()) # already incremented by the kernel + futures = [ + self._pool.submit(self._fetch_expert, layer_id, int(e), int(e)) + for e in disk.tolist() + ] + for f in futures: + f.result() + # Same bookkeeping the materialize kernel writes, per fetched expert. + flat = layer_id * cache.num_experts + disk + cache.slot_for_id[layer_id, disk] = disk + cache.id_of_slot[disk] = flat + cache.usage[disk] = step + + def fetch_pending(self, cache, layer_id: int) -> None: + """Fetch this layer's disk-resident misses into their slots; shrink the miss + list to the RAM-resident remainder for the existing PCIe copy path.""" + n = int(cache.num_indices.item()) + if n == 0: + return + src = cache.src_indices[:n].cpu() + slots = cache.evict_slots[:n].cpu() + disk = [i for i in range(n) if int(src[i]) >= self._ram] + if not disk: + return + futures = [ + self._pool.submit(self._fetch_expert, layer_id, int(src[i]), int(slots[i])) + for i in disk + ] + for f in futures: + f.result() + disk_set = set(disk) + ram = [i for i in range(n) if i not in disk_set] + if ram: + sel = torch.tensor(ram, dtype=torch.long) + cache.src_indices[:len(ram)].copy_(src[sel].to(cache.src_indices.dtype)) + cache.evict_slots[:len(ram)].copy_(slots[sel].to(cache.evict_slots.dtype)) + cache.num_indices.fill_(len(ram)) + + def refresh(self, cache) -> None: + """Rebind the slot-cache references after a runtime cache rebuild.""" + self._banks = list(cache.banks) + + def stats(self) -> dict: + return {"experts_fetched": self._fetches, "bytes_fetched": self._fetch_bytes} diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba8..ba99f45b8 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -49,6 +49,11 @@ class ExpertBanks: # streamed straight to its sink instead of staying materialized here) -- set by # convert.py's per-format streaming gate; ``sources`` may hold released tensors. streamed: bool = False + # Disk tier (None when off): a moe.disk_tier.Nvfp4DiskIndex over the original + # checkpoint plus how many experts per layer are RAM-resident (the rest are + # disk-resident and fetched on slot-cache miss). + disk_index: object | None = field(default=None) + disk_ram_experts: int = 0 _PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader @@ -160,7 +165,7 @@ def assemble(self, num_layers: int): return sources, gate_up_alpha, down_alpha -def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None, disk_tier=None) -> ExpertBanks: from freetoken.models.weight import load_nvfp4_moe_expert_sources from freetoken.moe.nvfp4_backends import ( b12x_repack_layer, @@ -183,6 +188,10 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, getattr(model_config, "nvfp4_backend", "auto"), activation=getattr(model_config, "hidden_act", "silu")) native = backend == "triton" + if disk_tier is not None and (not native or dummy): + raise NotImplementedError( + "disk tier requires the native NVFP4 layout (triton backend, decode_target=gpu, " + "not dummy)") repack_sink = None if not native and not dummy and layer_sink is not None: @@ -197,7 +206,7 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, # parallel only parallelizes the source read; the backend repack below is reused unchanged. sources = load_nvfp4_moe_expert_sources( model_path, model_config, dummy=dummy, parallel=parallel, workers=workers, chunk=chunk, - layer_sink=sink, + layer_sink=sink, disk_tier=disk_tier, ) # CPU-compute decode (cpu/hybrid) reads the native ModelOpt rows directly (its # dequant-in-GEMV kernel), so keep the native "nvfp4" layout and skip the GPU-tiled @@ -304,13 +313,14 @@ def _model_setup_override(model_config): } -def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None, disk_tier=None) -> ExpertBanks: """Dispatch to the model's setup-override or the per-quant provider. ``parallel=True`` is the parallel read; a provider that hasn't implemented it raises NotImplementedError (the caller falls back to serial). ``decode_target`` lets the cpu backend force CPU-readable (native, non-GPU-tiled) bank layouts. ``layer_sink`` (converter only) is forwarded to setups/providers that declare the parameter; the rest ignore it and stay on the - materialize-and-write path (``ExpertBanks.streamed`` reports which happened).""" + materialize-and-write path (``ExpertBanks.streamed`` reports which happened). + ``disk_tier`` (a ``moe.disk_tier.DiskTierSpec``) is forwarded the same way.""" setup = _model_setup_override(model_config) if setup is not None: import inspect @@ -330,6 +340,8 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel kw["decode_target"] = decode_target if "layer_sink" in params and layer_sink is not None: kw["layer_sink"] = layer_sink + if "disk_tier" in params and disk_tier is not None: + kw["disk_tier"] = disk_tier return setup(model_path, model_config, **kw) expert_quant = model_config.expert_quant @@ -341,7 +353,7 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel return _PROVIDERS[expert_quant]( model_path, model_config, device, dtype, dummy, parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target, - layer_sink=layer_sink, + layer_sink=layer_sink, disk_tier=disk_tier, ) @@ -418,6 +430,7 @@ def load_expert_banks( decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, + disk_tier=None, ) -> ExpertBanks: """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning the same normalized ``ExpertBanks`` and both pinning after fill: @@ -442,6 +455,10 @@ def load_expert_banks( from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks if model_path and is_ftw_checkpoint(model_path) and not dummy: + if disk_tier is not None: + raise NotImplementedError( + "disk tier v0 reads the original safetensors checkpoint; FTW checkpoints " + "are not supported yet (serve from the source path)") banks = load_ftw_banks( model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, layer_residency=layer_residency, @@ -483,13 +500,13 @@ def load_expert_banks( 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) + decode_target, layer_sink, disk_tier) 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) + decode_target, layer_sink, disk_tier) return _echo_residency(banks, layer_residency, residency_plan) diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 436f52d2d..2c25b8c75 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -141,6 +141,39 @@ def pin(self) -> None: ) from exc self._pinned = True + def pin_prefix(self, nrows: int) -> None: + """Pin only the first ``nrows`` rows (disk tier: the rest stays disk-resident). + + The unpinned tail keeps its filled pages until :meth:`release_range` drops + them; nothing may DMA from the tail (the disk tier's miss filter guarantees + the GPU never copies those rows).""" + if self._pinned: + return + from freetoken.kernel.pinned import host_register + + row_bytes = self.nbytes // self.tensor.shape[0] + nbytes = nrows * row_bytes + try: + host_register(self.addr, nbytes) + except RuntimeError as exc: + raise RuntimeError( + f"cudaHostRegister failed for {nbytes / 2**30:.1f} GiB prefix" + ) from exc + self._pinned = True + + def release_range(self, offset: int, nbytes: int) -> None: + """MADV_DONTNEED a byte range of the backing mmap (frees the resident pages). + + For the disk tier's unpinned bank tails: the address space stays valid but + the pages are gone, so any read of the range silently refaults as zeros -- + callers must guarantee nothing reads it (see :meth:`pin_prefix`).""" + import ctypes + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + MADV_DONTNEED = 4 + rc = libc.madvise(self.addr + offset, nbytes, MADV_DONTNEED) + if rc != 0: + raise OSError(ctypes.get_errno(), "madvise(MADV_DONTNEED) failed") def release(self) -> None: """Drop the resident pages; the address space stays valid, the contents become undefined. @@ -289,7 +322,8 @@ class PinPipeline: A clean context-manager exit drains the queue and re-raises the first settle failure. """ - def __init__(self) -> None: + def __init__(self, prefix_rows: int | None = None) -> None: + self._prefix_rows = prefix_rows self._q: queue.SimpleQueue = queue.SimpleQueue() self._exc: BaseException | None = None # the current device is thread-local: a fresh thread sits on device 0 and cudaHostRegister would build its context there -- carry the creator's (bound) device into the worker @@ -308,9 +342,16 @@ def _run(self) -> None: continue # drain without settling after a failure 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) + if self._prefix_rows is not None: + # Disk tier: pin only the RAM-resident expert prefix; the + # disk-resident tail is released by the caller. Takes + # precedence over the residency label (H2D needs the + # prefix page-locked regardless). + bank.pin_prefix(self._prefix_rows) + else: + _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 diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index e1f20dd2f..934822934 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -154,6 +154,9 @@ def __post_init__(self) -> None: # offload/PCIe path. Set by the engine after construction (empty = all-GPU, # all layers = the plain --moe-backend cpu case). self.cpu_layer_ids: frozenset = frozenset() + # Disk tier (None when off): a moe.disk_tier.DiskTier that fetches + # disk-resident slot-cache misses before the PCIe copy path. + self._disk_tier = None # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path. self.validate_rebuild(self.cache_size) assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, ( @@ -509,6 +512,8 @@ def rebuild(self, cache_size: int) -> None: self.prefill_overlap = False if self.prefill_overlap: self._init_prefill_overlap_buffers() + if self._disk_tier is not None: + self._disk_tier.refresh(self) # slot caches were reallocated def set_alphas( self, gate_up_alpha: torch.Tensor | None, down_alpha: torch.Tensor | None @@ -850,7 +855,18 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction ) - def materialize_layer(self, layer_id: int) -> None: + @property + def disk_tier_enabled(self) -> bool: + return self._disk_tier is not None + + def materialize_layer(self, layer_id: int, expert_ids: torch.Tensor | None = None) -> None: + if self._disk_tier is not None: + # Disk tier: stream only the RAM-resident prefix, then fetch the routed + # disk-resident experts into their identity slots (needs the routing). + assert expert_ids is not None, "disk-tier prefill needs the routed expert ids" + self._pending_src_layer = layer_id + self._disk_tier.materialize_layer(self, layer_id, expert_ids) + return from freetoken.moe.offload_kernels import materialize_layer self._pending_src_layer = layer_id @@ -985,11 +1001,25 @@ def decode_routing_stats(self) -> dict: "norm_entropy": norm_ent, } + def attach_disk_tier(self, index, ram_experts: int, workers: int = 8) -> None: + """Enable the NVMe tier: disk-resident slot-cache misses are fetched from the + original checkpoint before the PCIe copy path (see moe/disk_tier.py).""" + from freetoken.moe.disk_tier import DiskTier + + assert self.decode_target == "gpu", "disk tier v0 supports the gpu (offload) path only" + assert self.quant_format == "nvfp4", f"disk tier v0 supports native nvfp4 banks (got {self.quant_format!r})" + assert not self.prefill_overlap, "disk tier v0 does not support prefill overlap" + self._disk_tier = DiskTier(index, self, ram_experts, workers=workers) + 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 self._disk_tier is not None: + # Fetch this layer's disk-resident misses into their slots, then shrink the + # miss list to the RAM-resident remainder for the PCIe copy below. + self._disk_tier.fetch_pending(self, layer_id) + elif 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 " diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d..3b6d67fb7 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -188,7 +188,10 @@ def _ensure_experts_hybrid_cpu( flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item()) -def _materialize_layer_gpu(cache, layer_id: int) -> None: +def _materialize_layer_gpu(cache, layer_id: int, materialize_count: int | None = None) -> None: + # materialize_count < num_experts: the disk tier's RAM-resident prefix only; the + # flat-id base still uses the full num_experts (the id space is layer * E + expert). + count = cache.num_experts if materialize_count is None else materialize_count block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size)) _materialize_layer_kernel[(1,)]( cache.slot_for_id, @@ -200,6 +203,7 @@ def _materialize_layer_gpu(cache, layer_id: int) -> None: cache.num_indices, layer_id, cache.num_experts, + count, cache.cache_size, BLOCK=block, ) @@ -257,11 +261,12 @@ def _materialize_layer_kernel( num_indices_ptr, layer_id: tl.constexpr, num_experts: tl.constexpr, + materialize_count: tl.constexpr, cache_size: tl.constexpr, BLOCK: tl.constexpr, ): off = tl.arange(0, BLOCK) - expert_mask = off < num_experts + expert_mask = off < materialize_count slot_mask = off < cache_size slot = off @@ -282,7 +287,7 @@ def _materialize_layer_kernel( tl.store(usage_ptr + slot, step, mask=expert_mask) tl.store(evict_slots_ptr + off, slot, mask=expert_mask) tl.store(src_indices_ptr + off, off, mask=expert_mask) # layer-local row - tl.store(num_indices_ptr, num_experts) + tl.store(num_indices_ptr, materialize_count) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 5b4db587d..375e49ea3 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -547,6 +547,33 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The unified MoE cache eviction policy.", ) + parser.add_argument( + "--moe-disk-tier", + default=ServerArgs.moe_disk_tier, + choices=["off", "on"], + help=( + "NVMe tier for MoE experts (see moe/disk_tier.py): experts beyond " + "--expert-ram-experts per layer stay on disk in the original checkpoint " + "and are fetched on slot-cache miss. Requires native NVFP4 banks, " + "--moe-backend offload, --disable-moe-prefill-overlap and no cuda graphs." + ), + ) + parser.add_argument( + "--expert-ram-experts", + type=int, + default=ServerArgs.expert_ram_experts, + help=( + "With --moe-disk-tier on: experts per layer kept pinned in RAM " + "(0 < N < num_experts; the rest are disk-resident)." + ), + ) + parser.add_argument( + "--disk-fetch-workers", + type=int, + default=ServerArgs.disk_fetch_workers, + help="Disk-tier O_DIRECT fetch threads (default 8).", + ) + parser.add_argument( "--moe-cpu-threads", type=int, diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py new file mode 100644 index 000000000..4dd0ff1d1 --- /dev/null +++ b/tests/moe/test_disk_tier.py @@ -0,0 +1,246 @@ +"""CPU tests for the NVMe disk tier (moe/disk_tier.py). + +A synthetic NVFP4 MoE checkpoint (2 layers, 4 experts) is written as safetensors +shards with deterministic per-tensor content; the tests verify that +:class:`Nvfp4DiskIndex` resolves the right byte ranges and that +:class:`DiskTier` places the right bytes into slot-cache rows and rewrites the +miss list. No CUDA needed -- the "GPU" banks here are CPU tensors and the +staging pin is stubbed. +""" + +import json +import re +import struct +import types + +import pytest +import torch + +from freetoken.moe.disk_tier import DiskTier, Nvfp4DiskIndex +from freetoken.moe.host_banks import HostBank +from freetoken.models.nvfp4_banks import Nvfp4ExpertSourceSpec + +H, I, E, L = 16, 32, 4, 2 +SHARDS = ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors") + +SPEC = Nvfp4ExpertSourceSpec( + key_pattern=re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight|weight_scale|weight_scale_2)$" + ), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, + desc="disk-tier test", +) + +# (proj, kind, shape, dtype) -- the native NVFP4 per-expert tensor layout. +TENSOR_SPECS = ( + ("gate_proj", "weight", (I, H // 2), torch.uint8), + ("up_proj", "weight", (I, H // 2), torch.uint8), + ("down_proj", "weight", (H, I // 2), torch.uint8), + ("gate_proj", "weight_scale", (I, H // 16), torch.uint8), + ("up_proj", "weight_scale", (I, H // 16), torch.uint8), + ("down_proj", "weight_scale", (H, I // 16), torch.uint8), + ("gate_proj", "weight_scale_2", (I,), torch.float16), + ("up_proj", "weight_scale_2", (I,), torch.float16), + ("down_proj", "weight_scale_2", (H,), torch.float16), +) + +BANK_SHAPES = ( + (2 * I, H // 2), # gate_up_packed + (2 * I, H // 16), # gate_up_scale + (2 * I,), # gate_up_global + (H, I // 2), # down_packed + (H, I // 16), # down_scale + (H,), # down_global +) +BANK_DTYPES = (torch.uint8, torch.uint8, torch.float16, torch.uint8, torch.uint8, torch.float16) + + +def _name(layer, expert, proj, kind): + return f"model.language_model.layers.{layer}.mlp.experts.{expert}.{proj}.{kind}" + + +def _tensor_for(layer, expert, proj, kind): + """Deterministic content: a base offset per (layer, expert, proj, kind) so any + misplacement is visible.""" + proj_i = ("gate_proj", "up_proj", "down_proj").index(proj) + kind_i = ("weight", "weight_scale", "weight_scale_2").index(kind) + base = layer * 100000 + expert * 1000 + proj_i * 100 + kind_i * 10 + for p, k, shape, dtype in TENSOR_SPECS: + if p == proj and k == kind: + n = int(torch.tensor(shape).prod()) + if dtype == torch.uint8: + return torch.arange(n, dtype=torch.uint8).add_(base % 251).view(shape) + return (torch.arange(n, dtype=torch.float32) + base).to(dtype).view(shape) + raise AssertionError((proj, kind)) + + +@pytest.fixture() +def checkpoint(tmp_path): + import safetensors.torch + + by_shard = {s: {} for s in SHARDS} + weight_map = {} + for layer in range(L): + for expert in range(E): + for proj, kind, _shape, _dtype in TENSOR_SPECS: + shard = SHARDS[(layer * E + expert) % 2] + name = _name(layer, expert, proj, kind) + by_shard[shard][name] = _tensor_for(layer, expert, proj, kind) + weight_map[name] = shard + for shard, tensors in by_shard.items(): + safetensors.torch.save_file(tensors, str(tmp_path / shard), metadata={"format": "pt"}) + with open(tmp_path / "model.safetensors.index.json", "w", encoding="utf-8") as f: + json.dump({"weight_map": weight_map, "metadata": None}, f) + config = types.SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, + num_layers=L, first_k_dense_replace=0) + return tmp_path, config + + +def _index(checkpoint): + path, config = checkpoint + return Nvfp4DiskIndex(str(path), config, SPEC) + + +def test_index_segments_match_file_bytes(checkpoint): + import safetensors + + path, config = checkpoint + index = _index(checkpoint) + assert len(index.shard_paths) == 2 + # Every (bank, layer, expert) row segment must point at the exact tensor bytes. + for bank_idx in range(6): + for layer in range(L): + for expert in range(E): + segs = index.row_segments(bank_idx, layer, expert) + expected = { + 0: [("gate_proj", "weight"), ("up_proj", "weight")], + 1: [("gate_proj", "weight_scale"), ("up_proj", "weight_scale")], + 2: [("gate_proj", "weight_scale_2"), ("up_proj", "weight_scale_2")], + 3: [("down_proj", "weight")], + 4: [("down_proj", "weight_scale")], + 5: [("down_proj", "weight_scale_2")], + }[bank_idx] + assert len(segs) == len(expected) + for (shard_idx, off, nbytes), (proj, kind) in zip(segs, expected): + shard_path = index.shard_paths[shard_idx] + with open(shard_path, "rb") as f: + (hlen,) = struct.unpack(" slots 5, 6, 7. + cache.src_indices[:3] = torch.tensor([0, 2, 3], dtype=torch.int32) + cache.evict_slots[:3] = torch.tensor([5, 6, 7], dtype=torch.int32) + cache.num_indices.fill_(3) + + tier.fetch_pending(cache, layer) + + # Disk misses fetched into their slots... + expected = _expected_rows(layer, 2) + for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): + assert torch.equal( + gpu_cache[6].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) + expected = _expected_rows(layer, 3) + for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): + assert torch.equal( + gpu_cache[7].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) + # ...and the miss list shrank to the RAM-resident remainder. + assert cache.num_indices.item() == 1 + assert cache.src_indices[0].item() == 0 + assert cache.evict_slots[0].item() == 5 + + +def test_fetch_pending_all_ram_is_noop(checkpoint): + cache = _fake_cache() + tier = _tier(checkpoint, cache, ram_experts=2) + cache.src_indices[:2] = torch.tensor([0, 1], dtype=torch.int32) + cache.evict_slots[:2] = torch.tensor([0, 1], dtype=torch.int32) + cache.num_indices.fill_(2) + tier.fetch_pending(cache, 0) + assert cache.num_indices.item() == 2 + assert tier.stats()["experts_fetched"] == 0 + + +def test_fetch_pending_all_disk_clears_list(checkpoint): + cache = _fake_cache() + tier = _tier(checkpoint, cache, ram_experts=2) + cache.src_indices[:1] = torch.tensor([3], dtype=torch.int32) + cache.evict_slots[:1] = torch.tensor([4], dtype=torch.int32) + cache.num_indices.fill_(1) + tier.fetch_pending(cache, 0) + assert cache.num_indices.item() == 0 + expected = _expected_rows(0, 3) + for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): + assert torch.equal( + gpu_cache[4].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) From ae5f736982c286d906f718b3fd98af5c11a7c0b5 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:41:29 +0000 Subject: [PATCH 02/38] fix(moe): disk tier H2D must view staging bytes as the dst dtype --- python/freetoken/moe/disk_tier.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 3aa315a9c..77b59f89b 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -224,7 +224,7 @@ def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: row_off = off - a0 src = staging.tensor[row_off:row_off + nbytes] dst = row[d0:d1] - dst.copy_(src.view(dst.shape), non_blocking=True) + dst.copy_(src.view(dst.dtype).view(dst.shape), non_blocking=True) self._fetches += 1 self._fetch_bytes += sum(self._row_bytes) From 5be9f88111e258232eb594a584d57b034f6ed137 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:42:17 +0000 Subject: [PATCH 03/38] fix(moe): disk tier offsets must be absolute (header + data section base) --- python/freetoken/moe/disk_tier.py | 10 ++++++++-- tests/moe/test_disk_tier.py | 3 ++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 77b59f89b..8797f2a34 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -73,11 +73,17 @@ def release_bank_tails(banks_by_name: dict[str, list[HostBank]], num_experts: in def _read_safetensors_offsets(path: str) -> dict[str, tuple[int, int]]: - """{tensor_name: (start, end)} from a shard's safetensors header.""" + """{tensor_name: (start, end)} from a shard's safetensors header, as ABSOLUTE + file offsets (data_offsets are relative to the data section, i.e. after the + 8-byte length + header JSON).""" with open(path, "rb") as f: (hlen,) = struct.unpack(" Date: Mon, 31 Aug 2026 21:43:12 +0000 Subject: [PATCH 04/38] test(moe): disk tier test staging must be per-thread (mirror production) --- tests/moe/test_disk_tier.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index b718ea027..b97f11b60 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -11,6 +11,7 @@ import json import re import struct +import threading import types import pytest @@ -160,9 +161,18 @@ def _fake_cache(): def _tier(checkpoint, cache, ram_experts=2): index = _index(checkpoint) tier = DiskTier(index, cache, ram_experts=ram_experts, workers=2) - # Stub the pinned staging (HostBank.pin needs CUDA); the mmap buffer itself is fine. - staging = HostBank((tier._staging_size,), torch.uint8) - tier._staging_buf = lambda: staging + # Stub the pinned staging (HostBank.pin needs CUDA) with the same per-thread + # buffer semantics as production (threading.local). + local = threading.local() + + def _staging_buf(): + buf = getattr(local, "buf", None) + if buf is None: + buf = HostBank((tier._staging_size,), torch.uint8) + local.buf = buf + return buf + + tier._staging_buf = _staging_buf return tier From 01c068741871b8245e2310d9987332104e08bfca Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:01:19 +0000 Subject: [PATCH 05/38] fix(moe): disk tier guard must require --cuda-graph-max-bs 0 (None auto-enables graphs) --- python/freetoken/engine/engine.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 266022beb..f06fe6b1c 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -578,8 +578,9 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: "--moe-disk-tier v0 requires the gpu decode path (--moe-backend offload)") if config.moe_prefill_overlap: raise ValueError("--moe-disk-tier v0 requires --disable-moe-prefill-overlap") - if config.cuda_graph_max_bs is not None: - raise ValueError("--moe-disk-tier v0 requires cuda graphs disabled") + if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs >= 1: + raise ValueError( + "--moe-disk-tier v0 requires --cuda-graph-max-bs 0 (cuda graphs disabled)") disk_tier = DiskTierSpec(ram_experts=config.expert_ram_experts) if cache_factory is None: # Fast path: an FTW checkpoint loads its repacked banks directly. From d7235fe59e8387f668c18d8631cf7be42675fdf5 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:10:13 +0000 Subject: [PATCH 06/38] fix(moe): madvise needs ctypes argtypes (64-bit addr truncation) + regression test --- python/freetoken/moe/host_banks.py | 4 ++++ tests/moe/test_disk_tier.py | 25 +++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 2c25b8c75..521ac423d 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -170,6 +170,10 @@ def release_range(self, offset: int, nbytes: int) -> None: import ctypes libc = ctypes.CDLL("libc.so.6", use_errno=True) + # argtypes are mandatory: without them ctypes truncates the 64-bit + # address to a C int and madvise fails (ENOMEM/EINVAL on bogus addrs). + libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + libc.madvise.restype = ctypes.c_int MADV_DONTNEED = 4 rc = libc.madvise(self.addr + offset, nbytes, MADV_DONTNEED) if rc != 0: diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index b97f11b60..342398ff0 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -255,3 +255,28 @@ def test_fetch_pending_all_disk_clears_list(checkpoint): for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): assert torch.equal( gpu_cache[4].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) + + +def test_release_range_frees_pages(): + """madvise(MADV_DONTNEED) must actually drop the resident pages (and not fail + on the 64-bit address -- the no-argtypes ctypes truncation bug).""" + import ctypes as ct + + size = 4 * 1024 * 1024 + bank = HostBank((size,), torch.uint8) + bank.tensor.fill_(7) # fault every page in + libc = ct.CDLL("libc.so.6", use_errno=True) + libc.mincore.argtypes = [ct.c_void_p, ct.c_size_t, ct.POINTER(ct.c_ubyte)] + libc.mincore.restype = ct.c_int + + def resident_pages(addr, nbytes): + vec = (ct.c_ubyte * ((nbytes + 4095) // 4096))() + assert libc.mincore(addr, nbytes, vec) == 0 + return sum(1 for b in vec if b & 1) + + pages = size // 4096 + assert resident_pages(bank.addr, size) == pages + bank.release_range(0, size) + assert resident_pages(bank.addr, size) == 0 + # The mapping stays valid: the refaulted pages read back as zeros. + assert bank.tensor[0] == 0 From e166a532f4b684cf528dbea1cf74fd406ca16f3b Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:13:22 +0000 Subject: [PATCH 07/38] fix(moe): release_range must remap MAP_PRIVATE in place (shared /dev/zero ignores DONTNEED); index offsets need 64-bit --- python/freetoken/moe/disk_tier.py | 6 ++--- python/freetoken/moe/host_banks.py | 42 +++++++++++++++++++++--------- tests/moe/test_disk_tier.py | 5 ++-- 3 files changed, 35 insertions(+), 18 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 8797f2a34..12777bf64 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -123,7 +123,7 @@ def __init__(self, model_dir: str, config, spec) -> None: shard_idx = {s: i for i, s in enumerate(shards)} E = config.num_experts - seg_size = struct.calcsize(" packed segments per expert for bank_idx in range(len(_NVP4_BANK_SEGS)): per_layer = [] @@ -140,7 +140,7 @@ def __init__(self, model_dir: str, config, spec) -> None: ) name, shard = entry start, end = offsets[shard][name] - rows += struct.pack(" list[tuple[int base = expert * self._seg_size * len(_NVP4_BANK_SEGS[bank_idx]) raw = self.entries[bank_idx][layer][base:base + self._seg_size * len(_NVP4_BANK_SEGS[bank_idx])] return [ - struct.unpack_from(" None: self._pinned = True def release_range(self, offset: int, nbytes: int) -> None: - """MADV_DONTNEED a byte range of the backing mmap (frees the resident pages). - - For the disk tier's unpinned bank tails: the address space stays valid but - the pages are gone, so any read of the range silently refaults as zeros -- - callers must guarantee nothing reads it (see :meth:`pin_prefix`).""" + """Free a byte range of the backing mapping by replacing it IN PLACE with a + fresh MAP_PRIVATE anonymous mapping at the same virtual address. + + HostBank's buffer is a MAP_SHARED /dev/zero mapping (CPython's + ``mmap(-1)``), and the kernel silently ignores MADV_DONTNEED on shared + mappings -- the pages would stay resident. Replacing the range with a + private zero mapping frees them while keeping every existing pointer + and torch view valid (same address). The range must be page-aligned + and must not overlap a pinned prefix (the disk tier's unpinned tails). + """ import ctypes + _BLK = 4096 + assert offset % _BLK == 0 and nbytes % _BLK == 0, ( + "release_range: page-aligned range required") libc = ctypes.CDLL("libc.so.6", use_errno=True) - # argtypes are mandatory: without them ctypes truncates the 64-bit - # address to a C int and madvise fails (ENOMEM/EINVAL on bogus addrs). - libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] - libc.madvise.restype = ctypes.c_int - MADV_DONTNEED = 4 - rc = libc.madvise(self.addr + offset, nbytes, MADV_DONTNEED) - if rc != 0: - raise OSError(ctypes.get_errno(), "madvise(MADV_DONTNEED) failed") + libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + libc.munmap.restype = ctypes.c_int + libc.mmap.restype = ctypes.c_void_p + libc.mmap.argtypes = [ + ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long] + addr = self.addr + offset + if libc.munmap(addr, nbytes) != 0: + raise OSError(ctypes.get_errno(), "munmap failed") + PROT_READ_WRITE = 3 + MAP_PRIVATE_ANON = 0x22 # MAP_PRIVATE | MAP_ANONYMOUS + MAP_FIXED = 0x10 + MAP_FAILED = (1 << 64) - 1 + new_addr = libc.mmap(addr, nbytes, PROT_READ_WRITE, MAP_PRIVATE_ANON | MAP_FIXED, -1, 0) + if new_addr in (None, MAP_FAILED): + raise OSError(ctypes.get_errno(), "mmap(MAP_FIXED) failed") + assert new_addr == addr, "MAP_FIXED returned a different address" def release(self) -> None: """Drop the resident pages; the address space stays valid, the contents become undefined. diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index 342398ff0..0fdd8558a 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -258,8 +258,9 @@ def test_fetch_pending_all_disk_clears_list(checkpoint): def test_release_range_frees_pages(): - """madvise(MADV_DONTNEED) must actually drop the resident pages (and not fail - on the 64-bit address -- the no-argtypes ctypes truncation bug).""" + """release_range must actually drop the resident pages (MAP_SHARED /dev/zero + mappings ignore MADV_DONTNEED, so the in-place private remap is verified + with mincore).""" import ctypes as ct size = 4 * 1024 * 1024 From 2b33ae761a2f9a99c2612ced4ad5aace3ae4bf93 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:14:45 +0000 Subject: [PATCH 08/38] fix(moe): ExpertBanks is frozen -- use dataclasses.replace for the disk index --- python/freetoken/models/qwen3_5_moe/weight.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 3818e9aac..e5cd98411 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -867,10 +867,13 @@ def setup_offload_expert_banks( if eq != "nvfp4": raise NotImplementedError( f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") + import dataclasses + from freetoken.moe.disk_tier import Nvfp4DiskIndex - banks.disk_index = Nvfp4DiskIndex(model_path, model_config, _NVFP4_SOURCE_SPEC) - banks.disk_ram_experts = disk_tier.ram_experts + index = Nvfp4DiskIndex(model_path, model_config, _NVFP4_SOURCE_SPEC) + banks = dataclasses.replace( + banks, disk_index=index, disk_ram_experts=disk_tier.ram_experts) return banks if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only") From 97c7c2c8a5d782149bcf6e65a08e4ec564e2eaaf Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:49:26 +0000 Subject: [PATCH 09/38] chore(moe): disk tier preadv failure diagnostics --- python/freetoken/moe/disk_tier.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 12777bf64..ba9b28a62 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -226,7 +226,24 @@ def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: else: a0, slen = off, nbytes mv = (ctypes.c_char * slen).from_address(staging.addr) - os.preadv(fd, [mv], a0) + try: + os.preadv(fd, [mv], a0) + except OSError: + vma = "?" + try: + for line in open("/proc/self/maps"): + lo, hi = line.split()[0].split("-") + if int(lo, 16) <= staging.addr < int(hi, 16): + vma = line.strip()[:120] + break + except OSError: + pass + raise OSError( + f"disk-tier preadv failed: shard={shard_idx} off={off} a0={a0} " + f"slen={slen} direct={direct} buf={hex(staging.addr)} " + f"staging_size={self._staging_size} thread={threading.current_thread().name} " + f"vma={vma}" + ) from None row_off = off - a0 src = staging.tensor[row_off:row_off + nbytes] dst = row[d0:d1] From 4f561cdaf24cfbefe909b7ea9b5a60c45a0cea8b Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:54:10 +0000 Subject: [PATCH 10/38] fix(moe): disk tier staging size must use full row bytes (was dropping the row's leading dim -> 8KB buffer, EFAULT/segfault on 512KB segments) --- python/freetoken/moe/disk_tier.py | 5 ++--- tests/moe/test_disk_tier.py | 5 +++++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index ba9b28a62..a3c1605f9 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -25,7 +25,6 @@ import ctypes import json -import math import os import struct import threading @@ -163,8 +162,8 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int self._ram = ram_experts self._banks = list(cache.banks) # [(per_layer_host, gpu_cache)] in schema order self._row_bytes = [ - math.prod(b[0][0][0].shape[1:]) * b[0][0][0].element_size() for b in self._banks - ] + b[0][0][0].numel() * b[0][0][0].element_size() for b in self._banks + ] # full expert-row bytes per bank (staging must hold the biggest one) # Per-bank destination row slices (gate|up split at the row midpoint). self._dst_slices: list[list[tuple[int, int]]] = [] for bank_idx, (host_layer, _gpu) in enumerate(self._banks): diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index 0fdd8558a..166a128f7 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -161,6 +161,11 @@ def _fake_cache(): def _tier(checkpoint, cache, ram_experts=2): index = _index(checkpoint) tier = DiskTier(index, cache, ram_experts=ram_experts, workers=2) + # Staging must hold the largest bank row's O_DIRECT super-block (a size + # regression here overflows the buffer -> EFAULT/segfault at fetch time). + max_row = max( + b[0][0][0].numel() * b[0][0][0].element_size() for b in cache.banks) + assert tier._staging_size >= max_row + 2 * 4096, (tier._staging_size, max_row) # Stub the pinned staging (HostBank.pin needs CUDA) with the same per-thread # buffer semantics as production (threading.local). local = threading.local() From f27d126b328178fd6b37336a12a36ba0c6632491 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:55:18 +0000 Subject: [PATCH 11/38] fix(moe): staging size ceiling rounding (was truncating sub-page rows to 8KB) --- python/freetoken/moe/disk_tier.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index a3c1605f9..52a124a23 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -174,7 +174,7 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int else: self._dst_slices.append([(0, row.shape[0])]) max_row = max(self._row_bytes) - self._staging_size = ((max_row + 2 * _ALIGN) // _ALIGN) * _ALIGN + self._staging_size = ((max_row + _ALIGN - 1) // _ALIGN + 2) * _ALIGN self._staging = threading.local() self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="disk-tier") self._fd_lock = threading.Lock() From de070c342b2263faa2f49e98d639c24e290fb7d0 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 00:26:08 +0000 Subject: [PATCH 12/38] fix(moe): disk tier fetch pool threads must run under inference_mode (server tensors) --- python/freetoken/moe/disk_tier.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 52a124a23..e6e6e3e39 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -213,6 +213,12 @@ def _staging_buf(self) -> HostBank: # ---------------------------------------------------------------- fetch def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: + # The server runs under inference_mode; the fetch pool threads do not, + # so the H2D writes into the (inference) slot cache need their own scope. + with torch.inference_mode(): + self._fetch_expert_inner(layer, expert, slot) + + def _fetch_expert_inner(self, layer: int, expert: int, slot: int) -> None: staging = self._staging_buf() for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): row = gpu_cache[slot] From 97e45ddf35aa3e198231a2bf8b111df6c07616a9 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 00:59:22 +0000 Subject: [PATCH 13/38] fix(moe): global-scale banks store per-expert fp32 scalars -- convert+fill, not byte copy; test matches real checkpoint layout --- python/freetoken/moe/disk_tier.py | 44 ++++++++++++++++++++----------- tests/moe/test_disk_tier.py | 27 +++++++++++++------ 2 files changed, 48 insertions(+), 23 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index e6e6e3e39..dbfaed5fc 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -71,6 +71,26 @@ def release_bank_tails(banks_by_name: dict[str, list[HostBank]], num_experts: in ) + +def _preadv_error(tier, staging, shard_idx: int, off: int, a0: int, slen: int, + direct: bool) -> OSError: + vma = "?" + try: + for line in open("/proc/self/maps"): + lo, hi = line.split()[0].split("-") + if int(lo, 16) <= staging.addr < int(hi, 16): + vma = line.strip()[:120] + break + except OSError: + pass + return OSError( + f"disk-tier preadv failed: shard={shard_idx} off={off} a0={a0} " + f"slen={slen} direct={direct} buf={hex(staging.addr)} " + f"staging_size={tier._staging_size} thread={threading.current_thread().name} " + f"vma={vma}" + ) + + def _read_safetensors_offsets(path: str) -> dict[str, tuple[int, int]]: """{tensor_name: (start, end)} from a shard's safetensors header, as ABSOLUTE file offsets (data_offsets are relative to the data section, i.e. after the @@ -234,22 +254,16 @@ def _fetch_expert_inner(self, layer: int, expert: int, slot: int) -> None: try: os.preadv(fd, [mv], a0) except OSError: - vma = "?" - try: - for line in open("/proc/self/maps"): - lo, hi = line.split()[0].split("-") - if int(lo, 16) <= staging.addr < int(hi, 16): - vma = line.strip()[:120] - break - except OSError: - pass - raise OSError( - f"disk-tier preadv failed: shard={shard_idx} off={off} a0={a0} " - f"slen={slen} direct={direct} buf={hex(staging.addr)} " - f"staging_size={self._staging_size} thread={threading.current_thread().name} " - f"vma={vma}" - ) from None + raise _preadv_error(self, staging, shard_idx, off, a0, slen, direct) row_off = off - a0 + if bank_idx in (2, 5): + # Global-scale banks: the checkpoint stores a per-expert fp32 + # SCALAR (weight_scale_2); the bank row is that value as fp16 + # broadcast across the row -- convert + fill, no byte copy. + val = staging.tensor[row_off:row_off + 4].view(torch.float32)[0].to( + torch.float16) + row[d0:d1].fill_(val) + continue src = staging.tensor[row_off:row_off + nbytes] dst = row[d0:d1] dst.copy_(src.view(dst.dtype).view(dst.shape), non_blocking=True) diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index 166a128f7..c75a1340a 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -43,9 +43,11 @@ ("gate_proj", "weight_scale", (I, H // 16), torch.uint8), ("up_proj", "weight_scale", (I, H // 16), torch.uint8), ("down_proj", "weight_scale", (H, I // 16), torch.uint8), - ("gate_proj", "weight_scale_2", (I,), torch.float16), - ("up_proj", "weight_scale_2", (I,), torch.float16), - ("down_proj", "weight_scale_2", (H,), torch.float16), + # weight_scale_2 is a per-expert fp32 SCALAR in the real checkpoints; the + # bank row is its fp16 value broadcast across the row. + ("gate_proj", "weight_scale_2", (), torch.float32), + ("up_proj", "weight_scale_2", (), torch.float32), + ("down_proj", "weight_scale_2", (), torch.float32), ) BANK_SHAPES = ( @@ -71,6 +73,8 @@ def _tensor_for(layer, expert, proj, kind): base = layer * 100000 + expert * 1000 + proj_i * 100 + kind_i * 10 for p, k, shape, dtype in TENSOR_SPECS: if p == proj and k == kind: + if shape == (): # per-expert fp32 scalar (kept in fp16 range) + return torch.tensor(float(base % 50000), dtype=torch.float32) n = int(torch.tensor(shape).prod()) if dtype == torch.uint8: return torch.arange(n, dtype=torch.uint8).add_(base % 251).view(shape) @@ -185,13 +189,20 @@ def _expected_rows(layer, expert): """The 6 bank rows for one expert, as flat uint8, in schema order.""" rows = [] for bank_idx, (shape, dtype) in enumerate(zip(BANK_SHAPES, BANK_DTYPES)): - if bank_idx < 3: - gate = _tensor_for(layer, expert, "gate_proj", ("weight", "weight_scale", "weight_scale_2")[bank_idx]) - up = _tensor_for(layer, expert, "up_proj", ("weight", "weight_scale", "weight_scale_2")[bank_idx]) + if bank_idx == 2: + gate = _tensor_for(layer, expert, "gate_proj", "weight_scale_2").to(torch.float16) + up = _tensor_for(layer, expert, "up_proj", "weight_scale_2").to(torch.float16) + row = torch.cat([gate.expand(I), up.expand(I)]).view(shape) + elif bank_idx == 5: + row = (_tensor_for(layer, expert, "down_proj", "weight_scale_2") + .to(torch.float16).expand(H).view(shape)) + elif bank_idx < 3: + kind = ("weight", "weight_scale")[bank_idx] + gate = _tensor_for(layer, expert, "gate_proj", kind) + up = _tensor_for(layer, expert, "up_proj", kind) row = torch.cat([gate.reshape(-1), up.reshape(-1)]).view(shape) else: - kind = ("weight", "weight_scale", "weight_scale_2")[bank_idx - 3] - row = _tensor_for(layer, expert, "down_proj", kind) + row = _tensor_for(layer, expert, "down_proj", ("weight", "weight_scale")[bank_idx - 3]) rows.append(row.contiguous().view(torch.uint8).reshape(-1)) return rows From baea729d2636f3aceabb1d60faaab35dd7293783 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:22:32 +0000 Subject: [PATCH 14/38] fix(moe): sync default stream after disk-tier fetches (H2D copies on pool-thread stream must land before GEMM reads slots) --- python/freetoken/moe/disk_tier.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index dbfaed5fc..c54f7d7c2 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -270,6 +270,14 @@ def _fetch_expert_inner(self, layer: int, expert: int, slot: int) -> None: self._fetches += 1 self._fetch_bytes += sum(self._row_bytes) + def _sync_fetches(self) -> None: + """Wait for the pool threads' async H2D copies to land. + + The copies are enqueued on the pool threads' default stream; f.result() only + waits for them to be ENQUEUED. The GEMM's stream is not ordered with that + stream, so sync the default stream before the GEMM reads the slots.""" + torch.cuda.default_stream().synchronize() + def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots (the normal kernel restricted to K experts; the following ``copy_missing`` @@ -290,6 +298,7 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N ] for f in futures: f.result() + self._sync_fetches() # Same bookkeeping the materialize kernel writes, per fetched expert. flat = layer_id * cache.num_experts + disk cache.slot_for_id[layer_id, disk] = disk @@ -313,6 +322,7 @@ def fetch_pending(self, cache, layer_id: int) -> None: ] for f in futures: f.result() + self._sync_fetches() disk_set = set(disk) ram = [i for i in range(n) if i not in disk_set] if ram: From f10a6e6c374d1ade64deedb9476b70ce9fda66d3 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:23:16 +0000 Subject: [PATCH 15/38] fix(moe): guard disk-tier stream sync on cuda availability (CPU-only tests) --- python/freetoken/moe/disk_tier.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index c54f7d7c2..473b0fb25 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -276,6 +276,8 @@ def _sync_fetches(self) -> None: The copies are enqueued on the pool threads' default stream; f.result() only waits for them to be ENQUEUED. The GEMM's stream is not ordered with that stream, so sync the default stream before the GEMM reads the slots.""" + if not torch.cuda.is_available(): + return # CPU-only tests: the copies are synchronous CPU->CPU torch.cuda.default_stream().synchronize() def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: From e16c93d822d8247b4604854f11fe88a9212154cd Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:29:03 +0000 Subject: [PATCH 16/38] chore(moe): disk-tier debug instrumentation (per-layer fetch counts + one-shot slot verify) --- python/freetoken/moe/disk_tier.py | 44 +++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 473b0fb25..83092fedf 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -280,6 +280,45 @@ def _sync_fetches(self) -> None: return # CPU-only tests: the copies are synchronous CPU->CPU torch.cuda.default_stream().synchronize() + def _verify_slot(self, cache, layer: int, expert: int) -> None: + """One-shot debug: read back a fetched disk expert's slot rows and compare + against the checkpoint bytes (ground truth). Gated on FT_DISK_TIER_VERIFY.""" + import torch + slot = expert # identity mapping in prefill + for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): + slot_row = gpu_cache[slot].contiguous() + flat = slot_row.view(torch.uint8).reshape(-1) + segs = self._index.row_segments(bank_idx, layer, expert) + dsl = self._dst_slices[bank_idx] + # Reconstruct the reference row bytes from the checkpoint. + row_bytes = flat.numel() + ref = torch.zeros(row_bytes, dtype=torch.uint8) + row_el = slot_row.element_size() + row_leading = slot_row.numel() // slot_row.shape[0] if slot_row.dim() > 1 else 1 + for (d0, d1), (shard_idx, off, nbytes) in zip(dsl, segs): + fd, direct = self._fd(shard_idx) + a0 = off if not direct else (off & ~(_ALIGN - 1)) + slen = nbytes if not direct else (off + nbytes - a0 + _ALIGN - 1) & ~(_ALIGN - 1) + buf = os.pread(fd, slen, a0) + row_off = off - a0 + seg = buf[row_off:row_off + nbytes] + if bank_idx in (2, 5): + val = int.from_bytes(seg, "little") + import struct as _st + f32 = _st.unpack(" None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots (the normal kernel restricted to K experts; the following ``copy_missing`` @@ -291,6 +330,9 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N _materialize_layer_gpu(cache, layer_id, materialize_count=self._ram) routed = expert_ids.reshape(-1) disk = torch.unique(routed[routed >= self._ram]) + if layer_id < 3: + print(f"[disk-tier dbg] layer={layer_id} routed={routed.numel()} " + f"unique_disk={disk.numel()} disk={disk.tolist()[:12]}", flush=True) if disk.numel() == 0: return step = int(cache.step.item()) # already incremented by the kernel @@ -301,6 +343,8 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N for f in futures: f.result() self._sync_fetches() + if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0 and disk.numel() > 0: + self._verify_slot(cache, 0, int(disk[0].item())) # Same bookkeeping the materialize kernel writes, per fetched expert. flat = layer_id * cache.num_experts + disk cache.slot_for_id[layer_id, disk] = disk From dae3ff2164538fef2434cf9d824f9dbc407f4976 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:59:27 +0000 Subject: [PATCH 17/38] fix(moe): disk-tier verify device mismatch + no-crash guard --- python/freetoken/moe/disk_tier.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 83092fedf..97c023659 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -313,11 +313,15 @@ def _verify_slot(self, cache, layer: int, expert: int) -> None: else: dst_off = d0 * row_leading * row_el ref[dst_off:dst_off + len(seg)] = torch.frombuffer(seg, dtype=torch.uint8) - match = bool(torch.equal(flat, ref)) - print(f"[verify] bank={bank_idx} expert={expert} match={match} " - f"slot_norm={slot_row.float().norm().item():.4f} " - f"ref_norm={ref.view(slot_row.dtype).view(slot_row.shape).float().norm().item():.4f} " - f"slot_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", flush=True) + try: + ref = ref.to(flat.device) + match = bool(torch.equal(flat, ref)) + print(f"[verify] bank={bank_idx} expert={expert} match={match} " + f"slot_norm={slot_row.float().norm().item():.4f} " + f"ref_norm={ref.view(slot_row.dtype).view(slot_row.shape).float().norm().item():.4f} " + f"slot_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", flush=True) + except Exception as exc: # never crash the server in debug + print(f"[verify] bank={bank_idx} expert={expert} ERROR {exc!r}", flush=True) def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots From 3379fdcfd91c456c7cba4a30be28a0b84863792a Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 02:04:03 +0000 Subject: [PATCH 18/38] chore(moe): verify RAM-expert slots after PCIe copy --- python/freetoken/moe/disk_tier.py | 7 +++++++ python/freetoken/moe/offload_cache.py | 24 +++++++++++++----------- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 97c023659..045a20a89 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -323,6 +323,13 @@ def _verify_slot(self, cache, layer: int, expert: int) -> None: except Exception as exc: # never crash the server in debug print(f"[verify] bank={bank_idx} expert={expert} ERROR {exc!r}", flush=True) + def verify_ram(self, cache, layer: int) -> None: + """One-shot debug: after the PCIe copy, check a RAM-resident expert's slot rows + against the checkpoint reference. Gated on FT_DISK_TIER_VERIFY.""" + expert = min(10, self._ram - 1) # a RAM-resident expert + print(f"[verify-ram] layer={layer} expert={expert} (RAM prefix)", flush=True) + self._verify_slot(cache, layer, expert) + def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots (the normal kernel restricted to K experts; the following ``copy_missing`` diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 934822934..80de7c0a0 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -1046,18 +1046,20 @@ def copy_missing(self) -> None: self.src_indices, self.num_indices, ) - return - - from freetoken.kernel import fast_index_copy_jit + else: + from freetoken.kernel import fast_index_copy_jit - for per_layer, cache in self.banks: - fast_index_copy_jit( - cache, - self.evict_slots, - per_layer[layer_id], - self.src_indices, - self.num_indices, - ) + for per_layer, cache in self.banks: + fast_index_copy_jit( + cache, + self.evict_slots, + per_layer[layer_id], + self.src_indices, + self.num_indices, + ) + if (self._disk_tier is not None and layer_id == 0 + and os.environ.get("FT_DISK_TIER_VERIFY")): + self._disk_tier.verify_ram(self, layer_id) def iter_offload_moe_layers(model) -> Iterator: From 3de005398876314e577411131e50252c2b25aaa2 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 02:11:39 +0000 Subject: [PATCH 19/38] chore(moe): verify decode-fetched disk expert slot (one-shot) --- python/freetoken/moe/disk_tier.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 045a20a89..91a2c41cc 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -201,6 +201,7 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int self._fds: dict[int, tuple[int, bool]] = {} self._fetches = 0 self._fetch_bytes = 0 + self._decode_verify_done = False # ------------------------------------------------------------------ fds def _fd(self, shard_idx: int) -> tuple[int, bool]: @@ -280,11 +281,12 @@ def _sync_fetches(self) -> None: return # CPU-only tests: the copies are synchronous CPU->CPU torch.cuda.default_stream().synchronize() - def _verify_slot(self, cache, layer: int, expert: int) -> None: - """One-shot debug: read back a fetched disk expert's slot rows and compare - against the checkpoint bytes (ground truth). Gated on FT_DISK_TIER_VERIFY.""" + def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None) -> None: + """One-shot debug: read back an expert's slot rows and compare against the + checkpoint bytes (ground truth). Gated on FT_DISK_TIER_VERIFY.""" import torch - slot = expert # identity mapping in prefill + if slot is None: + slot = expert # identity mapping (prefill) for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): slot_row = gpu_cache[slot].contiguous() flat = slot_row.view(torch.uint8).reshape(-1) @@ -380,6 +382,12 @@ def fetch_pending(self, cache, layer_id: int) -> None: for f in futures: f.result() self._sync_fetches() + if (os.environ.get("FT_DISK_TIER_VERIFY") and not self._decode_verify_done + and layer_id == 0): + i0 = disk[0] + self._decode_verify_done = True + print(f"[verify-decode] layer=0 expert={int(src[i0])} slot={int(slots[i0])}", flush=True) + self._verify_slot(cache, layer_id, int(src[i0]), int(slots[i0])) disk_set = set(disk) ram = [i for i in range(n) if i not in disk_set] if ram: From ee8b9be9f3c8ed76194526d2a5108cf7cafa8808 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 02:16:54 +0000 Subject: [PATCH 20/38] chore(moe): verify prefill layer 20 + first decode fetch (any layer) --- python/freetoken/moe/disk_tier.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 91a2c41cc..2936c9608 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -356,8 +356,8 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N for f in futures: f.result() self._sync_fetches() - if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0 and disk.numel() > 0: - self._verify_slot(cache, 0, int(disk[0].item())) + if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id in (0, 20) and disk.numel() > 0: + self._verify_slot(cache, layer_id, int(disk[0].item())) # Same bookkeeping the materialize kernel writes, per fetched expert. flat = layer_id * cache.num_experts + disk cache.slot_for_id[layer_id, disk] = disk @@ -382,11 +382,11 @@ def fetch_pending(self, cache, layer_id: int) -> None: for f in futures: f.result() self._sync_fetches() - if (os.environ.get("FT_DISK_TIER_VERIFY") and not self._decode_verify_done - and layer_id == 0): + if os.environ.get("FT_DISK_TIER_VERIFY") and not self._decode_verify_done: i0 = disk[0] self._decode_verify_done = True - print(f"[verify-decode] layer=0 expert={int(src[i0])} slot={int(slots[i0])}", flush=True) + print(f"[verify-decode] layer={layer_id} expert={int(src[i0])} slot={int(slots[i0])} " + f"ndisk={len(disk)}", flush=True) self._verify_slot(cache, layer_id, int(src[i0]), int(slots[i0])) disk_set = set(disk) ram = [i for i in range(n) if i not in disk_set] From 191d38508736fd22c2608f5d230094969bfb7041 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 02:28:41 +0000 Subject: [PATCH 21/38] chore(moe): verify all 6 fetched disk experts per layer (confirm staging race) --- python/freetoken/moe/disk_tier.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 2936c9608..80106e6d5 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -357,7 +357,8 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N f.result() self._sync_fetches() if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id in (0, 20) and disk.numel() > 0: - self._verify_slot(cache, layer_id, int(disk[0].item())) + for e in disk.tolist()[:6]: + self._verify_slot(cache, layer_id, int(e)) # Same bookkeeping the materialize kernel writes, per fetched expert. flat = layer_id * cache.num_experts + disk cache.slot_for_id[layer_id, disk] = disk From 162098b6afc011d0db47c2a6a3dcdd4d00da1c34 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:10:55 +0000 Subject: [PATCH 22/38] chore(moe): verify ALL layer-0 disk experts (prefill) + 3 decode steps (race hunt) --- python/freetoken/moe/disk_tier.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 80106e6d5..fb670dbfd 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -201,7 +201,7 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int self._fds: dict[int, tuple[int, bool]] = {} self._fetches = 0 self._fetch_bytes = 0 - self._decode_verify_done = False + self._decode_verify_steps = 0 # ------------------------------------------------------------------ fds def _fd(self, shard_idx: int) -> tuple[int, bool]: @@ -357,7 +357,8 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N f.result() self._sync_fetches() if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id in (0, 20) and disk.numel() > 0: - for e in disk.tolist()[:6]: + limit = disk.numel() if layer_id == 0 else 6 # layer 0: ALL experts (race hunt) + for e in disk.tolist()[:limit]: self._verify_slot(cache, layer_id, int(e)) # Same bookkeeping the materialize kernel writes, per fetched expert. flat = layer_id * cache.num_experts + disk @@ -383,12 +384,13 @@ def fetch_pending(self, cache, layer_id: int) -> None: for f in futures: f.result() self._sync_fetches() - if os.environ.get("FT_DISK_TIER_VERIFY") and not self._decode_verify_done: - i0 = disk[0] - self._decode_verify_done = True - print(f"[verify-decode] layer={layer_id} expert={int(src[i0])} slot={int(slots[i0])} " - f"ndisk={len(disk)}", flush=True) - self._verify_slot(cache, layer_id, int(src[i0]), int(slots[i0])) + if (os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0 + and self._decode_verify_steps < 3): + self._decode_verify_steps += 1 + for i in disk[:8]: + print(f"[verify-decode] step={self._decode_verify_steps} " + f"expert={int(src[i])} slot={int(slots[i])} ndisk={len(disk)}", flush=True) + self._verify_slot(cache, layer_id, int(src[i]), int(slots[i])) disk_set = set(disk) ram = [i for i in range(n) if i not in disk_set] if ram: From df3dca778c1cb603681d27d392bb311f03b11f61 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:18:18 +0000 Subject: [PATCH 23/38] chore(moe): verify RAM-expert HOST row too (separate host vs PCIe-copy corruption) --- python/freetoken/moe/disk_tier.py | 62 ++++++++++++++++++------------- 1 file changed, 37 insertions(+), 25 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index fb670dbfd..49eaa7d92 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -290,31 +290,9 @@ def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None) for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): slot_row = gpu_cache[slot].contiguous() flat = slot_row.view(torch.uint8).reshape(-1) - segs = self._index.row_segments(bank_idx, layer, expert) - dsl = self._dst_slices[bank_idx] - # Reconstruct the reference row bytes from the checkpoint. - row_bytes = flat.numel() - ref = torch.zeros(row_bytes, dtype=torch.uint8) - row_el = slot_row.element_size() - row_leading = slot_row.numel() // slot_row.shape[0] if slot_row.dim() > 1 else 1 - for (d0, d1), (shard_idx, off, nbytes) in zip(dsl, segs): - fd, direct = self._fd(shard_idx) - a0 = off if not direct else (off & ~(_ALIGN - 1)) - slen = nbytes if not direct else (off + nbytes - a0 + _ALIGN - 1) & ~(_ALIGN - 1) - buf = os.pread(fd, slen, a0) - row_off = off - a0 - seg = buf[row_off:row_off + nbytes] - if bank_idx in (2, 5): - val = int.from_bytes(seg, "little") - import struct as _st - f32 = _st.unpack(" 1 else 1) try: ref = ref.to(flat.device) match = bool(torch.equal(flat, ref)) @@ -325,12 +303,46 @@ def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None) except Exception as exc: # never crash the server in debug print(f"[verify] bank={bank_idx} expert={expert} ERROR {exc!r}", flush=True) + def _ref_row(self, bank_idx: int, layer: int, expert: int, row_bytes: int, + row_el: int, row_leading: int) -> torch.Tensor: + """Reference row bytes for (bank, layer, expert) straight from the checkpoint.""" + ref = torch.zeros(row_bytes, dtype=torch.uint8) + segs = self._index.row_segments(bank_idx, layer, expert) + for (d0, d1), (shard_idx, off, nbytes) in zip(self._dst_slices[bank_idx], segs): + fd, direct = self._fd(shard_idx) + a0 = off if not direct else (off & ~(_ALIGN - 1)) + slen = nbytes if not direct else (off + nbytes - a0 + _ALIGN - 1) & ~(_ALIGN - 1) + buf = os.pread(fd, slen, a0) + row_off = off - a0 + seg = buf[row_off:row_off + nbytes] + if bank_idx in (2, 5): + import struct as _st + import numpy as _np + f16 = _np.float16(_st.unpack(" None: """One-shot debug: after the PCIe copy, check a RAM-resident expert's slot rows against the checkpoint reference. Gated on FT_DISK_TIER_VERIFY.""" expert = min(10, self._ram - 1) # a RAM-resident expert print(f"[verify-ram] layer={layer} expert={expert} (RAM prefix)", flush=True) self._verify_slot(cache, layer, expert) + # Also check the HOST row (CPU) -- separates host-side corruption from the + # GPU PCIe copy. + for bank_idx, (host_layer, _gpu) in enumerate(self._banks): + host_row = host_layer[layer][expert].contiguous() + flat = host_row.view(torch.uint8).reshape(-1) + ref = self._ref_row(bank_idx, layer, expert, flat.numel(), + host_row.element_size(), + host_row.numel() // host_row.shape[0] if host_row.dim() > 1 else 1) + match = bool(torch.equal(flat, ref)) + print(f"[verify-host] bank={bank_idx} expert={expert} match={match} " + f"host_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", flush=True) def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots From 90726fdeb80d12c5979ba4eb662c67d07f6a5bed Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:29:49 +0000 Subject: [PATCH 24/38] chore(moe): verify ALL RAM experts post-copy + identify corrupted-slot byte source --- python/freetoken/moe/disk_tier.py | 67 ++++++++++++++++++++++++++----- 1 file changed, 56 insertions(+), 11 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 49eaa7d92..6a33edd90 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -332,17 +332,62 @@ def verify_ram(self, cache, layer: int) -> None: expert = min(10, self._ram - 1) # a RAM-resident expert print(f"[verify-ram] layer={layer} expert={expert} (RAM prefix)", flush=True) self._verify_slot(cache, layer, expert) - # Also check the HOST row (CPU) -- separates host-side corruption from the - # GPU PCIe copy. - for bank_idx, (host_layer, _gpu) in enumerate(self._banks): - host_row = host_layer[layer][expert].contiguous() - flat = host_row.view(torch.uint8).reshape(-1) - ref = self._ref_row(bank_idx, layer, expert, flat.numel(), - host_row.element_size(), - host_row.numel() // host_row.shape[0] if host_row.dim() > 1 else 1) - match = bool(torch.equal(flat, ref)) - print(f"[verify-host] bank={bank_idx} expert={expert} match={match} " - f"host_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", flush=True) + # Check ALL RAM experts: slot vs host row (host correctness established + # separately). Count mismatches; identify the source of the first one. + n_bad = 0 + identified = False + for e in range(self._ram): + for bank_idx, (host_layer, gpu_cache) in enumerate(self._banks): + slot_row = gpu_cache[e].contiguous() + flat = slot_row.view(torch.uint8).reshape(-1) + hflat = host_layer[layer][e].contiguous().view(torch.uint8).reshape(-1) + if flat.numel() != hflat.numel() or not bool(torch.equal(flat.cpu(), hflat)): + n_bad += 1 + if n_bad <= 12: + print(f"[verify-ram] MISMATCH e={e} bank={bank_idx} " + f"slot_head={flat[:8].tolist()} host_head={hflat[:8].tolist()}", + flush=True) + if not identified: + identified = True + self._identify_source(layer, bank_idx, flat) + print(f"[verify-ram] layer={layer} mismatches={n_bad}/{self._ram * len(self._banks)}", + flush=True) + + def _identify_source(self, layer: int, bank_idx: int, flat: torch.Tensor) -> None: + """Debug: find where a corrupted slot's bytes came from (GPU slot, host row, + or checkpoint row).""" + flat_cpu = flat.cpu() + head = flat_cpu[:16] + found = [] + _host_layer, gpu_cache = self._banks[bank_idx] + gpu_flat = gpu_cache.view(torch.uint8).reshape(gpu_cache.shape[0], -1) + if gpu_flat.shape[1] == flat_cpu.numel(): + cand = torch.nonzero( + (gpu_flat[:, :16].cpu() == head.unsqueeze(0)).all(dim=1)).flatten().tolist() + for s in cand[:16]: + if bool(torch.equal(gpu_flat[s].cpu(), flat_cpu)): + found.append(f"gpu_slot={s}(L{layer},B{bank_idx})") + for e in range(self._ram): + hflat = _host_layer[layer][e].contiguous().view(torch.uint8).reshape(-1) + if hflat.numel() == flat_cpu.numel() and bool(torch.equal(hflat, flat_cpu)): + found.append(f"host_row_e{e}(L{layer},B{bank_idx})") + import itertools + num_layers = len(self._banks[bank_idx][0]) + targets = sorted(set(itertools.product([layer], range(256))) + | set(itertools.product(range(num_layers), [10]))) + for (L, e) in targets: + row_el, row_leading = None, None + hrow = _host_layer[layer][0] + row_el = hrow.element_size() + row_leading = hrow.numel() // hrow.shape[0] if hrow.dim() > 1 else 1 + try: + ref = self._ref_row(bank_idx, L, e, flat_cpu.numel(), row_el, row_leading) + except Exception: + continue + if bool(torch.equal(ref, flat_cpu)): + found.append(f"checkpoint_L{L}_e{e}") + print(f"[identify] L{layer} B{bank_idx} n={flat_cpu.numel()} " + f"source={found if found else 'UNKNOWN'}", flush=True) def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots From a9b93b6d9e4a4116691b84b8a3b98a60534e0814 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:39:07 +0000 Subject: [PATCH 25/38] chore(moe): instrument copy_missing/fetch_pending (layer 0) for the all-RAM-mismatch hunt --- python/freetoken/moe/disk_tier.py | 3 +++ python/freetoken/moe/offload_cache.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 6a33edd90..e8a43c9ca 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -432,6 +432,9 @@ def fetch_pending(self, cache, layer_id: int) -> None: src = cache.src_indices[:n].cpu() slots = cache.evict_slots[:n].cpu() disk = [i for i in range(n) if int(src[i]) >= self._ram] + if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0: + print(f"[fetch-pend] layer=0 n={n} ndisk={len(disk)} " + f"src_head={src[:4].tolist()}", flush=True) if not disk: return futures = [ diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 80de7c0a0..cfe439f87 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -1031,6 +1031,11 @@ def copy_missing(self) -> None: for per_layer, cache in self.banks: cache[: self.num_experts].copy_(per_layer[layer_id]) return + if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0: + print(f"[copy-miss] layer={layer_id} fused={self._copy_fused_ok} " + f"n={int(self.num_indices.item())} " + f"evict={self.evict_slots[:4].cpu().tolist()} " + f"src={self.src_indices[:4].cpu().tolist()}", flush=True) if self._copy_fused_ok: from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit From 5fd7d85ef0a81ef5449033ecfd5041c3791dec26 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:52:08 +0000 Subject: [PATCH 26/38] chore(moe): verify decode slot-mapping end-to-end (slot data vs bookkeeping expert) --- python/freetoken/layers/moe.py | 3 +++ python/freetoken/moe/disk_tier.py | 38 +++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 0293a82d3..dc6dc37b5 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -311,6 +311,9 @@ def _decode_routed( return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) cache.ensure_experts(self.layer_id, topk_ids) cache.copy_missing() + if (cache.disk_tier_enabled and self.layer_id == 0 + and os.environ.get("FT_DISK_TIER_VERIFY")): + cache._disk_tier.verify_decode_mapping(cache, self.layer_id, topk_ids) return self._expert_gemm( cache, hidden_states, diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index e8a43c9ca..4f2ae982e 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -202,6 +202,8 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int self._fetches = 0 self._fetch_bytes = 0 self._decode_verify_steps = 0 + self._map_verify_steps = 0 + self._cache = cache # ------------------------------------------------------------------ fds def _fd(self, shard_idx: int) -> tuple[int, bool]: @@ -389,6 +391,41 @@ def _identify_source(self, layer: int, bank_idx: int, flat: torch.Tensor) -> Non print(f"[identify] L{layer} B{bank_idx} n={flat_cpu.numel()} " f"source={found if found else 'UNKNOWN'}", flush=True) + def verify_decode_mapping(self, cache, layer_id: int, topk_ids: torch.Tensor) -> None: + """Debug: after the LRU rewrite + fetch/copy, check that every slot the GEMM + will read actually holds the expert the bookkeeping says it holds. Gated on + FT_DISK_TIER_VERIFY; first 4 decode steps only.""" + if self._map_verify_steps >= 4: + return + self._map_verify_steps += 1 + slots = torch.unique(topk_ids.reshape(-1)) + nbad = 0 + for s in slots.tolist(): + s = int(s) + flat_id = int(cache.id_of_slot[s].item()) + if flat_id < 0: + print(f"[verify-map] step={self._map_verify_steps} slot={s} id_of_slot=-1", + flush=True) + nbad += 1 + continue + expert = flat_id % cache.num_experts + for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): + slot_row = gpu_cache[s].contiguous() + flat = slot_row.view(torch.uint8).reshape(-1) + ref = self._ref_row(bank_idx, layer_id, expert, flat.numel(), + slot_row.element_size(), + slot_row.numel() // slot_row.shape[0] + if slot_row.dim() > 1 else 1) + if not bool(torch.equal(flat.cpu(), ref)): + nbad += 1 + if nbad <= 8: + print(f"[verify-map] step={self._map_verify_steps} slot={s} " + f"expert={expert} bank={bank_idx} MISMATCH " + f"slot_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", + flush=True) + print(f"[verify-map] step={self._map_verify_steps} slots={slots.numel()} bad={nbad}", + flush=True) + def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: """Disk-tier prefill: materialize the RAM-resident prefix into identity slots (the normal kernel restricted to K experts; the following ``copy_missing`` @@ -462,6 +499,7 @@ def fetch_pending(self, cache, layer_id: int) -> None: def refresh(self, cache) -> None: """Rebind the slot-cache references after a runtime cache rebuild.""" self._banks = list(cache.banks) + self._cache = cache def stats(self) -> dict: return {"experts_fetched": self._fetches, "bytes_fetched": self._fetch_bytes} From 541ea78662ae5c5bbff8f459eb61375103610f71 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 06:50:09 +0000 Subject: [PATCH 27/38] fix(moe): clear stale disk-slot bookkeeping in disk-tier prefill materialize The GPU slot cache is one shared pool across all layers. Prefill identity mapping owns all of slots [0, E) per layer, but the materialize kernel only scans slots < materialize_count (the RAM prefix), so disk slots [ram, E) that held a previous prefill layer's experts kept their slot_for_id entries. Decode then took phantom hits on those slots and read another layer's weights (expert 112 verified correct at prefill, corrupted by decode step 1). Clear the stale entries device-side before the kernel. --- python/freetoken/moe/disk_tier.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 4f2ae982e..2d7f1aa13 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -434,6 +434,18 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N preserved, so the prefill GEMM is unchanged.""" from freetoken.moe.offload_kernels import _materialize_layer_gpu + # Prefill identity mapping owns ALL of slots [0, E) for this layer, but + # the kernel only scans slots < materialize_count, so the disk slots + # [ram, E) that still hold a previous layer's experts (previous prefill + # layer or decode LRU) would keep their slot_for_id entries -- phantom + # decode hits that read another layer's weights. Clear them first + # (device-side, no sync). + seg = cache.id_of_slot[self._ram:cache.num_experts] + valid = seg >= 0 + cache.slot_for_id.view(-1)[seg[valid].long()] = -1 + seg[valid] = -1 + cache.usage[self._ram:cache.num_experts][valid] = 0 + _materialize_layer_gpu(cache, layer_id, materialize_count=self._ram) routed = expert_ids.reshape(-1) disk = torch.unique(routed[routed >= self._ram]) From 1c4f4275082d9e994ef180441abf647f1ba791f7 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 07:53:47 +0000 Subject: [PATCH 28/38] chore(moe): run disk-tier verify_ram on prefill only The RAM-slot identity check (slot e == host row e) only holds under the prefill identity mapping; during decode the LRU owns slots 0..63. It was firing on every decode step for layer 0 (384 experts x banks of D2H + CPU compare per step), dragging the instrumented E2E from ~11 tok/s to 1.9. Track whether the pending miss list came from materialize_layer (prefill) or ensure_experts (decode) and gate on it. --- python/freetoken/moe/offload_cache.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index cfe439f87..17e220bbe 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -865,6 +865,7 @@ def materialize_layer(self, layer_id: int, expert_ids: torch.Tensor | None = Non # disk-resident experts into their identity slots (needs the routing). assert expert_ids is not None, "disk-tier prefill needs the routed expert ids" self._pending_src_layer = layer_id + self._pending_whole_layer = True self._disk_tier.materialize_layer(self, layer_id, expert_ids) return from freetoken.moe.offload_kernels import materialize_layer @@ -1063,6 +1064,7 @@ def copy_missing(self) -> None: self.num_indices, ) if (self._disk_tier is not None and layer_id == 0 + and self._pending_whole_layer and os.environ.get("FT_DISK_TIER_VERIFY")): self._disk_tier.verify_ram(self, layer_id) From de5963b67fd6b8a576ce696b5946afb3f602995d Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:12:27 +0000 Subject: [PATCH 29/38] chore(moe): log per-rank bank row shapes vs disk-index row bytes at DiskTier init TP=2 disktier debug (step 2): prove empirically whether the host bank rows are full-per-rank or TP-sharded, and that the disk index's full-row segments match them byte-for-byte. Gated on FT_DISK_TIER_VERIFY. --- python/freetoken/moe/disk_tier.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 2d7f1aa13..5aff22395 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -193,6 +193,19 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int self._dst_slices.append([(0, mid), (mid, row.shape[0])]) else: self._dst_slices.append([(0, row.shape[0])]) + if os.environ.get("FT_DISK_TIER_VERIFY"): + # TP=2 debug: prove per-rank whether the host bank rows are full or + # TP-sharded, and that the disk index's full-row segments match them. + from freetoken.distributed import try_get_tp_info + tp = try_get_tp_info() + host_shapes = [tuple(b[0][0][0].shape) for b in self._banks] + disk_bytes = [ + sum(nb for _, _, nb in self._index.row_segments(bi, 0, 100)) + for bi in range(len(self._banks)) + ] + print(f"[disk-tier-init] tp_rank={getattr(tp, 'rank', '?')}/{getattr(tp, 'size', '?')} " + f"ram={ram_experts} host_row_shapes={host_shapes} " + f"host_row_bytes={self._row_bytes} disk_row_bytes={disk_bytes}", flush=True) max_row = max(self._row_bytes) self._staging_size = ((max_row + _ALIGN - 1) // _ALIGN + 2) * _ALIGN self._staging = threading.local() From 76fa506c0533fe32ac198a9137f60cdba0d76b75 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:04:29 +0000 Subject: [PATCH 30/38] fix(moe): close disk-tier staging-buffer reuse race with a per-thread ring Each worker thread preadved the next bank's bytes into the SAME pinned staging buffer while the previous bank's async H2D copy was still DMAing from it. The copy is cudaMemcpyAsync: copy_ returns after ENQUEUE, and the GPU reads the host bytes later, so the next preadv could land mid-DMA and the slot row came out as a partial mix of two experts' data. TP=1 rarely hit it (single rank's disk load, idle-ish stream); TP=2 doubles the NVMe load and adds NCCL stream work, so the window opened and the run degenerated (2/5412 verified rows wrong, both gate|up weight_scale). Fix: a per-thread ring of _STAGING_RING pinned buffers, each armed with a CUDA event recorded after the copy that used it; reusing a buffer waits on that event. Exact, and the wait is a no-op whenever the ring outruns the DMA (the normal case), so disk/GPU overlap is preserved. Also: [verify] lines now carry phase/layer/slot and, on mismatch, the diff span plus an overwriter hunt (which checkpoint row the slot actually holds); _preadv_error call now passes its full signature. Test stub updated to the ring. --- python/freetoken/moe/disk_tier.py | 90 ++++++++++++++++++++++++++----- tests/moe/test_disk_tier.py | 17 +++--- 2 files changed, 87 insertions(+), 20 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 5aff22395..b38b91d76 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -239,13 +239,26 @@ def _fd(self, shard_idx: int) -> tuple[int, bool]: return ent # --------------------------------------------------------------- staging - def _staging_buf(self) -> HostBank: - buf = getattr(self._staging, "buf", None) - if buf is None: - buf = HostBank((self._staging_size,), torch.uint8) - buf.pin() # small; pin once per worker thread - self._staging.buf = buf - return buf + # Staging ring depth per worker thread. A buffer must not be overwritten by the + # next preadv until the async H2D copy that read it has finished (pinned-memory + # reuse race -- the copy is DMA, still reading host bytes after copy_ returns). + # The depth only needs to cover one copy's DMA time in host-side preadv time; + # the per-slot CUDA event below makes any shallower lap correct, just slower. + _STAGING_RING = 8 + + def _staging_ring(self) -> list: + ring = getattr(self._staging, "ring", None) + if ring is None: + ring = [] + for _ in range(self._STAGING_RING): + buf = HostBank((self._staging_size,), torch.uint8) + buf.pin() # pin once per worker thread + ev = torch.cuda.Event() if torch.cuda.is_available() else None + if ev is not None: + ev.record() # start "complete"; re-recorded after each copy + ring.append([buf, ev]) + self._staging.ring = ring + return ring # ---------------------------------------------------------------- fetch def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: @@ -255,11 +268,18 @@ def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: self._fetch_expert_inner(layer, expert, slot) def _fetch_expert_inner(self, layer: int, expert: int, slot: int) -> None: - staging = self._staging_buf() + ring = self._staging_ring() + ri = getattr(self._staging, "ri", 0) for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): row = gpu_cache[slot] segs = self._index.row_segments(bank_idx, layer, expert) for (d0, d1), (shard_idx, off, nbytes) in zip(self._dst_slices[bank_idx], segs): + staging, ev = ring[ri] + if ev is not None: + # This buffer's last async H2D copy must be done before the + # preadv below overwrites it (pinned-memory reuse race). + ev.synchronize() + ri = (ri + 1) % len(ring) fd, direct = self._fd(shard_idx) if direct: a0 = off & ~(_ALIGN - 1) @@ -283,6 +303,10 @@ def _fetch_expert_inner(self, layer: int, expert: int, slot: int) -> None: src = staging.tensor[row_off:row_off + nbytes] dst = row[d0:d1] dst.copy_(src.view(dst.dtype).view(dst.shape), non_blocking=True) + if ev is not None: + # Arm: the next reuse of this buffer waits for this copy. + ev.record() + self._staging.ri = ri self._fetches += 1 self._fetch_bytes += sum(self._row_bytes) @@ -296,7 +320,8 @@ def _sync_fetches(self) -> None: return # CPU-only tests: the copies are synchronous CPU->CPU torch.cuda.default_stream().synchronize() - def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None) -> None: + def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None, + phase: str = "prefill") -> None: """One-shot debug: read back an expert's slot rows and compare against the checkpoint bytes (ground truth). Gated on FT_DISK_TIER_VERIFY.""" import torch @@ -311,10 +336,19 @@ def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None) try: ref = ref.to(flat.device) match = bool(torch.equal(flat, ref)) - print(f"[verify] bank={bank_idx} expert={expert} match={match} " + if match: + print(f"[verify] {phase} L{layer} bank={bank_idx} expert={expert} " + f"slot={slot} match=True", flush=True) + continue + diff = (flat != ref) + nz = torch.nonzero(diff).flatten() + print(f"[verify] {phase} L{layer} bank={bank_idx} expert={expert} " + f"slot={slot} match=False n_diff={int(diff.sum())}/{flat.numel()} " + f"first_off={int(nz[0])} last_off={int(nz[-1])} " f"slot_norm={slot_row.float().norm().item():.4f} " f"ref_norm={ref.view(slot_row.dtype).view(slot_row.shape).float().norm().item():.4f} " f"slot_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", flush=True) + self._identify_overwriter(layer, bank_idx, expert, flat) except Exception as exc: # never crash the server in debug print(f"[verify] bank={bank_idx} expert={expert} ERROR {exc!r}", flush=True) @@ -341,6 +375,37 @@ def _ref_row(self, bank_idx: int, layer: int, expert: int, row_bytes: int, ref[dst_off:dst_off + len(seg)] = torch.frombuffer(seg, dtype=torch.uint8) return ref + def _identify_overwriter(self, layer: int, bank_idx: int, expert: int, + flat: torch.Tensor) -> None: + """Debug: on a verify mismatch, find whose row the slot actually holds. + + Compares the slot's first 64 bytes against (a) every other expert of the + same layer and (b) the same expert in every other layer, straight from the + checkpoint. A hit names the overwriter (e.g. a staging-buffer reuse race + landing a neighbour's preadv); no hit means a partial mix. Only runs on + mismatch, so the ~300 extra preads are free otherwise.""" + head = flat[:64].cpu() + host_row = self._banks[bank_idx][0][0][layer] + row_el = host_row.element_size() + row_leading = host_row.numel() // host_row.shape[0] if host_row.dim() > 1 else 1 + num_experts = host_row.shape[0] + num_layers = len(self._banks[bank_idx][0]) + hits = [] + for e in range(num_experts): + if e == expert: + continue + ref = self._ref_row(bank_idx, layer, e, flat.numel(), row_el, row_leading) + if bool(torch.equal(ref[:64], head)): + hits.append(f"L{layer}_e{e}") + for L in range(num_layers): + if L == layer: + continue + ref = self._ref_row(bank_idx, L, expert, flat.numel(), row_el, row_leading) + if bool(torch.equal(ref[:64], head)): + hits.append(f"L{L}_e{expert}") + print(f"[overwriter] L{layer} B{bank_idx} e{expert} head64 matches: " + f"{hits if hits else 'NONE (partial mix?)'}", flush=True) + def verify_ram(self, cache, layer: int) -> None: """One-shot debug: after the PCIe copy, check a RAM-resident expert's slot rows against the checkpoint reference. Gated on FT_DISK_TIER_VERIFY.""" @@ -478,7 +543,7 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id in (0, 20) and disk.numel() > 0: limit = disk.numel() if layer_id == 0 else 6 # layer 0: ALL experts (race hunt) for e in disk.tolist()[:limit]: - self._verify_slot(cache, layer_id, int(e)) + self._verify_slot(cache, layer_id, int(e), phase="prefill") # Same bookkeeping the materialize kernel writes, per fetched expert. flat = layer_id * cache.num_experts + disk cache.slot_for_id[layer_id, disk] = disk @@ -512,7 +577,8 @@ def fetch_pending(self, cache, layer_id: int) -> None: for i in disk[:8]: print(f"[verify-decode] step={self._decode_verify_steps} " f"expert={int(src[i])} slot={int(slots[i])} ndisk={len(disk)}", flush=True) - self._verify_slot(cache, layer_id, int(src[i]), int(slots[i])) + self._verify_slot(cache, layer_id, int(src[i]), int(slots[i]), + phase="decode") disk_set = set(disk) ram = [i for i in range(n) if i not in disk_set] if ram: diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index c75a1340a..96e28b407 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -171,17 +171,18 @@ def _tier(checkpoint, cache, ram_experts=2): b[0][0][0].numel() * b[0][0][0].element_size() for b in cache.banks) assert tier._staging_size >= max_row + 2 * 4096, (tier._staging_size, max_row) # Stub the pinned staging (HostBank.pin needs CUDA) with the same per-thread - # buffer semantics as production (threading.local). + # ring semantics as production (threading.local, no CUDA events on CPU). local = threading.local() - def _staging_buf(): - buf = getattr(local, "buf", None) - if buf is None: - buf = HostBank((tier._staging_size,), torch.uint8) - local.buf = buf - return buf + def _staging_ring(): + ring = getattr(local, "ring", None) + if ring is None: + ring = [[HostBank((tier._staging_size,), torch.uint8), None] + for _ in range(tier._STAGING_RING)] + local.ring = ring + return ring - tier._staging_buf = _staging_buf + tier._staging_ring = _staging_ring return tier From 623ca1d15b193762878f0d6f57fb3347e5d1cc60 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Wed, 2 Sep 2026 00:33:17 +0000 Subject: [PATCH 31/38] fix(moe): disk-tier fetch workers must run on the rank's CUDA device Fresh ThreadPoolExecutor threads default to CUDA device 0, but a TP>1 rank lives on another device. The H2D slot copies land on the destination tensor's device stream (copy_ guards on dst), while ev.record() uses the thread's CURRENT stream -- so on TP=2 rank 1 the staging-ring reuse guard waited on an idle device-0 stream and a preadv could overwrite a pinned buffer mid-DMA. Seen on Bandit (4.3 GB/s NVMe, TP=2, RAM=64): 6/12288 layer-0 slot rows corrupted in one prefill (e4m3 scale banks mixed into e4m3 NaN encodings; weight banks partial mixes; scalar-fill banks clean). Rudi's 1.8 GB/s disk made the preadv slower than the DMA, hiding the window. - _staging_ring: torch.cuda.set_device(rank device) once per worker thread, so ev.record() lands on the stream the copies use. - _sync_fetches: sync the rank device's default stream explicitly. - _identify_overwriter: num_experts came from an expert ROW's shape[0] (1024/128/32/2048 per bank) -- the scan ran past the index and died with struct.error, masking overwriter attribution. --- python/freetoken/moe/disk_tier.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index b38b91d76..9f1606371 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -249,6 +249,15 @@ def _fd(self, shard_idx: int) -> tuple[int, bool]: def _staging_ring(self) -> list: ring = getattr(self._staging, "ring", None) if ring is None: + # Fresh worker threads default to CUDA device 0, but the rank may + # live on another device (TP>1). The H2D copies land on the + # destination tensor's device stream, while ev.record() below uses + # the thread's CURRENT stream -- without this, the ring's reuse + # guard waits on an idle stream and a preadv can overwrite the + # buffer mid-DMA (corrupted slot rows on TP=2 rank 1). + dev = self._banks[0][1].device + if dev.type == "cuda": + torch.cuda.set_device(dev) ring = [] for _ in range(self._STAGING_RING): buf = HostBank((self._staging_size,), torch.uint8) @@ -318,7 +327,7 @@ def _sync_fetches(self) -> None: stream, so sync the default stream before the GEMM reads the slots.""" if not torch.cuda.is_available(): return # CPU-only tests: the copies are synchronous CPU->CPU - torch.cuda.default_stream().synchronize() + torch.cuda.default_stream(self._banks[0][1].device).synchronize() def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None, phase: str = "prefill") -> None: @@ -385,10 +394,10 @@ def _identify_overwriter(self, layer: int, bank_idx: int, expert: int, landing a neighbour's preadv); no hit means a partial mix. Only runs on mismatch, so the ~300 extra preads are free otherwise.""" head = flat[:64].cpu() - host_row = self._banks[bank_idx][0][0][layer] + host_row = self._banks[bank_idx][0][0][0] # expert-0 row (all rows share its shape) row_el = host_row.element_size() row_leading = host_row.numel() // host_row.shape[0] if host_row.dim() > 1 else 1 - num_experts = host_row.shape[0] + num_experts = self._cache.num_experts num_layers = len(self._banks[bank_idx][0]) hits = [] for e in range(num_experts): From de4346908a719dd5530e06339ffa54d330adf097 Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:29:21 +0000 Subject: [PATCH 32/38] fix(moe): apply reviewer fixes from PR #337 (MT-z) - never materialize disk-resident expert rows at load: serial loader skips rows >= K before get_tensor, parallel loader filters at the reader, and the completion tracker / placed assert count rows_per_layer instead of E. Load-time RAM peak is now K/E of the expert set instead of the full set. - thread disk_tier= through the remaining six NVFP4 family loaders (gemma4, glm4_moe, glm5_next, minimax_m2, minimax_m3, qwen4_exp) so --moe-disk-tier on no longer TypeErrors outside qwen3_5_moe. - key _sync_fetches on the banks' device type instead of torch.cuda.is_available(): CPU-bank unit tests no longer fall through to torch.cuda.default_stream(cpu_device) on CUDA machines. Diffs from MT-z's PR #337 review comment, applied verbatim. Verified: tests/moe/test_disk_tier.py 6/6 pass in a CUDA container (freetoken:local, Rudi GPU0), incl. the two previously failing fetch_pending tests; inspect.signature confirms disk_tier on all 13 wrappers. --- python/freetoken/models/gemma4/weight.py | 6 ++-- python/freetoken/models/glm4_moe/weight.py | 6 ++-- python/freetoken/models/glm5_next/weight.py | 6 +++- python/freetoken/models/minimax_m2/weight.py | 6 ++-- python/freetoken/models/minimax_m3/weight.py | 6 ++-- python/freetoken/models/nvfp4_banks.py | 31 ++++++++++++++++---- python/freetoken/models/qwen4_exp/weight.py | 6 ++-- python/freetoken/moe/disk_tier.py | 10 +++++-- 8 files changed, 58 insertions(+), 19 deletions(-) diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index 43c2355d6..ad9c43eb2 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -276,7 +276,7 @@ def _expert_name(raw_name: str) -> str | None: def load_nvfp4_expert_sources( - model_path: str, config, *, layer_sink=None + model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, list[torch.Tensor]]: """CPU NVFP4 expert source banks for the offload cache; see load_nvfp4_expert_source_banks.""" return load_nvfp4_expert_source_banks( @@ -286,11 +286,12 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -304,6 +305,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py index 7efc6e6be..a00005ec2 100644 --- a/python/freetoken/models/glm4_moe/weight.py +++ b/python/freetoken/models/glm4_moe/weight.py @@ -191,7 +191,7 @@ def _iter_resident_weights(reader, config, primary) -> Iterator[tuple[str, torch # -------------------------------------------------------------------------------------- # Routed expert host banks (NVFP4) for the offload cache. # -------------------------------------------------------------------------------------- -def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> dict[str, torch.Tensor]: +def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None) -> dict[str, torch.Tensor]: """Build the pinned CPU NVFP4 banks for GLM-4's routed experts. experts exist only for layers [first_k_dense_replace, num_layers) and pack by MoE layer @@ -205,11 +205,12 @@ def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> di drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded O_DIRECT reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -223,6 +224,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/glm5_next/weight.py b/python/freetoken/models/glm5_next/weight.py index fda852256..2e3b7f684 100644 --- a/python/freetoken/models/glm5_next/weight.py +++ b/python/freetoken/models/glm5_next/weight.py @@ -95,7 +95,10 @@ def _select_expert_source_spec(model_path: str) -> Nvfp4ExpertSourceSpec: _KDA_IN_PROJ = ("q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj") -def load_nvfp4_expert_sources(model_path: str, config, layer_sink=None): +def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None): + # disk_tier is threaded through the same way qwen3_5_moe does it: the caller builds the + # NVMe tier and every NVFP4 family's loader has to pass it down, or the tail of the bank + # is never registered and load fails with an unexpected-kwarg TypeError. return load_nvfp4_expert_source_banks( model_path, config, @@ -103,6 +106,7 @@ def load_nvfp4_expert_sources(model_path: str, config, layer_sink=None): drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/minimax_m2/weight.py b/python/freetoken/models/minimax_m2/weight.py index 103f8ec56..a010225bf 100644 --- a/python/freetoken/models/minimax_m2/weight.py +++ b/python/freetoken/models/minimax_m2/weight.py @@ -92,7 +92,7 @@ def load_nvfp4_expert_sources( model_path: str, config, *, - layer_sink=None, + layer_sink=None, disk_tier=None, ) -> dict[str, torch.Tensor]: """CPU NVFP4 expert source banks for the offload cache; see load_nvfp4_expert_source_banks.""" return load_nvfp4_expert_source_banks( @@ -102,11 +102,12 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded O_DIRECT reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -120,6 +121,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/minimax_m3/weight.py b/python/freetoken/models/minimax_m3/weight.py index e10a7237e..bfcfbbd87 100644 --- a/python/freetoken/models/minimax_m3/weight.py +++ b/python/freetoken/models/minimax_m3/weight.py @@ -249,7 +249,7 @@ def iter_weights( def load_nvfp4_expert_sources( - model_path: str, config, *, layer_sink=None + model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, list[torch.Tensor]]: """CPU NVFP4 expert source banks for the offload cache; see load_nvfp4_expert_source_banks.""" return load_nvfp4_expert_source_banks( @@ -259,11 +259,12 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded O_DIRECT reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -277,6 +278,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 68a5cfbf0..4386da857 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -152,6 +152,13 @@ def load_nvfp4_expert_source_banks( _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill K = disk_tier.ram_experts if disk_tier is not None else None + # Rows this loader actually materializes per layer. The banks are still allocated + # at the full [E, ...] shape -- rows K..E-1 are the disk tier's fetch destination and + # must exist as address space -- but they are never read or written here, so they + # stay unbacked. Filling them and releasing afterwards (the previous order) made the + # load-time peak the FULL expert set, which is exactly the case the tier exists for: + # GLM-5.3-Flash (166 GiB of experts) swapped a 61 GB box to a halt at shard 73/118. + rows_per_layer = E if K is None else min(K, E) if disk_tier is not None and layer_sink is not None: raise NotImplementedError("disk tier: the converter (layer_sink) path is not supported yet") gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] @@ -164,7 +171,7 @@ def load_nvfp4_expert_source_banks( from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline def _load(sink) -> int: - tracker = LayerCompletionTracker(E * 6, _hb, sink) + tracker = LayerCompletionTracker(rows_per_layer * 6, _hb, sink) placed = 0 for shard in tqdm(sorted(weight_shards), desc=f"Loading {spec.desc}", disable=not primary): path = os.path.join(folder, shard) @@ -172,6 +179,8 @@ def _load(sink) -> int: for name, match, bank_layer_id in weight_shards[shard]: layer = int(match.group("layer")) expert = int(match.group("expert")) + if expert >= rows_per_layer: + continue # disk-resident row: fetched on demand, never loaded here proj = match.group("proj") role = spec.proj_to_role[proj] kind = _canon_kind(spec, match.group("kind")) @@ -213,7 +222,7 @@ def _load(sink) -> int: release_bank_tails(_hb, E, K) - expected = num_layers * E * 6 + expected = num_layers * rows_per_layer * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" return { "gate_up_packed": gate_up_packed, @@ -284,6 +293,13 @@ def load_nvfp4_expert_source_banks_parallel( _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill K = disk_tier.ram_experts if disk_tier is not None else None + # Rows this loader actually materializes per layer. The banks are still allocated + # at the full [E, ...] shape -- rows K..E-1 are the disk tier's fetch destination and + # must exist as address space -- but they are never read or written here, so they + # stay unbacked. Filling them and releasing afterwards (the previous order) made the + # load-time peak the FULL expert set, which is exactly the case the tier exists for: + # GLM-5.3-Flash (166 GiB of experts) swapped a 61 GB box to a halt at shard 73/118. + rows_per_layer = E if K is None else min(K, E) if disk_tier is not None and layer_sink is not None: raise NotImplementedError("disk tier: the converter (layer_sink) path is not supported yet") gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] @@ -297,10 +313,15 @@ def load_nvfp4_expert_source_banks_parallel( # Pass 2: bulk weight/weight_scale via the common parallel reader; place by name. def _load(sink) -> int: - tracker = LayerCompletionTracker(E * 6, _hb, sink) + tracker = LayerCompletionTracker(rows_per_layer * 6, _hb, sink) placed = 0 + def _wanted(n: str) -> bool: + info = weight_info.get(n) + # Filter at the reader so disk-resident rows cost no I/O at all, not just no write. + return info is not None and int(info[0].group("expert")) < rows_per_layer + for name, tensor in iter_expert_tensors_parallel( - folder, lambda n: n in weight_info, workers=workers, chunk=chunk + folder, _wanted, workers=workers, chunk=chunk ): match, bank_layer_id = weight_info[name] layer = int(match.group("layer")) @@ -340,7 +361,7 @@ def _load(sink) -> int: release_bank_tails(_hb, E, K) - expected = num_layers * E * 6 + expected = num_layers * rows_per_layer * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" return { "gate_up_packed": gate_up_packed, diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index f8d2a7494..edd38f4b2 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -289,7 +289,7 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> dict: +def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None) -> dict: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" return load_nvfp4_expert_source_banks( model_path, @@ -298,11 +298,12 @@ def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> di drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ) -> dict: """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -316,6 +317,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 9f1606371..b5c7f1453 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -325,9 +325,13 @@ def _sync_fetches(self) -> None: The copies are enqueued on the pool threads' default stream; f.result() only waits for them to be ENQUEUED. The GEMM's stream is not ordered with that stream, so sync the default stream before the GEMM reads the slots.""" - if not torch.cuda.is_available(): - return # CPU-only tests: the copies are synchronous CPU->CPU - torch.cuda.default_stream(self._banks[0][1].device).synchronize() + # Key this on where the BANKS live, not on whether the machine has a GPU: the + # disk-tier unit tests build CPU banks on a CUDA box, and default_stream() rejects + # a CPU device outright. There is nothing to order for a CPU->CPU copy anyway. + device = self._banks[0][1].device + if device.type != "cuda": + return # CPU banks: the copies are synchronous CPU->CPU + torch.cuda.default_stream(device).synchronize() def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None, phase: str = "prefill") -> None: From 795659de07fdbafb08ab1214a5c945b1b57cba6b Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:58:33 +0000 Subject: [PATCH 33/38] chore(moe): disk-tier review cleanups (PR #337 items 4-5) - gate the layer<3 prefill debug print behind FT_DISK_TIER_DEBUG instead of printing unconditionally. - drop the per-layer device->host sync: cache.usage[disk] now takes the 0-d cache.step tensor directly (same dtype/device) instead of .item(). - validate all --moe-disk-tier v0 preconditions at once and raise a single ValueError listing every unmet flag (each used to cost a full boot to discover); list the exact flags in --moe-disk-tier --help. E2E on Rudi (freetoken:standalone-mtz = current tree, Qwen3.6-35B-A3B-NVFP4, TP=1, GPU0): RAM=64 4296/4296 verify match, RAM=32 4902/4902 verify match, 0 mismatches; decode 11.69 tok/s median at RAM=32 (11.85 at RAM=64 hist). --- python/freetoken/engine/engine.py | 14 ++++++++++---- python/freetoken/moe/disk_tier.py | 8 +++++--- python/freetoken/server/args.py | 6 ++++-- 3 files changed, 19 insertions(+), 9 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index f06fe6b1c..038f7a264 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -570,17 +570,23 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: from freetoken.moe.disk_tier import DiskTierSpec E = config.model_config.num_experts + # Collect ALL unmet preconditions and raise once: each used to surface as a + # separate boot-time ValueError, costing a full boot per missing flag. + problems = [] if not 0 < config.expert_ram_experts < E: - raise ValueError( + problems.append( f"--expert-ram-experts must be in (0, {E}) with --moe-disk-tier on") if decode_target != "gpu": - raise ValueError( + problems.append( "--moe-disk-tier v0 requires the gpu decode path (--moe-backend offload)") if config.moe_prefill_overlap: - raise ValueError("--moe-disk-tier v0 requires --disable-moe-prefill-overlap") + problems.append("--moe-disk-tier v0 requires --disable-moe-prefill-overlap") if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs >= 1: - raise ValueError( + problems.append( "--moe-disk-tier v0 requires --cuda-graph-max-bs 0 (cuda graphs disabled)") + if problems: + raise ValueError( + "--moe-disk-tier on: unmet preconditions:\n - " + "\n - ".join(problems)) disk_tier = DiskTierSpec(ram_experts=config.expert_ram_experts) if cache_factory is None: # Fast path: an FTW checkpoint loads its repacked banks directly. diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index b5c7f1453..0e56fd319 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -540,12 +540,14 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N _materialize_layer_gpu(cache, layer_id, materialize_count=self._ram) routed = expert_ids.reshape(-1) disk = torch.unique(routed[routed >= self._ram]) - if layer_id < 3: + if os.environ.get("FT_DISK_TIER_DEBUG") and layer_id < 3: print(f"[disk-tier dbg] layer={layer_id} routed={routed.numel()} " f"unique_disk={disk.numel()} disk={disk.tolist()[:12]}", flush=True) if disk.numel() == 0: return - step = int(cache.step.item()) # already incremented by the kernel + # cache.step was already incremented by the kernel; assign the 0-d tensor + # device-side (same dtype/device as usage) instead of .item()-ing it, which + # would sync the stream once per layer on the prefill/decode path. futures = [ self._pool.submit(self._fetch_expert, layer_id, int(e), int(e)) for e in disk.tolist() @@ -561,7 +563,7 @@ def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> N flat = layer_id * cache.num_experts + disk cache.slot_for_id[layer_id, disk] = disk cache.id_of_slot[disk] = flat - cache.usage[disk] = step + cache.usage[disk] = cache.step def fetch_pending(self, cache, layer_id: int) -> None: """Fetch this layer's disk-resident misses into their slots; shrink the miss diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 375e49ea3..affcda2b9 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -554,8 +554,10 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help=( "NVMe tier for MoE experts (see moe/disk_tier.py): experts beyond " "--expert-ram-experts per layer stay on disk in the original checkpoint " - "and are fetched on slot-cache miss. Requires native NVFP4 banks, " - "--moe-backend offload, --disable-moe-prefill-overlap and no cuda graphs." + "and are fetched on slot-cache miss. Requires native NVFP4 banks. " + "v0 preconditions (all enforced at once at boot): --moe-backend offload " + "(gpu decode), --disable-moe-prefill-overlap, --cuda-graph-max-bs 0, " + "and 0 < --expert-ram-experts < num_experts." ), ) parser.add_argument( From b602e6d5534149527a597bca8d75517934b0a3fb Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:37:45 +0000 Subject: [PATCH 34/38] test(moe): mincore startup check that the released bank tails stay unbacked Implements the follow-up MT-z proposed in the PR #337 review: the lazy-tail invariant (nothing reads/writes rows [K, E) after release_bank_tails, so they cost no RAM) is now checked rather than assumed. - tail_resident_bytes(): mincore(2) over one bank's tail byte range. - check_tail_unbacked(): runs in both NVFP4 loaders right after release_bank_tails; logs 'tail check rank=r/n: resident X MiB of Y MiB' and warns above the THP bound (one 2 MiB huge page per bank layer -- shmem_enabled=always|force can back the prefix/tail boundary as a huge page; more than that means something touched the tail). - test_tail_unbacked_after_release: small bank, fill prefix, release tail, assert zero resident pages in [K, E); one tail write backs exactly one page. Verified on Rudi (shmem_enabled=always -- the config MT-z flagged as the risky one): real boot at RAM=32 logs 'resident 0 MiB of 15172 MiB', no warning; E2E 4902/4902 slot verifies match, 11.59 tok/s median. --- python/freetoken/models/nvfp4_banks.py | 6 ++- python/freetoken/moe/disk_tier.py | 61 ++++++++++++++++++++++++++ tests/moe/test_disk_tier.py | 17 +++++++ 3 files changed, 82 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 4386da857..a23adcb38 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -218,9 +218,10 @@ def _load(sink) -> int: with PinPipeline(prefix_rows=K) as pins: placed = _load(pins) if K is not None: - from freetoken.moe.disk_tier import release_bank_tails + from freetoken.moe.disk_tier import check_tail_unbacked, release_bank_tails release_bank_tails(_hb, E, K) + check_tail_unbacked(_hb, E, K) expected = num_layers * rows_per_layer * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" @@ -357,9 +358,10 @@ def _wanted(n: str) -> bool: with PinPipeline(prefix_rows=K) as pins: placed = _load(pins) if K is not None: - from freetoken.moe.disk_tier import release_bank_tails + from freetoken.moe.disk_tier import check_tail_unbacked, release_bank_tails release_bank_tails(_hb, E, K) + check_tail_unbacked(_hb, E, K) expected = num_layers * rows_per_layer * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 0e56fd319..2b50aff01 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -57,6 +57,67 @@ def release_bank_tails(banks_by_name: dict[str, list[HostBank]], num_experts: in row_bytes = bank.nbytes // num_experts bank.release_range(ram_experts * row_bytes, bank.nbytes - ram_experts * row_bytes) + +def tail_resident_bytes(bank: HostBank, num_experts: int, ram_experts: int) -> int: + """Bytes the kernel currently backs in the released tail rows ``[ram_experts, E)``. + + mincore(2) over the tail's byte range: one syscall, per-page residency. + Conservative -- mincore also reports private-anon pages mapped from the + shared zero page (a plain READ of a tail row), so this overcounts what + actually costs RAM (cgroup memory.stat shmem is the real number). + Returns -1 if mincore itself fails.""" + row_bytes = bank.nbytes // num_experts + off = ram_experts * row_bytes + size = bank.nbytes - off + if size <= 0: + return 0 + _PAGE = 4096 + vec = (ctypes.c_ubyte * (size // _PAGE))() + libc = ctypes.CDLL("libc.so.6", use_errno=True) + libc.mincore.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_ubyte)] + if libc.mincore(ctypes.c_void_p(bank.addr + off), size, vec) != 0: + return -1 + return sum(vec) * _PAGE + + +def check_tail_unbacked(banks_by_name: dict[str, list[HostBank]], num_experts: int, + ram_experts: int) -> None: + """Startup sanity check for the lazy-tail invariant, right after + ``release_bank_tails``: nothing reads or writes the tail rows, so the kernel + should be backing ~none of them. The expected worst case is one 2 MiB THP + huge page per bank layer (shmem_enabled=always|force can back the + prefix/tail boundary as a huge page); more than that means something touched + the tail and the disk-tier RAM math no longer holds. Always logs, warns + above the bound.""" + _HUGE = 2 << 20 + resident = tail = n_banks = 0 + for layer_banks in banks_by_name.values(): + for bank in layer_banks: + row_bytes = bank.nbytes // num_experts + t = bank.nbytes - ram_experts * row_bytes + if t <= 0: + continue + r = tail_resident_bytes(bank, num_experts, ram_experts) + if r < 0: + continue # mincore failed: skip rather than warn on our own probe + resident += r + tail += t + n_banks += 1 + if n_banks == 0: + return + from freetoken.distributed import try_get_tp_info + tp = try_get_tp_info() + rank = getattr(tp, "rank", "?") + size_ = getattr(tp, "size", "?") + bound = n_banks * _HUGE + print(f"[disk-tier] tail check rank={rank}/{size_}: resident {resident >> 20} MiB " + f"of {tail >> 20} MiB (warn bound {bound >> 20} MiB = 1x2MiB per bank layer)", + flush=True) + if resident > bound: + print(f"[disk-tier] WARNING: tail rows more resident than the THP bound -- " + f"something is reading the released rows; the disk-tier RAM math no longer holds", + flush=True) + # Native NVFP4 bank order (== _BANK_SCHEMAS["nvfp4"]) and, per bank, the # checkpoint segments that make up one expert row: (proj, kind, dst_row_start, # dst_row_end). The gate|up-fused banks splice gate rows then up rows on the diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index 96e28b407..acea3361d 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -298,3 +298,20 @@ def resident_pages(addr, nbytes): assert resident_pages(bank.addr, size) == 0 # The mapping stays valid: the refaulted pages read back as zeros. assert bank.tensor[0] == 0 + + +def test_tail_unbacked_after_release(): + """Lazy-tail invariant (the disk-tier RAM math): after release_range, the + tail rows [K, E) back NO pages -- until something writes them. mincore over + the tail is the cheap startup check check_tail_unbacked() runs for real.""" + from freetoken.moe.disk_tier import release_bank_tails, tail_resident_bytes + + E, K = 8, 4 + bank = HostBank((E, 4096), torch.uint8) # page-sized rows + bank.tensor[:K].fill_(1) # touch only the prefix + release_bank_tails({"b": [bank]}, E, K) + assert tail_resident_bytes(bank, E, K) == 0 + # The mapping still works: one tail write backs exactly one page (the + # invariant is "nothing touches the tail", not "the tail refuses to back"). + bank.tensor[K].fill_(2) + assert tail_resident_bytes(bank, E, K) == 4096 From a6bd5c07a6d3c16a1fdfb4633993114c396c711c Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:45:37 +0000 Subject: [PATCH 35/38] fix(moe): apply MT-z PR #337 round-3 review findings - host_banks: bank mmap is now MAP_PRIVATE|MAP_ANONYMOUS (CPython's default MAP_SHARED is backed by an internal shmem object, so release_range's munmap+MAP_FIXED gave back the address range but not the pages -- a touched tail stays charged for the life of the process). release_range is one madvise(MADV_DONTNEED) + an assert that the range does not overlap the cudaHostRegister'd prefix (_pinned_bytes tracking). Patch verbatim from MT-z (issuecomment-5518521155). - offload_cache: gate the [copy-miss] FT_DISK_TIER_VERIFY print on self._disk_tier is not None and not-capturing -- .item()/.cpu() inside a captured CUDA graph crashed every graph-capturing boot with the env var set, tier off (issuecomment-5519070434). - engine: --moe-disk-tier on a family that builds no disk index now raises NotImplementedError instead of silently serving zeroed experts (issuecomment-5520770513 item 1). - disk_tier: release_bank_tails warns and skips banks whose row boundary is not page-aligned (odd --expert-ram-experts used to AssertionError at load); init probe reads expert 0, not expert 100 (items 2 and 4). - qwen3_5_moe: the default fp8_block path raises instead of silently dropping the disk-tier flag (item 3). - tests: regression test for the unaligned row boundary; updated the release-range test docstring (MAP_PRIVATE, not the private remap). --- python/freetoken/engine/engine.py | 8 +++ python/freetoken/models/qwen3_5_moe/weight.py | 5 ++ python/freetoken/moe/disk_tier.py | 19 +++++- python/freetoken/moe/host_banks.py | 60 +++++++++---------- python/freetoken/moe/offload_cache.py | 4 +- tests/moe/test_disk_tier.py | 41 ++++++++++++- 6 files changed, 98 insertions(+), 39 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 038f7a264..1e1269433 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -658,6 +658,14 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"disk tier: {banks.disk_ram_experts}/{config.model_config.num_experts} " f"experts/layer pinned in RAM; the rest fetched from " f"{config.model_path} on slot-cache miss") + elif disk_tier is not None: + # The loader released the tail rows but this family built no disk + # index (only qwen3_5_moe and re-exporters attach one): without a + # fetcher the released rows would serve as zeros -- wrong logits, + # no error. Fail loudly instead. + raise NotImplementedError( + "--moe-disk-tier on: this model family builds no disk index; experts " + "[K, E) are released at load and would never be refetched") cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) else: cache = cache_factory(config, self.device) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index e5cd98411..424c40015 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -877,6 +877,11 @@ def setup_offload_expert_banks( return banks if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only") + if disk_tier is not None: + # the fp8_block path builds no disk index; without this the flag would be + # accepted and silently dropped (nothing released, nothing fetched) + raise NotImplementedError( + f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") from freetoken.moe.expert_banks import ExpertBanks mode = os.environ.get("FREETOKEN_FP8_EXPERTS", "fp8").strip().lower() diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py index 2b50aff01..959ce319a 100644 --- a/python/freetoken/moe/disk_tier.py +++ b/python/freetoken/moe/disk_tier.py @@ -51,11 +51,24 @@ class DiskTierSpec: def release_bank_tails(banks_by_name: dict[str, list[HostBank]], num_experts: int, ram_experts: int) -> None: - """MADV_DONTNEED the unpinned tail rows of every bank layer (post-load).""" + """MADV_DONTNEED the unpinned tail rows of every bank layer (post-load). + + The release is an optimization, not an invariant: the tail rows were never + written at load, so when a row boundary is not page-aligned (the small scale + banks) we warn and skip that bank instead of failing the boot.""" + _PAGE = 4096 for layer_banks in banks_by_name.values(): for bank in layer_banks: row_bytes = bank.nbytes // num_experts - bank.release_range(ram_experts * row_bytes, bank.nbytes - ram_experts * row_bytes) + offset = ram_experts * row_bytes + size = bank.nbytes - offset + if offset % _PAGE or size % _PAGE: + print(f"[disk-tier] WARNING: bank row boundary not page-aligned " + f"(ram_experts={ram_experts}, row_bytes={row_bytes}); skipping " + f"the release for this bank -- the tail rows were never written, " + f"so nothing is lost", flush=True) + continue + bank.release_range(offset, size) def tail_resident_bytes(bank: HostBank, num_experts: int, ram_experts: int) -> int: @@ -261,7 +274,7 @@ def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int tp = try_get_tp_info() host_shapes = [tuple(b[0][0][0].shape) for b in self._banks] disk_bytes = [ - sum(nb for _, _, nb in self._index.row_segments(bi, 0, 100)) + sum(nb for _, _, nb in self._index.row_segments(bi, 0, 0)) for bi in range(len(self._banks)) ] print(f"[disk-tier-init] tp_rank={getattr(tp, 'rank', '?')}/{getattr(tp, 'size', '?')} " diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 388893c5c..08217ed14 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -79,7 +79,7 @@ class HostBank: 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", "_locked") + __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned", "_pinned_bytes", "_locked") def __init__(self, shape: tuple[int, ...], dtype: torch.dtype, *, backing: str | None = None): @@ -104,11 +104,19 @@ def __init__(self, shape: tuple[int, ...], dtype: torch.dtype, self.addr = raw.data_ptr() + off assert self.addr % _BLK == 0 self._pinned = True # born pinned+mapped; pin() is a no-op + self._pinned_bytes = asize else: - self._buf = mmap.mmap(-1, asize) # lazy: address space only, no resident pages yet + # MAP_PRIVATE, not CPython's default MAP_SHARED: on a shared anonymous mapping a + # *read* fault allocates a page (no zero-page sharing) and MADV_DONTNEED is ignored, + # so an untouched region is only free by convention and a freed one never comes back. + # Private anonymous gives both for real: reads map the shared zero page, and + # release_range() actually returns memory. Nothing needs the mapping to be shared -- + # the loaders are thread pools and ranks are mp-spawned, each with its own banks. + self._buf = mmap.mmap(-1, asize, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS) _LIVE_BUFFERS.append(self._buf) self.addr = ctypes.addressof(ctypes.c_char.from_buffer(self._buf)) self._pinned = False + self._pinned_bytes = 0 self.tensor = torch.frombuffer(self._buf, dtype=dtype, count=self.nbytes // elsize).view(*shape) self._locked = False @@ -140,6 +148,7 @@ def pin(self) -> None: f"cudaHostRegister failed for {len(self._buf) / 2**30:.1f} GiB" ) from exc self._pinned = True + self._pinned_bytes = len(self._buf) def pin_prefix(self, nrows: int) -> None: """Pin only the first ``nrows`` rows (disk tier: the rest stays disk-resident). @@ -160,44 +169,31 @@ def pin_prefix(self, nrows: int) -> None: f"cudaHostRegister failed for {nbytes / 2**30:.1f} GiB prefix" ) from exc self._pinned = True + self._pinned_bytes = nbytes def release_range(self, offset: int, nbytes: int) -> None: - """Free a byte range of the backing mapping by replacing it IN PLACE with a - fresh MAP_PRIVATE anonymous mapping at the same virtual address. - - HostBank's buffer is a MAP_SHARED /dev/zero mapping (CPython's - ``mmap(-1)``), and the kernel silently ignores MADV_DONTNEED on shared - mappings -- the pages would stay resident. Replacing the range with a - private zero mapping frees them while keeping every existing pointer - and torch view valid (same address). The range must be page-aligned - and must not overlap a pinned prefix (the disk tier's unpinned tails). - """ - import ctypes + """Free a byte range of the backing mapping with MADV_DONTNEED. + + The bank is a MAP_PRIVATE anonymous mapping, so dropping a range frees the + pages outright and a later read faults the shared zero page again; every + existing pointer and torch view stays valid (the mapping is never replaced). - _BLK = 4096 + The range must be page-aligned and must not overlap the pinned prefix: + dropping pages under a cudaHostRegister'd range corrupts silently, so it is + asserted here rather than left to the caller (the disk tier's unpinned tails). + """ assert offset % _BLK == 0 and nbytes % _BLK == 0, ( "release_range: page-aligned range required") - libc = ctypes.CDLL("libc.so.6", use_errno=True) - libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t] - libc.munmap.restype = ctypes.c_int - libc.mmap.restype = ctypes.c_void_p - libc.mmap.argtypes = [ - ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long] - addr = self.addr + offset - if libc.munmap(addr, nbytes) != 0: - raise OSError(ctypes.get_errno(), "munmap failed") - PROT_READ_WRITE = 3 - MAP_PRIVATE_ANON = 0x22 # MAP_PRIVATE | MAP_ANONYMOUS - MAP_FIXED = 0x10 - MAP_FAILED = (1 << 64) - 1 - new_addr = libc.mmap(addr, nbytes, PROT_READ_WRITE, MAP_PRIVATE_ANON | MAP_FIXED, -1, 0) - if new_addr in (None, MAP_FAILED): - raise OSError(ctypes.get_errno(), "mmap(MAP_FIXED) failed") - assert new_addr == addr, "MAP_FIXED returned a different address" + assert offset >= self._pinned_bytes, ( + f"release_range: [{offset}, {offset + nbytes}) overlaps the pinned prefix " + f"[0, {self._pinned_bytes})") + if nbytes: + self._buf.madvise(mmap.MADV_DONTNEED, offset, nbytes) def release(self) -> None: """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.""" + For buffers that are done being read (the converter). No-op for born-pinned banks: registered pages cannot be dropped. + (This frees memory only because the mapping is MAP_PRIVATE; the kernel ignores MADV_DONTNEED on shared ones.)""" if self._pinned: return self._buf.madvise(mmap.MADV_DONTNEED) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 17e220bbe..e12f4fd37 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -1032,7 +1032,9 @@ def copy_missing(self) -> None: for per_layer, cache in self.banks: cache[: self.num_experts].copy_(per_layer[layer_id]) return - if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0: + if (self._disk_tier is not None and layer_id == 0 + and os.environ.get("FT_DISK_TIER_VERIFY") + and not torch.cuda.is_current_stream_capturing()): print(f"[copy-miss] layer={layer_id} fused={self._copy_fused_ok} " f"n={int(self.num_indices.item())} " f"evict={self.evict_slots[:4].cpu().tolist()} " diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py index acea3361d..aea439d14 100644 --- a/tests/moe/test_disk_tier.py +++ b/tests/moe/test_disk_tier.py @@ -275,9 +275,9 @@ def test_fetch_pending_all_disk_clears_list(checkpoint): def test_release_range_frees_pages(): - """release_range must actually drop the resident pages (MAP_SHARED /dev/zero - mappings ignore MADV_DONTNEED, so the in-place private remap is verified - with mincore).""" + """release_range must actually drop the resident pages. The bank is a + MAP_PRIVATE anonymous mapping, so MADV_DONTNEED frees for real; mincore + verifies the pages are gone (a MAP_SHARED mapping would keep them).""" import ctypes as ct size = 4 * 1024 * 1024 @@ -315,3 +315,38 @@ def test_tail_unbacked_after_release(): # invariant is "nothing touches the tail", not "the tail refuses to back"). bank.tensor[K].fill_(2) assert tail_resident_bytes(bank, E, K) == 4096 + + +def test_release_bank_tails_unaligned_row_boundary(): + """A row boundary that is not page-aligned (the small scale banks) must not + fail the boot: release_bank_tails warns and skips that bank instead of + asserting in release_range. The tail rows were never written, so skipping + loses nothing. Aligned boundaries still release.""" + import ctypes as ct + + from freetoken.moe.disk_tier import release_bank_tails + + libc = ct.CDLL("libc.so.6", use_errno=True) + libc.mincore.argtypes = [ct.c_void_p, ct.c_size_t, ct.POINTER(ct.c_ubyte)] + libc.mincore.restype = ct.c_int + + def resident_pages(addr, nbytes): + vec = (ct.c_ubyte * ((nbytes + 4095) // 4096))() + assert libc.mincore(addr, nbytes, vec) == 0 + return sum(1 for b in vec if b & 1) + + # A real Ornith gate_up_scale row size: 2048 bytes/row, NOT page-aligned. + E, K = 256, 127 + bank = HostBank((E, 2048), torch.uint8) + bank.tensor.fill_(7) # fault every page in + assert resident_pages(bank.addr, bank.nbytes) == bank.nbytes // 4096 + # K=127: offset = 127*2048 = 259072, not % 4096 -> warn+skip, no AssertionError. + release_bank_tails({"gate_up_scale": [bank]}, E, K) + # Skipped: the (already resident) tail pages are untouched, not freed. + assert resident_pages(bank.addr, bank.nbytes) == bank.nbytes // 4096 + + # Aligned K on the same bank shape: offset = 128*2048 = 262144 (% 4096) -> releases. + bank2 = HostBank((E, 2048), torch.uint8) + bank2.tensor.fill_(7) + release_bank_tails({"gate_up_scale": [bank2]}, E, 128) + assert resident_pages(bank2.addr + 128 * 2048, bank2.nbytes - 128 * 2048) == 0 From 0f4c26d6faa807676db5bc3a1b71ff18c7a87abe Mon Sep 17 00:00:00 2001 From: CraigStone-Dev <321209063+CraigStone-Dev@users.noreply.github.com> Date: Thu, 3 Sep 2026 23:57:16 +0000 Subject: [PATCH 36/38] docs(args): --expert-ram-experts help notes the page-alignment rule Per MT-z (PR #337 issuecomment-5527183880): which K values release the tail cleanly is a per-model arithmetic rule decided by the smallest bank row (the fp32 global scales) -- K * min(row_bytes) must be a multiple of 4096, else the small scale bands stay resident (warns, does not abort). Qwen3.8-Flash-Next needs a multiple of 8; Ornith-1.5-35B a multiple of 2. --- python/freetoken/server/args.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index affcda2b9..89ea3b005 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -566,7 +566,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.expert_ram_experts, help=( "With --moe-disk-tier on: experts per layer kept pinned in RAM " - "(0 < N < num_experts; the rest are disk-resident)." + "(0 < N < num_experts; the rest are disk-resident). Keep " + "N * (smallest bank row bytes) page-aligned (a multiple of 4096) or " + "the small scale banks' tail rows stay resident instead of released " + "(warns, does not abort). The rule is per-model: e.g. Qwen3.8-Flash-Next " + "needs a multiple of 8, Ornith-1.5-35B a multiple of 2." ), ) parser.add_argument( From 7915ac19fbe82b38069bee5b66c27c6b3c09c3ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Thu, 3 Sep 2026 14:58:10 +0900 Subject: [PATCH 37/38] feat(moe): build the disk index for every NVFP4 family, not just qwen3_5_moe The loader releases expert rows [K, E) for every NVFP4 family, but the disk index was constructed in one place -- qwen3_5_moe's setup_offload_expert_banks. Every other family (glm5_next, glm4_moe, gemma4, minimax_m2, minimax_m3, deepseek_v4, ...) reached that release with disk_index=None, so the engine's attach was skipped, the "disk tier: K/E pinned" line never printed, and copy_missing fed the grouped GEMM zeroed expert rows. Wrong logits, no error, no diagnostic. Each family already owns the Nvfp4ExpertSourceSpec its loader reads with, so expose it (nvfp4_expert_source_spec, re-exported from each package like the loaders are) and build the index in the shared provider, which is also where the release happens. glm5_next keeps choosing between its compressed-tensors and modelopt namings per checkpoint. Guards, so the silent path cannot come back: - _nvfp4_banks raises when a family exposes no spec (it would release and never refetch). - the engine raises when the tier is on and the provider returned no index, which also covers the non-NVFP4 providers (bf16, ds_fp4, q4_0). - qwen3_5_moe checks "nvfp4 experts only" before its fp8_block branch instead of inside the nvfp4 one, where the default fp8_block path silently dropped --moe-disk-tier. Verified on an RTX 4090: the index now builds for GLM-5.3-Flash-NVFP4 (glm5_next, 288 experts, 42 MoE layers, 118 shards) and its rows are byte-identical to the checkpoint tensors read through safetensors -- packed weights, per-row scales and the fp32 global scales, 9 of 9 spot checks. Ornith-1.5 (qwen3_5_moe) is unchanged end to end. Co-Authored-By: Claude Fable 5.1 --- python/freetoken/engine/engine.py | 11 ++-- python/freetoken/models/gemma4/__init__.py | 3 +- python/freetoken/models/gemma4/weight.py | 9 +++ python/freetoken/models/glm4_moe/__init__.py | 3 +- python/freetoken/models/glm4_moe/weight.py | 9 +++ python/freetoken/models/glm5_next/__init__.py | 3 +- python/freetoken/models/glm5_next/weight.py | 9 +++ .../freetoken/models/minimax_m2/__init__.py | 3 +- python/freetoken/models/minimax_m2/weight.py | 9 +++ .../freetoken/models/minimax_m3/__init__.py | 3 +- python/freetoken/models/minimax_m3/weight.py | 9 +++ .../freetoken/models/qwen3_5_moe/__init__.py | 3 +- python/freetoken/models/qwen3_5_moe/weight.py | 27 ++++---- python/freetoken/models/qwen4_exp/__init__.py | 3 +- python/freetoken/models/qwen4_exp/weight.py | 9 +++ python/freetoken/models/weight.py | 14 ++++ python/freetoken/moe/expert_banks.py | 21 +++++- tests/moe/test_disk_tier_families.py | 64 +++++++++++++++++++ 18 files changed, 186 insertions(+), 26 deletions(-) create mode 100644 tests/moe/test_disk_tier_families.py diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 1e1269433..e103c6a3b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -659,13 +659,12 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"experts/layer pinned in RAM; the rest fetched from " f"{config.model_path} on slot-cache miss") elif disk_tier is not None: - # The loader released the tail rows but this family built no disk - # index (only qwen3_5_moe and re-exporters attach one): without a - # fetcher the released rows would serve as zeros -- wrong logits, - # no error. Fail loudly instead. + # The loader released experts [K, E) but no fetcher came back: serving would + # multiply by zeroed rows and log nothing. Fail where the flag was set. raise NotImplementedError( - "--moe-disk-tier on: this model family builds no disk index; experts " - "[K, E) are released at load and would never be refetched") + "--moe-disk-tier on: this checkpoint's expert provider returned no disk " + "index, so experts released at load would never be refetched " + f"(quant_format={banks.quant_format!r})") cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) else: cache = cache_factory(config, self.device) diff --git a/python/freetoken/models/gemma4/__init__.py b/python/freetoken/models/gemma4/__init__.py index 37a8dcb39..5ee8b87a5 100644 --- a/python/freetoken/models/gemma4/__init__.py +++ b/python/freetoken/models/gemma4/__init__.py @@ -12,11 +12,12 @@ from .weight import ( iter_weights, iter_weights_parallel, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, ) -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Gemma4Attention", "Gemma4ForCausalLM", "Gemma4MultimodalEmbedder", diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index ad9c43eb2..af4b8f705 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -275,6 +275,15 @@ def _expert_name(raw_name: str) -> str | None: yield _expert_name(raw_name), tensor + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, list[torch.Tensor]]: diff --git a/python/freetoken/models/glm4_moe/__init__.py b/python/freetoken/models/glm4_moe/__init__.py index a07134c34..14db4d12e 100644 --- a/python/freetoken/models/glm4_moe/__init__.py +++ b/python/freetoken/models/glm4_moe/__init__.py @@ -1,8 +1,9 @@ from .config import parse_config from .model import Glm4MoeForCausalLM from .weight import iter_weights, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel +from .weight import nvfp4_expert_source_spec -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Glm4MoeForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py index a00005ec2..3f983d0e3 100644 --- a/python/freetoken/models/glm4_moe/weight.py +++ b/python/freetoken/models/glm4_moe/weight.py @@ -191,6 +191,15 @@ def _iter_resident_weights(reader, config, primary) -> Iterator[tuple[str, torch # -------------------------------------------------------------------------------------- # Routed expert host banks (NVFP4) for the offload cache. # -------------------------------------------------------------------------------------- + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None) -> dict[str, torch.Tensor]: """Build the pinned CPU NVFP4 banks for GLM-4's routed experts. diff --git a/python/freetoken/models/glm5_next/__init__.py b/python/freetoken/models/glm5_next/__init__.py index c1c05140a..927b7a9e8 100644 --- a/python/freetoken/models/glm5_next/__init__.py +++ b/python/freetoken/models/glm5_next/__init__.py @@ -1,8 +1,9 @@ from .config import parse_config from .model import Glm5NextForCausalLM from .weight import iter_weights, load_nvfp4_expert_sources +from .weight import nvfp4_expert_source_spec -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Glm5NextForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/glm5_next/weight.py b/python/freetoken/models/glm5_next/weight.py index 2e3b7f684..8a349a351 100644 --- a/python/freetoken/models/glm5_next/weight.py +++ b/python/freetoken/models/glm5_next/weight.py @@ -95,6 +95,15 @@ def _select_expert_source_spec(model_path: str) -> Nvfp4ExpertSourceSpec: _KDA_IN_PROJ = ("q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj") + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _select_expert_source_spec(model_path) + def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None): # disk_tier is threaded through the same way qwen3_5_moe does it: the caller builds the # NVMe tier and every NVFP4 family's loader has to pass it down, or the tail of the bank diff --git a/python/freetoken/models/minimax_m2/__init__.py b/python/freetoken/models/minimax_m2/__init__.py index e4e1211ce..2f39f6bcb 100644 --- a/python/freetoken/models/minimax_m2/__init__.py +++ b/python/freetoken/models/minimax_m2/__init__.py @@ -1,8 +1,9 @@ from .config import parse_config from .model import MiniMaxM2ForCausalLM from .weight import iter_weights, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel +from .weight import nvfp4_expert_source_spec -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "MiniMaxM2ForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/minimax_m2/weight.py b/python/freetoken/models/minimax_m2/weight.py index a010225bf..aedba6bf4 100644 --- a/python/freetoken/models/minimax_m2/weight.py +++ b/python/freetoken/models/minimax_m2/weight.py @@ -88,6 +88,15 @@ def raw() -> Iterator[tuple[str, torch.Tensor]]: yield from iter_merged_tensors(raw(), _MERGE_RULES, model_name="minimax_m2") + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( model_path: str, config, diff --git a/python/freetoken/models/minimax_m3/__init__.py b/python/freetoken/models/minimax_m3/__init__.py index 71c0bb51f..926aeca96 100644 --- a/python/freetoken/models/minimax_m3/__init__.py +++ b/python/freetoken/models/minimax_m3/__init__.py @@ -2,11 +2,12 @@ from .model import MiniMaxM3ForCausalLM from .weight import ( iter_weights, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, ) -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "MiniMaxM3ForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/minimax_m3/weight.py b/python/freetoken/models/minimax_m3/weight.py index bfcfbbd87..1b96c3b68 100644 --- a/python/freetoken/models/minimax_m3/weight.py +++ b/python/freetoken/models/minimax_m3/weight.py @@ -248,6 +248,15 @@ def iter_weights( reader.close() + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, list[torch.Tensor]]: diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 98936e9f2..589fc0ea7 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -3,12 +3,13 @@ from .weight import ( iter_weights, iter_weights_parallel, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, setup_offload_expert_banks, ) -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Qwen3_5MoEForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 424c40015..a27f7c552 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -856,6 +856,11 @@ def setup_offload_expert_banks( :class:`~freetoken.moe.disk_tier.Nvfp4DiskIndex` so the offload cache can fetch the disk-resident experts on miss.""" eq = getattr(model_config, "expert_quant", "none") + # Checked before the branch below, not inside it: the fp8_block path never passes + # disk_tier on, so a guard inside the nvfp4 branch silently accepted the flag there. + if disk_tier is not None and eq != "nvfp4": + raise NotImplementedError( + f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") if eq != "fp8_block": from freetoken.moe.expert_banks import _PROVIDERS # nvfp4 -> _nvfp4_banks, none -> _bf16_banks @@ -863,17 +868,8 @@ def setup_offload_expert_banks( parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target, layer_sink=layer_sink, disk_tier=disk_tier) - if disk_tier is not None: - if eq != "nvfp4": - raise NotImplementedError( - f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") - import dataclasses - - from freetoken.moe.disk_tier import Nvfp4DiskIndex - - index = Nvfp4DiskIndex(model_path, model_config, _NVFP4_SOURCE_SPEC) - banks = dataclasses.replace( - banks, disk_index=index, disk_ram_experts=disk_tier.ram_experts) + # The disk index rides back on the banks: _nvfp4_banks builds it for every family + # through nvfp4_expert_source_spec, so there is nothing family-specific left here. return banks if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only") @@ -1083,6 +1079,15 @@ def _load(sink) -> None: return ExpertBanks("bf16", banks, streamed=layer_sink is not None) + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, torch.Tensor]: diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py index 04c29d778..e18b007b9 100644 --- a/python/freetoken/models/qwen4_exp/__init__.py +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -13,6 +13,7 @@ from .model import Qwen4ExpForCausalLM from .weight import ( iter_weights, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, load_ple_table, @@ -24,7 +25,7 @@ # load_nvfp4_expert_sources via the model spec. from freetoken.models.qwen3_5_moe.weight import setup_offload_expert_banks -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Qwen4ExpForCausalLM", "iter_weights", "load_nvfp4_expert_sources", diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index edd38f4b2..476929f08 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -289,6 +289,15 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None) -> dict: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" return load_nvfp4_expert_source_banks( diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 5c10e92a3..075904c9c 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -293,6 +293,20 @@ def load_moe_expert_sources( ) +def nvfp4_moe_expert_source_spec(model_path: str, model_config): + """The family's ``Nvfp4ExpertSourceSpec`` for this checkpoint, or None if it has none. + + The disk tier indexes the same rows the loader placed, so it needs the very spec the + family loader used -- which only the family knows (glm5_next picks between a + compressed-tensors and a modelopt naming per checkpoint). None means the family has no + NVFP4 expert source spec at all, which the caller must treat as "no disk tier here" + rather than silently continue. + """ + _config, spec = _spec_for_model_path(model_path) + getter = _model_override(spec, "nvfp4_expert_source_spec") + return None if getter is None else getter(model_path, model_config) + + def load_nvfp4_moe_expert_sources( model_path: str, model_config, diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index ba99f45b8..13bc4f38d 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -192,6 +192,23 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, raise NotImplementedError( "disk tier requires the native NVFP4 layout (triton backend, decode_target=gpu, " "not dummy)") + # Build the index HERE, not in a per-family setup: the loader below releases expert rows + # [K, E) for every family, so a family that reaches the release without an index would + # serve zeroed experts. Resolving the spec through the family hook keeps the index and + # the loader reading the same rows. + disk_kw = {} + if disk_tier is not None: + from freetoken.models.weight import nvfp4_moe_expert_source_spec + from freetoken.moe.disk_tier import Nvfp4DiskIndex + + source_spec = nvfp4_moe_expert_source_spec(model_path, model_config) + if source_spec is None: + raise NotImplementedError( + f"--moe-disk-tier on: {type(model_config).__name__} exposes no " + "nvfp4_expert_source_spec, so the tier cannot locate expert rows in the " + "checkpoint (the loader would release them and never refetch)") + disk_kw = {"disk_index": Nvfp4DiskIndex(model_path, model_config, source_spec), + "disk_ram_experts": disk_tier.ram_experts} repack_sink = None if not native and not dummy and layer_sink is not None: @@ -213,7 +230,7 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, # marlin/b12x repacks (which only the GPU W4A16 kernels can read). if decode_target == "cpu": return ExpertBanks("nvfp4", {name: sources[name] for name in _BANK_SCHEMAS["nvfp4"]}, - streamed=sink is not None) + streamed=sink is not None, **disk_kw) # Pick the expert-GEMM backend by compute capability (and MoE width: auto keeps # small-I MoE on the Triton M=1 GEMV, which beats b12x's tensor cores at single-stream # decode) and repack the banks (in place; the tiled blocks are byte-identical per @@ -221,7 +238,7 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, logger.info(f"NVFP4 expert backend: {backend}") if backend == "triton": return ExpertBanks("nvfp4", {name: sources[name] for name in _BANK_SCHEMAS["nvfp4"]}, - streamed=sink is not None) + streamed=sink is not None, **disk_kw) quant_format = f"nvfp4_{backend}" if repack_sink is not None: # Streamed conversion: each layer was already repacked + written by the wrapper. diff --git a/tests/moe/test_disk_tier_families.py b/tests/moe/test_disk_tier_families.py new file mode 100644 index 000000000..c7a0f349e --- /dev/null +++ b/tests/moe/test_disk_tier_families.py @@ -0,0 +1,64 @@ +"""Every NVFP4 family must hand the disk tier a source spec. + +The loader releases expert rows ``[K, E)`` for every family, so a family that reaches +that release without an index would serve zeroed experts silently. These tests pin the +hook that keeps the index and the loader reading the same rows. +""" +from __future__ import annotations + +import importlib + +import pytest + +# Families whose loader passes an Nvfp4ExpertSourceSpec to load_nvfp4_expert_source_banks. +NVFP4_FAMILIES = [ + "qwen3_5_moe", "qwen4_exp", "glm4_moe", "glm5_next", "gemma4", "minimax_m2", "minimax_m3", +] + + +@pytest.mark.parametrize("family", NVFP4_FAMILIES) +def test_family_exposes_a_source_spec_hook(family): + mod = importlib.import_module(f"freetoken.models.{family}.weight") + getter = getattr(mod, "nvfp4_expert_source_spec", None) + assert callable(getter), ( + f"{family} defines _NVFP4_SOURCE_SPEC but exposes no nvfp4_expert_source_spec, so " + "the shared provider cannot build a disk index for it") + + +@pytest.mark.parametrize("family", [f for f in NVFP4_FAMILIES if f != "glm5_next"]) +def test_hook_returns_the_spec_the_loader_uses(family): + # glm5_next is excluded here only because its hook reads the checkpoint config to pick + # between the compressed-tensors and modelopt namings; it is covered by the test below. + mod = importlib.import_module(f"freetoken.models.{family}.weight") + spec = mod.nvfp4_expert_source_spec("unused/for/these/families", None) + assert spec is mod._NVFP4_SOURCE_SPEC + assert spec.key_pattern.groupindex.keys() >= {"layer", "expert", "proj", "kind"} + assert set(spec.proj_to_role.values()) == {"gate", "up", "down"} + + +def test_glm5_next_hook_follows_the_checkpoint_quant_method(monkeypatch): + mod = importlib.import_module("freetoken.models.glm5_next.weight") + + class _Cfg: + def __init__(self, method): + self.quantization_config = {"quant_method": method} + + monkeypatch.setattr(mod, "cached_load_hf_config", lambda path: _Cfg("compressed-tensors")) + assert mod.nvfp4_expert_source_spec("p", None) is mod._NVFP4_CT_SOURCE_SPEC + monkeypatch.setattr(mod, "cached_load_hf_config", lambda path: _Cfg("modelopt")) + assert mod.nvfp4_expert_source_spec("p", None) is mod._NVFP4_SOURCE_SPEC + + +def test_provider_refuses_a_family_without_a_spec(monkeypatch): + """A family with no hook must fail at load, not release rows and serve zeros.""" + from freetoken.moe import expert_banks + from freetoken.moe.disk_tier import DiskTierSpec + + monkeypatch.setattr("freetoken.models.weight.nvfp4_moe_expert_source_spec", + lambda path, config: None) + # select_nvfp4_backend is reached before the spec lookup; decode_target="cpu" keeps the + # path "native" so the earlier NotImplementedError does not mask the one under test. + with pytest.raises(NotImplementedError, match="nvfp4_expert_source_spec"): + expert_banks._nvfp4_banks( + "does/not/matter", object(), None, None, False, + decode_target="cpu", disk_tier=DiskTierSpec(ram_experts=1)) From c3abff6d84ad2cfff54b3798ca0d6689d4fc1403 Mon Sep 17 00:00:00 2001 From: pi agent Date: Fri, 4 Sep 2026 04:57:14 +0000 Subject: [PATCH 38/38] test(moe): regression test for the [copy-miss] probe's disk-tier gate The probe's .item()/.cpu() syncs crash CUDA graph capture when FT_DISK_TIER_VERIFY is left set with the tier off (PR #337 issuecomment-5519070434). Verified on Rudi (test_offload.py 25/25); this test was developed against a6bd5c0 but never committed. --- tests/moe/test_offload.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 422ca8675..3b346c1d3 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -868,3 +868,22 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +def test_copy_miss_verify_probe_gated_on_disk_tier(monkeypatch, capsys): + """FT_DISK_TIER_VERIFY must not fire the [copy-miss] probe when the disk tier + is off: the probe's .item()/.cpu() are device-to-host syncs, which crash any + CUDA graph capture (PR #337 issuecomment-5519070434 -- every graph-capturing + boot died with the env var left over from a tier session). Gated on the tier + like its neighbours, and skipped while a stream is capturing.""" + layer, cache = _make_layer_and_cache() + cache._pending_src_layer = 0 + cache.evict_slots = torch.tensor([0, 1], dtype=torch.int32) + cache.src_indices = torch.tensor([2, 3], dtype=torch.int32) + cache.num_indices = torch.tensor(2) + cache._copy_fused_ok = False + monkeypatch.setattr("freetoken.kernel.fast_index_copy_jit", lambda *a, **k: None) + monkeypatch.setenv("FT_DISK_TIER_VERIFY", "1") + + cache.copy_missing() + assert "[copy-miss]" not in capsys.readouterr().out