feat(moe): estimate expert-bank bytes from the model config

python/freetoken/engine/engine.py

@@ -535,10 +535,12 @@ class Engine:
         )
         if config.moe_backend == "cpu" and not split_residency:
             # cpu mode pins every bank for the prefill double buffer; over the pin cap that dies in cudaHostRegister, so lock everything instead
-            from freetoken.moe.expert_banks import ftw_bank_bytes
+            from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes
 
             budget = _pin_budget_bytes()
-            bank_bytes = ftw_bank_bytes(config.model_path) if budget is not None else None
+            bank_bytes = None
+            if budget is not None:
+                bank_bytes = ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config)
             if bank_bytes and bank_bytes > budget:
                 split_residency = True
                 logger.info_rank0(
@@ -1159,9 +1161,9 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[int
     """Pick CPU (locked) MoE layers automatically when the banks exceed the pin budget.
 
     Locks just enough head+tail layers: per-layer decode miss rates are U-shaped, so the ends are the cheapest to move off the slot cache."""
-    from freetoken.moe.expert_banks import ftw_bank_bytes
+    from freetoken.moe.expert_banks import bank_bytes_estimate, ftw_bank_bytes
 
-    bank_bytes = ftw_bank_bytes(config.model_path)
+    bank_bytes = ftw_bank_bytes(config.model_path) or bank_bytes_estimate(config.model_config)
     if not bank_bytes:
         return frozenset()
     budget = _pin_budget_bytes()

python/freetoken/moe/expert_banks.py

@@ -24,7 +24,7 @@ import torch
 
 from freetoken.utils import init_logger
 
-from .offload_cache import _BANK_SCHEMAS
+from .offload_cache import _BANK_BYTES_PER_EXPERT, _BANK_SCHEMAS
 
 logger = init_logger(__name__)
 
@@ -386,6 +386,25 @@ def ftw_bank_bytes(model_path: str) -> int | None:
     return sum(t["nbytes"] for t in tensors if t.get("kind") == "experts_bank")
 
 
+def bank_bytes_estimate(model_config) -> int | None:
+    """Estimated total expert-bank bytes of a raw checkpoint, from the model config alone.
+
+    Sizes the pin-budget decisions where FTW metadata is not available; ``None`` for unknown formats or missing dims (callers then skip the pre-load sizing).
+    nvfp4 uses the native-row formula, a slight over-estimate for the repacked backends."""
+    expert_quant = getattr(model_config, "expert_quant", "none")
+    fmt = expert_quant if expert_quant != "none" else (
+        getattr(model_config, "moe_weight_format", None) or "bf16"
+    )
+    per_expert = _BANK_BYTES_PER_EXPERT.get(fmt)
+    layers = getattr(model_config, "num_moe_layers", None)
+    experts = getattr(model_config, "num_experts", None)
+    hidden = getattr(model_config, "hidden_size", None)
+    inter = getattr(model_config, "moe_intermediate_size", None)
+    if per_expert is None or not all((layers, experts, hidden, inter)):
+        return None
+    return layers * experts * per_expert(hidden, inter)
+
+
 def load_expert_banks(
     model_path: str,
     model_config,

python/freetoken/moe/offload_cache.py

@@ -77,6 +77,17 @@ _BANK_SCHEMAS: dict[str, tuple[str, ...]] = {
     "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"),
 }
 
+# bytes per (expert, layer) as f(hidden, moe_intermediate), from the bank shapes above; keep in sync with _BANK_SCHEMAS
+# keyed by the config-time format tag (expert_quant / moe_weight_format), not quant_format: "mxfp4" sizes the mxfp4_triton banks, "nvfp4" also covers its repacked variants
+_BANK_BYTES_PER_EXPERT = {
+    "bf16": lambda H, I: 3 * I * H * 2,
+    "fp8_block": lambda H, I: 3 * I * H + ((2 * I // 128) * (H // 128) + (H // 128) * (I // 128)) * 2,
+    "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18,
+    "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2),
+    "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2),
+    "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32),
+}
+
 # vLLM's marlin grouped-GEMM hands the full [cache_size] slot cache as its expert
 # dimension; moe_align_block_size requires round_up(experts, 32) < 1024, i.e. <= 992.
 MARLIN_MAX_CACHE_SIZE = 992