From eebb3f53f357315eff8a92baba2e61ce1336c333 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 23 Aug 2026 23:06:56 +0000 Subject: [PATCH] feat(moe): estimate expert-bank bytes from the model config --- python/freetoken/engine/engine.py | 10 ++++++---- python/freetoken/moe/expert_banks.py | 21 ++++++++++++++++++++- python/freetoken/moe/offload_cache.py | 11 +++++++++++ 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 9ce193a..22c4a6c 100644 --- a/python/freetoken/engine/engine.py +++ b/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() diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 0c8be6a..b6fadff 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/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, diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index b4d1d89..6ee7640 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/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