From 0111d42107c73ccfcac72636f936e2db88d389b4 Mon Sep 17 00:00:00 2001 From: bernimccoy <85797556+bernimccoy@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:47:10 -0400 Subject: [PATCH] fix(moe): match the benchbw verdict to the served model's expert size load_backend_recommendation keys the profile on the expert format only and prefers the ft bench bw --dtype verdict, which is measured on one geometry per format (nvfp4: 3072x1536, 7.97 MB experts). The CPU MoE kernel does not scale down to small experts: an RTX 4070 SUPER box benches 89.7 GB/s there and 32.2 GB/s at Qwen3.6-35B-A3B's 2048x512 (1.78 MB), below the 2x PCIe threshold. So auto picks hybrid and decodes 3x slower than offload (27.2 vs 83.7 tok/s at bs=1). ft bench bw --model qwen3.6-moe already records offload for this geometry, but the reader consults the per-format dtypes entry first and never reaches the per-model one. load_backend_recommendation / load_hybrid_fetch_fraction take the served model's geometry (H, I, E, top_k) and expert bytes: a per-model workload entry benched at that geometry wins; otherwise the dtype entry applies only if its expert_bytes is within EXPERT_BYTES_TOLERANCE (2x) of the model's, other workload entries are filtered the same way, and a skipped entry logs once with the ft bench bw --model remedy. Callers that pass no geometry, and profiles written before expert_bytes was recorded, get today's format-only match. The engine passes both at the backend pick and at --moe-hybrid-max-fetch auto. No profile format change. Tested on RTX 4070 SUPER (sm_89), driver 591.86, WSL2, torch 2.11.0+cu130, nvidia/Qwen3.6-35B-A3B-NVFP4: with the existing dtype profile, auto now logs the mismatch and picks offload: 91.0 tok/s streamed on the request that gave 26.4 tok/s under auto -> hybrid. tests/moe/test_hybrid_fetch.py: new case with the measured profile numbers, 6/6 pass. Assisted-by: Claude Fable 5 --- python/freetoken/engine/engine.py | 41 +++++++++++- python/freetoken/moe/bench_profile.py | 92 +++++++++++++++++++++++---- tests/moe/test_hybrid_fetch.py | 40 ++++++++++++ 3 files changed, 160 insertions(+), 13 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2d..1f4f7a64b 100644 --- a/python/freetoken/engine/engine.py +++ b/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") diff --git a/python/freetoken/moe/bench_profile.py b/python/freetoken/moe/bench_profile.py index d7b720738..7e4a77501 100644 --- a/python/freetoken/moe/bench_profile.py +++ b/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 ` 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: diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index 169f8e2b6..31e7ba8b8 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/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)