Repository navigation
feat(moe): bench hybrid fetch fractions per moe layer kind - #734
Conversation
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.
|
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):
1. An exact key match is required, so most layers end up at cap 1.
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:
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. 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 Nits
|
…l), PR FlashML-org#734 (bench hybrid fetch), PR FlashML-org#732, PR FlashML-org#731
Part of #671.
Summary
Hybrid decode takes its fetch split from the
ft bench bwprofile. 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:
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.--dtypeand the--modelpresets still work and write rows at their fixed shapes.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
layer_kinds. Older profiles are ignored; rerunft bench bw.--moe-strategy autono longer switches tohybridfrom the profile. Use--moe-strategy hybrid.dtype_kernelsandworkloadsfields 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.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.Bench rows (GB/s, 12 CPU threads):
Other checks
mxfp4_tritonandds_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 autodoes not choose hybrid per layer yet; that needs each layer's expert quant inModelConfig.