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

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

python/freetoken/engine/engine.py

@@ -565,6 +565,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-backend 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)
         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 +613,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 +650,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)
         else:
             cache = cache_factory(config, self.device)

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,
@@ -397,7 +400,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/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",

python/freetoken/models/gemma4/weight.py

@@ -275,8 +275,17 @@ 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
+    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 +295,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 +314,7 @@ def load_nvfp4_expert_sources_parallel(
         workers=workers,
         chunk=chunk,
         layer_sink=layer_sink,
+        disk_tier=disk_tier,
     )
 
 

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",

python/freetoken/models/glm4_moe/weight.py

@@ -191,7 +191,16 @@ 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 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.
 
     experts exist only for layers [first_k_dense_replace, num_layers) and pack by MoE layer
@@ -205,11 +214,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 +233,7 @@ def load_nvfp4_expert_sources_parallel(
         workers=workers,
         chunk=chunk,
         layer_sink=layer_sink,
+        disk_tier=disk_tier,
     )
 
 

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",

python/freetoken/models/glm5_next/weight.py

@@ -95,14 +95,27 @@ 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 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
+    # is never registered and load fails with an unexpected-kwarg TypeError.
     return load_nvfp4_expert_source_banks(
         model_path,
         config,
         _select_expert_source_spec(model_path),
         drop_page_cache=drop_page_cache,
         primary=get_tp_info().is_primary(),
         layer_sink=layer_sink,
+        disk_tier=disk_tier,
     )
 
 

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",

python/freetoken/models/minimax_m2/weight.py

@@ -88,11 +88,20 @@ 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,
     *,
-    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 +111,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 +130,7 @@ def load_nvfp4_expert_sources_parallel(
         workers=workers,
         chunk=chunk,
         layer_sink=layer_sink,
+        disk_tier=disk_tier,
     )
 
 

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",

python/freetoken/models/minimax_m3/weight.py

@@ -248,8 +248,17 @@ 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
+    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 +268,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 +287,7 @@ def load_nvfp4_expert_sources_parallel(
         workers=workers,
         chunk=chunk,
         layer_sink=layer_sink,
+        disk_tier=disk_tier,
     )