Skip to content

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

Open
bernimccoy wants to merge 1 commit into
FlashML-org:mainfrom
house-of-estel:fix/benchbw-geometry
Open

fix(moe): match the benchbw verdict to the served model's expert size#278
bernimccoy wants to merge 1 commit into
FlashML-org:mainfrom
house-of-estel:fix/benchbw-geometry

Conversation

@bernimccoy

@bernimccoy bernimccoy commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

Fixes the 3x decode regression of --moe-backend auto on small-expert NVFP4 models
(Qwen3.6-35B-A3B-NVFP4 on an RTX 4070 SUPER: auto -> hybrid 27 tok/s, offload 84 tok/s).

Problem

load_backend_recommendation keys the profile on the expert format only and prefers the
ft bench bw --dtype verdict, measured on one geometry per format (nvfp4: 3072x1536, 7.6 MiB
experts). The CPU MoE kernel does not scale down to small experts: the same box benches
89.7 GB/s there and 32.2 GB/s at Qwen3.6-35B-A3B's 2048x512 (1.7 MiB), below the 2x PCIe
threshold. So auto picks hybrid and decodes 3x slower than offload. 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.

Fix

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; 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 12 GB (sm_89), driver 591.86, WSL2, torch 2.11.0+cu130, 24-core host, 108 GB
RAM. nvidia/Qwen3.6-35B-A3B-NVFP4, --moe-cache-auto (3831 slots), bs=1.

benchmarks/bench_decode_moe.py --backend offload,cpu,hybrid (256 decode tokens):

backend decode tok/s
offload 83.7
cpu 25.9
hybrid (what auto picked on main) 27.2

ft bench bw: dtype nvfp4 89.7 GB/s CPU vs 25.3 PCIe (hybrid); --model qwen3.6-moe: 32.2
vs 25.9 (offload). With this change and the same profile, ft serve (auto) logs
benchbw profile: the 'nvfp4' entry was benched on 7.61 MB experts, this model's are 1.69 MB (4.5x apart); not applying its verdict, picks offload, and streams 91.0 tok/s 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 (exact-geometry
entry wins; a 4.5x-off dtype verdict is skipped; a 1.3x-off one still applies; no-geometry
callers unchanged); 6/6 pass. bf16 Qwen3-30B-A3B (9.4 MiB experts) is within 2x of the bf16
dtype geometry, so its verdict is unchanged.

Caveats

The 2x tolerance is a judgment call from one box (the kernel ratio here was 2.8x between the
two geometries); a per-model bench stays the precise answer and the warning points at it.

Assisted by: Claude Code

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant