feat(moe): NVMe disk tier for MoE expert banks

python/freetoken/engine/config.py

@@ -49,6 +49,13 @@ class EngineConfig:
     # CPU MoE backend (--moe-strategy cpu): number of CPU worker threads computing
     # the decode experts. 0 = auto (physical cores). Ignored by other backends.
     moe_cpu_threads: int = 0
+    # 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
     # Hybrid CPU/GPU decode (--moe-strategy offload only): which MoE layers decode on
     # the CPU executor instead of the GPU offload/PCIe path. Spec is an explicit id
     # list ("3,7,11"), a count ("8" -> 8 layers evenly strided across depth), or a

python/freetoken/engine/engine.py

@@ -544,6 +544,29 @@ 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
+            # 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:
+                problems.append(
+                    f"--expert-ram-experts must be in (0, {E}) with --moe-disk-tier on")
+            if decode_target != "gpu":
+                problems.append(
+                    "--moe-disk-tier v0 requires the gpu decode path (--moe-strategy offload)")
+            if config.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:
+                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)
         # Fast path: an FTW checkpoint loads its repacked banks directly.
         # Slow path: load_expert_banks auto-picks parallel vs serial baseline by
         # expert-tensor granularity. Both pin-after-fill.
@@ -570,6 +593,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,
             )
         except PinFailed as exc:
             raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc
@@ -619,6 +643,21 @@ 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")
+        elif disk_tier is not None:
+            # 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 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)
         if decode_target == "hybrid":
             self._resolve_hybrid_fetch(config, cache)

python/freetoken/layers/moe.py

@@ -278,6 +278,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,
@@ -364,7 +367,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,

python/freetoken/models/nvfp4_banks.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import collections
+import itertools
 import json
 import os
 import re
@@ -74,13 +75,20 @@ def iter_nvfp4_expert_pieces(
     chunk: int = 8 << 20,
     drop_page_cache: DropPageCache | None = None,
     primary: bool = True,
+    skip_experts_from: int | None = None,
 ):
     """One piece per routed expert: ``gate`` / ``up`` / ``down`` codes plus their ``_scale``
     (fp8 block scales) and ``_global`` (the per-tensor scale, reciprocal for quant-side dialects,
     fp16) companions, straight from the safetensors shards.
 
     Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either
     way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land.
+
+    ``skip_experts_from`` (the disk tier): experts ``[skip_experts_from, E)`` are disk-resident.
+    The serial reader never calls ``get_tensor`` for them (no I/O); the parallel reader filters
+    them at the reader (the whole-shard read is unchanged, but no per-expert work happens). Each
+    skipped expert still yields an EMPTY piece so the bank fill completes the layer without
+    touching the (released) tail rows.
     """
     from freetoken.models.loader import drop_page_cache as _drop
     from freetoken.moe.expert_pieces import per_expert_pieces
@@ -95,6 +103,8 @@ def iter_nvfp4_expert_pieces(
         match = spec.key_pattern.match(name)
         if match is None:
             continue
+        if skip_experts_from is not None and int(match.group("expert")) >= skip_experts_from:
+            continue  # disk-resident: never read
         bank_layer = _bank_layer(spec, int(match.group("layer")), config)
         if bank_layer is None:
             continue
@@ -105,7 +115,8 @@ def iter_nvfp4_expert_pieces(
         if kind not in ("weight", "weight_scale", "weight_scale_2"):
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
         wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind))
-    expected = _num_moe_layers(config) * config.num_experts * 9
+    experts = min(skip_experts_from, config.num_experts) if skip_experts_from is not None else config.num_experts
+    expected = _num_moe_layers(config) * experts * 9
     if len(wanted) != expected:
         raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}")
 
@@ -133,7 +144,19 @@ def _parallel():
                 tensor = _ingest_global(spec, tensor)
             yield name, tensor
 
-    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    pieces = per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    if skip_experts_from is not None:
+        pieces = itertools.chain(pieces, _disk_skip_pieces(config, skip_experts_from))
+    return pieces
+
+
+def _disk_skip_pieces(config, skip_from: int):
+    """Empty pieces for the disk-resident experts ``[skip_from, E)``: the rows are
+    complete without a checkpoint read (``build_expert_banks`` marks them written,
+    pins only the prefix, and releases the tail)."""
+    for bank_layer in range(_num_moe_layers(config)):
+        for e in range(skip_from, config.num_experts):
+            yield bank_layer, e, e + 1, {}
 
 
 __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]