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
- hybrid fetch split: upstream per-layer-kind bench rows (FlashML-org#734) first, our startup split now measured per layer kind without a row (before the CPU executor), then cap 1; dropped the profile-driven auto->hybrid pick and its cache-share guard with upstream - qwen3_5_moe: upstream fp32 runtime norm +1 (FlashML-org#755) under our DFlash hidden-layer return and LinearOProj - qwen4_exp: upstream quant-config reader and MXFP8 scale aliases (FlashML-org#759) replace our .weight_scale rename; our validating scan_ple_table and NVFP4 PLE tables gain the bf16 table support - nvfp4_banks: upstream W4A4 input scales (FlashML-org#736) combined with our TP sharding and disk-tier skip - ftw: fp16 nvfp4 global refusal (FlashML-org#739) on top of our eager index validation - kernel/nvfp4_moe: drop the e2m1 LUT argument our bit-placement decode removed - benchbw: upstream layer-kind rows plus our THP bench banks, decode-shaped gather and custom profile hint - tests: disk-tier spec test skips the quant-config-driven qwen4_exp hook Assisted-by: Claude Opus 5.5
FlashML-org#764 FlashML-org#763 FlashML-org#708 FlashML-org#687 deferred, FlashML-org#688 kept as a feature; FlashML-org#169 kept after its upstream close, FlashML-org#494/FlashML-org#154/FlashML-org#315 replaced upstream by FlashML-org#690/FlashML-org#698; FlashML-org#278 dropped with upstream's auto->hybrid pick (FlashML-org#734). Assisted-by: Claude Opus 5.5
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.