Skip to content

feat(moe): bench hybrid fetch fractions per moe layer kind - #734

Merged
jason-fxz merged 2 commits into
mainfrom
feat/benchbw-layer-kinds
Oct 18, 2026
Merged

jason-fxz merged 2 commits into
mainfrom
feat/benchbw-layer-kinds

Conversation

@jason-fxz

Copy link
Copy Markdown
Collaborator

Part of #671.

Summary

Hybrid decode takes its fetch split from the ft bench bw profile. On main the profile has one entry per expert format (bf16, nvfp4, fp8_block, mxfp4, ds_fp4) and none for GGUF, so every GGUF model runs hybrid at a fixed fetch cap of 1. A model with several expert formats uses one fraction for all layers.

This PR stores one profile row per MoE layer kind and sets the fetch split per layer. ft bench bw --model <checkpoint> benches every layer kind of a GGUF or HF checkpoint. On Gemma-4 26B-A4B UD-Q8_K_XL GGUF, hybrid decode goes from 96.5 to 112.1 tokens/s.

Design

Row key:

<cpu format>[<options>]:<hidden>x<moe_intermediate>:top<top_k>:t<cpu threads>
gguf[down_type=8,gate_up_type=8]:2816x704:top8:t12
bf16:2816x704:top8:t12
  • Format and options come from the layer's quant_method.cpu_spec(). The thread count is the CPU executor's worker count; CPU bandwidth depends on it and on top_k.
  • ft bench bw --model <checkpoint> builds the model on the meta device, collects its layer kinds and benches each on synthetic banks. --dtype and the --model presets still work and write rows at their fixed shapes.
  • A row stores the CPU MoE and PCIe gather bandwidths, alone and overlapped, plus metadata outside the key (ISA, cores, activation, sources, threshold, time, commit).
  • A run replaces the rows it benched and keeps the rest.
  • At startup each hybrid layer looks up its row and fetches pcie_ov / (pcie_ov + cpu_ov) of its misses. A layer without a row keeps the fixed cap of 1, with a warning.

Breaking changes

  • Profile version 6 holds only layer_kinds. Older profiles are ignored; rerun ft bench bw.
  • --moe-strategy auto no longer switches to hybrid from the profile. Use --moe-strategy hybrid.
  • FreeToken Desktop reads the old dtype_kernels and workloads fields and needs an update.

Tested on

H100 80GB HBM3, Xeon Platinum 8480+ (13 cores, one NUMA node), driver 580.95.05, CUDA 13.0, torch 2.14.1+cu130, Triton 3.8.0, based on main 31c6ef4.

Checkpoints: unsloth/gemma-4-26B-A4B-it-GGUF (UD-Q8_K_XL: 29 layers Q8_0, 1 layer BF16), nvidia/Gemma-4-26B-A4B-NVFP4-GGUF, ggml-org/gemma-4-26B-A4B-it-MXFP4-GGUF, Qwen/Qwen3-30B-A3B.

ft bench bw --dtype all --model /path/to/gemma-4-26B-A4B-it-UD-Q8_K_XL.gguf,/path/to/Gemma-4-26B-A4B-NVFP4.gguf,/path/to/Qwen3-30B-A3B
ft bench bw --model /path/to/gemma-4-26B-A4B-it-MXFP4_MOE.gguf
ft serve --model /path/to/gemma-4-26B-A4B-it-UD-Q8_K_XL.gguf --moe-strategy hybrid --moe-cache-mem 4GB --max-running-requests 1 --cuda-graph-max-bs 1

Results

AIME25 problem 1 through the chat API, temperature 1.0, top-p 0.95, top-k 64, 5 runs per arm. Gemma-4 26B-A4B UD-Q8_K_XL, --moe-strategy hybrid --moe-cache-mem 4GB (660 of 3840 experts cached). Fixed cap 1 is what main uses for GGUF.

Fetch split Decode tokens/s, median (min - max) Answer 70
fixed cap 1 96.5 (95.6 - 97.8) 5 / 5
per layer: 25.7% on 29 Q8_0 layers, cap 1 on the BF16 layer 112.1 (110.2 - 114.1) 5 / 5

Bench rows (GB/s, 12 CPU threads):

Layer kind CPU MoE PCIe gather Overlapped CPU / PCIe Fraction Recommended
Gemma-4 UD-Q8_K_XL, Q8_0 101.6 31.5 89.1 / 31.1 25.9% hybrid
Gemma-4 UD-Q8_K_XL, BF16 120.6 31.6 107.1 / 32.1 23.0% hybrid
Gemma-4 MXFP4_MOE 68.2 31.7 60.4 / 31.5 34.3% hybrid
Gemma-4 NVFP4 GGUF 57.0 31.6 53.5 / 32.3 37.7% offload
Qwen3-30B-A3B, BF16 121.1 32.4 106.5 / 32.0 23.1% hybrid

Other checks

  • Servers find a row for every layer: all 30 on Gemma-4 UD-Q8_K_XL (two kinds), all 48 on Qwen3-30B-A3B.
  • Re-benching one checkpoint replaced only its two rows.
  • GGUF MXFP4 blocks start with an e8m0 scale byte; zero-filled banks made it a denormal and the CPU bench read 6.2 GB/s. The bench now writes a unit scale, as it does for mxfp4_triton and ds_fp4.
  • PYTHONPATH=python pytest tests/moe/test_hybrid_fetch.py tests/moe/test_offload.py tests/moe/test_mixed_expert_formats.py tests/checkpoint tests/engine/test_cache_budget.py tests/server --deselect tests/moe/test_offload.py::test_graph_capture_reuses_warm_offload_cache_before_capture: 708 passed, 7 skipped. The deselected test also fails on main.

Not covered

  • --moe-strategy auto does not choose hybrid per layer yet; that needs each layer's expert quant in ModelConfig.
  • With TP > 1 no layer finds a row: the engine keys on the per-rank intermediate size, and the bench runs at TP 1.
  • Mixed-kind checkpoints were tested on one file only.
  • Only H100 was tested. ROCm was not run.

The ft bench bw profile (version 6) holds only layer_kinds rows, keyed by CPU format, options, expert shape, top_k and CPU worker threads; older profiles are ignored. ft bench bw --model also takes checkpoint paths and benches each of their layer kinds. Hybrid decode reads each layer's own row for its fetch fraction, else keeps a fixed cap of 1. --moe-strategy auto no longer upgrades to hybrid from the profile.
@andy-yang-1

Copy link
Copy Markdown
Collaborator

The mechanism works and I found no bug. Two design choices can silently cost hybrid users decode speed, though, and the new engine path has no test.

Checked (H100, Xeon 8480+, torch 2.9.0+cu128):

  • When the fraction is > 0, the hybrid kernel replaces max_fetch with it, so the engine's hybrid_max_fetch = 1 is safe.
  • executor.num_threads is the worker count after the coordinator core is reserved, which matches the bench's len(core_ids).
  • Nothing else reads hybrid_fetch_fraction as a float. The gguf_scaled banks match the layout.
  • ft bench bw --model on 5 local checkpoints (Qwen3-30B-A3B, gemma-4-26B-A4B, Qwen3.6-35B-A3B, Gemma-4-26B-A4B-NVFP4, Qwen3.6-35B-A3B-FP8) read every layer kind and wrote one row each in 33 s. Re-runs merge rows correctly.
  • With a bf16 + NVFP4 model and a profile that has only the bf16 row, the bf16 layers get their fraction, the NVFP4 layers fall back to cap 1, and the warning is clear.
  • tests/moe/test_hybrid_fetch.py tests/moe/test_offload.py tests/moe/test_mixed_expert_formats.py tests/checkpoint tests/engine tests/server: 798 passed and the same 2 failures as on the base 31c6ef4f.

1. An exact key match is required, so most layers end up at cap 1. _resolve_hybrid_fetch uses a row only if format, options, shape, top_k and threads all match; otherwise the layer gets cap 1. The PR's own run puts cap 1 at -14% (96.5 vs 112.1 tok/s). On my box the thread count moves the fraction a lot, but the shape barely does:

bf16 rows fraction
2048x768 / 2816x704 / 2048x512, all at 12 threads 34.1% / 33.9% / 35.8%
2048x768 at 12 / 8 / 6 threads 34.1% / 39.0% / 45.4%

Your table shows the same: bf16 is 23.0% at 2816x704 and 23.1% at 2048x768. So keeping threads in the key is right, but matching on the shape costs a lot:

  • Plain ft bench bw (canonical shapes) writes rows that almost no real model matches; bf16 2048x768 happens to be Qwen3-30B-A3B's shape.
  • After upgrading, a user whose old profile gave nvfp4 models a fraction drops to cap 1 until they re-bench, and only --model <checkpoint> brings it back.

Suggestion: when there is no exact row, fall back to the row with the same format, options and threads and the nearest shape, and log it at info level.

2. --moe-strategy auto no longer picks hybrid. Users who ran the bench silently go from hybrid to offload. If per-layer auto has to wait, a cheap middle ground is a hint after the model is built ("the profile recommends hybrid for these layers; pass --moe-strategy hybrid").

3. No tests for the new paths. The updated tests cover only the profile reader. Nothing tests the per-layer mapping and cap-1 fallback in _resolve_hybrid_fetch, _merge_previous, or checkpoint_layer_kinds. A fake executor plus a temp profile (FREETOKEN_BENCHBW_PATH) covers the first one in a few lines.

Nits

  • _model_list accepts only local paths and presets. A hub id such as Qwen/Qwen3-30B-A3B is rejected, while ft serve accepts it.
  • With the default thread count, the bench reserves a coordinator core (13 cores -> t12). A server started with an explicit --moe-cpu-threads 13 keys t13 and finds no row. The docs could say: leave both on auto, or pass the same explicit value to both.

@jason-fxz jason-fxz changed the title feat(moe): bench hybrid fetch fractions per moe layer kind feat(moe): bench hybrid fetch fractions per moe layer kind Oct 18, 2026
@jason-fxz
jason-fxz merged commit 8f49e0a into main Oct 18, 2026
2 checks passed
pretor pushed a commit to pretor/FreeToken that referenced this pull request Oct 18, 2026
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.

2 participants