fix(moe): match the benchbw verdict to the served model's expert size

python/freetoken/engine/engine.py

@@ -652,8 +652,11 @@ def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None:
         from freetoken.moe.bench_profile import load_hybrid_fetch_fraction
 
         gpu_name, gpu_uuid = _profile_gpu(self.device.index)
+        model_config = config.model_config
         fraction = load_hybrid_fetch_fraction(
-            cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid
+            cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid,
+            expert_bytes=_model_expert_bytes(cache.quant_format, model_config),
+            geometry=_model_geometry(model_config),
         )
         if fraction is None:
             cache.hybrid_max_fetch = 1
@@ -997,6 +1000,32 @@ def shutdown(self) -> None:
         destroy_distributed()
 
 
+def _model_geometry(model_config) -> dict | None:
+    """The served model's MoE geometry in benchbw's workload terms, or None for a dense model."""
+    try:
+        geometry = {
+            "hidden": int(model_config.hidden_size),
+            "inter": int(model_config.moe_intermediate_size),
+            "experts": int(model_config.num_experts),
+            "top_k": int(model_config.num_experts_per_tok),
+        }
+    except (AttributeError, TypeError, ValueError):
+        return None
+    return geometry if geometry["experts"] > 0 and geometry["inter"] > 0 else None
+
+
+def _model_expert_bytes(bench_fmt: str, model_config) -> int | None:
+    """Per-expert offload-bank bytes for this model in ``bench_fmt`` (the same per-format
+    sizing ``bank_bytes_estimate`` uses for the pin budget), or None for a format without one."""
+    from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT
+
+    geometry = _model_geometry(model_config)
+    per_expert = _BANK_BYTES_PER_EXPERT.get(bench_fmt)
+    if geometry is None or per_expert is None:
+        return None
+    return int(per_expert(geometry["hidden"], geometry["inter"]))
+
+
 def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]:
     """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA."""
     if not torch.cuda.is_available():
@@ -1371,7 +1400,15 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
         from freetoken.moe.bench_profile import load_backend_recommendation
 
         gpu_name, gpu_uuid = _profile_gpu()
-        if load_backend_recommendation(bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid) == "hybrid":
+        # The profile is keyed on the expert format; the model's geometry / expert size picks the
+        # entry that was benched on comparable experts (a small-expert model must not inherit a
+        # verdict measured on 4x larger ones).
+        recommendation = load_backend_recommendation(
+            bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid,
+            expert_bytes=_model_expert_bytes(bench_fmt, model_config),
+            geometry=_model_geometry(model_config),
+        )
+        if recommendation == "hybrid":
             from freetoken.moe.cpu_executor import compiled_extension_supports
 
             _act = getattr(model_config, "hidden_act", "silu")

python/freetoken/moe/bench_profile.py

@@ -5,8 +5,14 @@
 itself. ``benchbw.py`` writes the profile; this module only reads it.
 
 The join key is the expert *format*: the CPU-MoE-vs-PCIe-gather bandwidth ratio the choice
-rides on is dominated by ``(format, hardware)``, not by the exact model, so a profile benched
-on one workload transfers to any model with the same expert format on the same GPU.
+rides on is dominated by ``(format, hardware)``, so a profile benched on one workload transfers
+to other models with the same expert format on the same GPU -- as long as their experts are of
+comparable size. The CPU side has a per-expert fixed cost (thread fan-out over a short GEMV),
+so a verdict measured on much larger experts overstates it for a small-expert model: Qwen3.6-
+35B-A3B-NVFP4 (1.7 MB experts) benches at 32 GB/s where the nvfp4 dtype geometry (8.0 MB)
+gives 90 GB/s on the same box, and hybrid then decodes 3x slower than offload. So a per-model
+bench of the served geometry wins, and a dtype verdict is only applied when its expert size is
+within :data:`EXPERT_BYTES_TOLERANCE` of the model's.
 """
 
 from __future__ import annotations
@@ -31,6 +37,52 @@
 }
 
 
+# A dtype verdict / fetch fraction transfers to a model whose experts are within this factor
+# of the benched ones (either way). Beyond it the entry is skipped and the caller falls back to
+# offload; `ft bench bw --model` for the model's geometry is the remedy.
+EXPERT_BYTES_TOLERANCE = 2.0
+_GEOMETRY_KEYS = ("hidden", "inter", "experts", "top_k")
+_warned: set = set()
+
+
+def _geometry_matches(entry_model, geometry) -> bool:
+    if not isinstance(entry_model, dict) or not geometry:
+        return False
+    return all(entry_model.get(k) == geometry.get(k) for k in _GEOMETRY_KEYS)
+
+
+def _model_entry(prof: dict, fmt: str, geometry) -> dict | None:
+    """``kernels[fmt]`` of the per-model workload benched at exactly ``geometry``, else None."""
+    for wl in (prof.get("workloads") or {}).values():
+        if isinstance(wl, dict) and _geometry_matches(wl.get("model"), geometry):
+            entry = (wl.get("kernels") or {}).get(fmt)
+            if isinstance(entry, dict):
+                return entry
+    return None
+
+
+def _comparable(entry, fmt: str, expert_bytes) -> bool:
+    """Whether ``entry`` (a dtype_kernels / workload kernels entry) was benched on experts of a
+    size comparable to ``expert_bytes``. Unknown on either side -> assumed comparable (older
+    profiles carry no ``expert_bytes``; callers that pass none keep the format-only join)."""
+    benched = entry.get("expert_bytes") if isinstance(entry, dict) else None
+    if not benched or not expert_bytes:
+        return True
+    ratio = max(benched, expert_bytes) / min(benched, expert_bytes)
+    if ratio <= EXPERT_BYTES_TOLERANCE:
+        return True
+    key = (fmt, benched, expert_bytes)
+    if key not in _warned:
+        _warned.add(key)
+        logger.warning(
+            f"benchbw profile: the {fmt!r} entry was benched on {benched / 2**20:.2f} MB experts, "
+            f"this model's are {expert_bytes / 2**20:.2f} MB ({ratio:.1f}x apart); not applying "
+            f"its verdict. Run `ft bench bw --model <preset>` for this model's geometry "
+            f"(see `ft bench bw --help`)"
+        )
+    return False
+
+
 def _cache_dir() -> str:
     cache = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
     return os.path.join(cache, "freetoken")
@@ -116,6 +168,9 @@ def load_backend_recommendation(
     gpu_name: str | None = None,
     path: str | None = None,
     gpu_uuid: str | None = None,
+    *,
+    expert_bytes: int | None = None,
+    geometry: dict | None = None,
 ) -> str | None:
     """Bench-recommended offload-family backend for ``quant_format`` on this GPU, or ``None``.
 
@@ -124,20 +179,30 @@ def load_backend_recommendation(
     near-threshold format) resolves conservatively to ``"offload"``. ``None`` means "no usable
     profile" (see ``_usable_profile``) or no entry for this format. The caller keeps its own
     default (offload) on ``None``.
+
+    ``geometry`` (``{hidden, inter, experts, top_k}``) selects a per-model bench of exactly the
+    served model first; ``expert_bytes`` (the model's per-expert bank bytes) then gates every
+    other entry through :func:`_comparable`.
     """
     fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format)
     prof = _usable_profile(gpu_name, path, gpu_uuid)
     if prof is None:
         return None
 
-    # Preferred: the per-dtype tuning verdicts (`ft bench bw --dtype`), a direct format->backend
-    # map -- the axis the backend pick is meant to key on.
+    # Best: this exact geometry was benched (`ft bench bw --model`).
+    entry = _model_entry(prof, fmt, geometry)
+    if entry is not None and entry.get("recommended") in ("hybrid", "offload"):
+        return entry["recommended"]
+
+    # Next: the per-dtype tuning verdicts (`ft bench bw --dtype`), a direct format->backend
+    # map -- the axis the backend pick keys on -- when benched on comparable experts.
     dtypes = prof.get("dtypes")
     if isinstance(dtypes, dict) and dtypes.get(fmt) in ("hybrid", "offload"):
-        return dtypes[fmt]
+        if _comparable((prof.get("dtype_kernels") or {}).get(fmt), fmt, expert_bytes):
+            return dtypes[fmt]
 
-    # Fallback: a per-model profile (`ft bench bw --model`). Aggregate the workloads sharing this
-    # format -- unanimous hybrid -> hybrid; any offload (a near-threshold split) -> offload.
+    # Fallback: other per-model workloads sharing this format, comparable experts only --
+    # unanimous hybrid -> hybrid; any offload (a near-threshold split) -> offload.
     workloads = prof.get("workloads")
     if not isinstance(workloads, dict):
         return None
@@ -147,6 +212,7 @@ def load_backend_recommendation(
         if isinstance(wl, dict)
         for entry in [(wl.get("kernels") or {}).get(fmt)]
         if isinstance(entry, dict) and entry.get("recommended")
+        and _comparable(entry, fmt, expert_bytes)
     ]
     if not picks:
         return None
@@ -158,6 +224,9 @@ def load_hybrid_fetch_fraction(
     gpu_name: str | None = None,
     path: str | None = None,
     gpu_uuid: str | None = None,
+    *,
+    expert_bytes: int | None = None,
+    geometry: dict | None = None,
 ) -> float | None:
     """Benched hybrid fetch fraction for ``quant_format``, or ``None``.
 
@@ -167,20 +236,21 @@ def load_hybrid_fetch_fraction(
     running concurrently -- the real contention regime): fetched/misses = pcie_ov /
     (pcie_ov + cpu_ov). Older profiles without it fall back to the standalone bandwidths
     under a full-DRAM-contention assumption (cpu keeps cpu - pcie under DMA), which
-    reduces to pcie/cpu. Per-dtype entry first, then any per-model entry with this format.
-    ``None`` = no usable profile; clamped to [0, 1].
+    reduces to pcie/cpu. A per-model entry of the served ``geometry`` first, then the
+    per-dtype entry, then any per-model entry with this format -- each only when benched on
+    experts comparable to ``expert_bytes``. ``None`` = no usable profile; clamped to [0, 1].
     """
     fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format)
     prof = _usable_profile(gpu_name, path, gpu_uuid)
     if prof is None:
         return None
-    entries = [(prof.get("dtype_kernels") or {}).get(fmt)] + [
+    entries = [_model_entry(prof, fmt, geometry), (prof.get("dtype_kernels") or {}).get(fmt)] + [
         (wl.get("kernels") or {}).get(fmt)
         for wl in (prof.get("workloads") or {}).values()
         if isinstance(wl, dict)
     ]
     for entry in entries:
-        if not isinstance(entry, dict):
+        if not isinstance(entry, dict) or not _comparable(entry, fmt, expert_bytes):
             continue
         cpu_ov, pcie_ov = entry.get("cpu_moe_overlap_gbs"), entry.get("pcie_gather_overlap_gbs")
         if cpu_ov and pcie_ov:

tests/moe/test_hybrid_fetch.py

@@ -66,6 +66,46 @@ def test_load_hybrid_fetch_fraction(tmp_path):
     assert load_hybrid_fetch_fraction("bf16", gpu_name="OTHER", path=str(path)) is None
 
 
+def test_backend_pick_follows_the_served_expert_geometry(tmp_path):
+    """Numbers from an RTX 4070 SUPER: the nvfp4 dtype bench (3072x1536, 7.97 MB experts) says
+    hybrid at 89.7 GB/s CPU vs 25.2 PCIe, but Qwen3.6-35B-A3B's own geometry (2048x512,
+    1.78 MB) benches at 32.2 GB/s -> offload, and hybrid decodes 3x slower than offload there."""
+    small = {"hidden": 2048, "inter": 512, "experts": 256, "top_k": 8}
+    prof = {
+        "gpu": {"name": "FAKE GPU"},
+        "dtypes": {"nvfp4": "hybrid"},
+        "dtype_kernels": {"nvfp4": {
+            "expert_bytes": 7974912, "cpu_moe_gbs": 89.7, "pcie_gather_gbs": 25.2,
+            "cpu_moe_overlap_gbs": 74.1, "pcie_gather_overlap_gbs": 25.2,
+        }},
+        "workloads": {"qwen3.6-moe": {
+            "model": dict(small),
+            "kernels": {"nvfp4": {
+                "expert_bytes": 1775616, "recommended": "offload", "cpu_moe_gbs": 32.2,
+                "pcie_gather_gbs": 25.9, "cpu_moe_overlap_gbs": 28.2, "pcie_gather_overlap_gbs": 25.4,
+            }},
+        }},
+    }
+    path = tmp_path / "benchbw.json"
+    path.write_text(json.dumps(prof))
+    p = str(path)
+
+    # a bench of exactly this geometry beats the dtype verdict, for the pick and the split
+    assert load_backend_recommendation("nvfp4", path=p, expert_bytes=1775616, geometry=small) == "offload"
+    assert load_hybrid_fetch_fraction("nvfp4", path=p, expert_bytes=1775616, geometry=small) == pytest.approx(25.4 / (25.4 + 28.2))
+    # callers that pass no geometry keep the format-only join
+    assert load_backend_recommendation("nvfp4", path=p) == "hybrid"
+
+    # no bench of this geometry: a verdict from 4.5x larger experts is not applied ...
+    del prof["workloads"]
+    path.write_text(json.dumps(prof))
+    assert load_backend_recommendation("nvfp4", path=p, expert_bytes=1775616, geometry=small) is None
+    assert load_hybrid_fetch_fraction("nvfp4", path=p, expert_bytes=1775616, geometry=small) is None
+    # ... while comparable experts (within 2x) inherit it as before
+    assert load_backend_recommendation("nvfp4", path=p, expert_bytes=6_000_000, geometry={"hidden": 3072, "inter": 1024, "experts": 128, "top_k": 8}) == "hybrid"
+    assert load_hybrid_fetch_fraction("nvfp4", path=p, expert_bytes=6_000_000) == pytest.approx(25.2 / (25.2 + 74.1))
+
+
 def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch):
     monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path))
     monkeypatch.delenv("FREETOKEN_BENCHBW_PATH", raising=False)