Pull requests / #2645
#2645 Metal: opt-in fp16-slab expert path (mpp) for streaming prompt chunks
open · @catalinb13 · 0 コメント · GitHub で見る
BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsDocumentation
本文
## What it changes The streaming prompt path reads each expert's quantized weights inside every gate/up and down GEMM's K loop. On the M1 Max that measured 2.9 TFLOP/s useful while the matrix units can do 6.2. `STRATA_METAL_PROMPT_MOE=mpp` decodes each batch of experts into fp16 slabs **once per chunk** (with `mm_expand`, which the engine did not dispatch before), then runs gate/up and down as one dynamic-K `matmul2d` each (`mm_gu_mpp_16_128_4`, `mm_dn_mpp_16_64_2`). Two side changes: - `mm_scan` now also writes `eoff[n_expert]` (one more word; the three `eoff` allocations grow by it), and `mm_expand` skips experts with zero rows, so a chunk only decodes what it routes to. - `STRATA_METAL_PROMPT_MOE_BATCH=N` sets the slab batch, default 64 experts (5.12 GB of prompt buffers in the arena region instead of 4.49; 128 = 5.8 GB). It measured faster than 128: the tile-range early exits cost less than the extra dispatches. ## What it leaves alone The default path is untouched and byte-identical (see below). Decode, the window fallback, CUDA/HIP/SYCL, and the dense/GDN/QSA prompt kernels are not touched; only Metal prompt experts change, and only when the variable is set. The paired `mm_gu2` experiment was removed along with its knob; parity's `--gu2p` became `--v2`. ## How I checked Apple M1 Max, 64 GB, GSQ-RCO pack, IQ3_S gate/up and IQ4_NL down, interleaved with same-session v2 anchors (the two anchors repeated within 56 us): - experts-only parity bench: 626 vs 290 tok/s (48 layers) - a 4096-token chunk's GPU time: 12243 -> 8818 ms (-28%); with the 356-token chunk 1695 -> 1488 ms (-12%, the skip) - warm real prefill 4096: 225.9 -> 256.6 tok/s (+13.6%); real 356: 60.0 -> 55.8 tok/s (-7%, the decode pays whole experts, so it stays opt-in) - `moe_mm_parity --mpp`: y vs the exact float reference rel L2 4.9e-04 (IQ2_XS), 4.9e-04 (IQ3_S), 3.2e-04 (Q2_0) at T=256, and PASS at T=8 and T=32 where most experts are untouched; v1 and v2 parity PASS unchanged - default-path outputs byte-identical to the pre-change build (8 tokens, same prompt file) `docs/MAC.md` carries the two switches with these numbers. The probe scripts live at the repo root locally; the result directories (`prefill-mpp2-1015-1702/`) can be published under `bench/results/` if wanted — say so in review.
関連リンク
インストール・モデル・リリースへの站内リンク。