Repository navigation
Conversation
Adds `mlx::int8_int4_gemm`: dynamic symmetric per-token int8 activations x symmetric per-channel int4 weights, INT32 accumulation via Metal Performance Primitives (Metal 4). Two lowerings: * full_k=False: one fused kernel that quantizes activations and unpacks int4 weights into threadgroup memory per BK tile, then runs MPP matmul2d. * full_k=True: per-token quantize kernel, then a full-K MPP `int8 x int4b_format` GEMM reading packed weights straight from device memory. Tile sizes (BM, BN, BK) are Metal template parameters. Adds bench_int8_int4_gemm.py, which sweeps tile configs for every Muse Glimmer GEMM shape, compares the best config against MLX 4-bit affine quantized_matmul (gs=32) for fp16 and bf16, and estimates total Muse Glimmer GEMM time per forward.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24467
Note: Links to docs will display an error until the docs builds have been completed. ❌ 7 New Failures, 1 Cancelled Job, 1 PendingAs of commit d92fee3 with merge base 822cb60 ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
- Default --bm now 8 16 32 64 128; BMs that are not a multiple of 8 are skipped (MPP matmul2d requires M % 8 == 0). - --warmup (default 1): untimed runs per config / shape / M before the timed --iters runs, so kernels are compiled first. - --dtype-pause (default 60s) between activation dtypes.
…fficial re-run - Time each config as a dependent chain (each call consumes the previous output) and subtract the chain-only cost, so calls cannot overlap on the GPU as they cannot between model layers. Independent calls hid per-call latency and overstated small-M speedups. - Pick the tile in two rounds: every config with --iters (10), then the --refine-top (3) fastest per lowering (staged, full-K) with --refine-iters (20). - Re-time the refined winner against MLX alternately --final-repeats (5) times with --final-iters (100) and report the medians; MLX is no longer timed after the sweep.
…; prune bench sweep Kernels (int8_int4_gemm.py): - Quantize kernel: one threadgroup of TPR (256) threads per row instead of one SIMD group; can zero-pad xq to whole BM tiles. Full-K GEMM reads M from x_scale so padded xq works. - Split-K full-K GEMM: SPLITS threadgroups per output tile each cover K/SPLITS and write INT32 partials; a reduce kernel sums them and applies the scales. - Matvec for small M: streams packed weights like MLX's qmv with exact integer accumulation. - Both new paths are bit-identical to the full-K kernel. The handler does not emit them yet. Benchmark (bench_int8_int4_gemm.py): - Sweeps split-K (--splits 1 2 4 8) and matvec (--matvec-max-m 2); refine keeps the top configs of each lowering. - --pad-m and --quant-threads. - Smaller default sweep: staged (BK > 0) is off (never won a serialized run), BM <= max(8, M), M = 4 8 16 256. - Time-based iteration counts: each timing targets --time-budget-ms (50) of GPU time, clamped to [--min-iters, round cap], so large-M timings stay short and the GPU stays cooler.
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Draft / experimental. Not for landing yet. This PR is mainly here so the benchmark can be run on other Apple GPUs, especially M5.
What this adds
mlx::int8_int4_gemm(backends/mlx/custom_kernel_ops/int8_int4_gemm.py): an MLX-backend custom op that multiplies 8-bit dynamically quantized activations by 4-bit per-channel weights, with exact INT32 accumulation.int4b_format.int8 x int4b_formatmatmul over the full K, reading packed weights straight from device memory. BM and BN are template parameters.qmv, using exact integer arithmetic instead of MPP.full_k=False). Not swept by default; it never won.bench_int8_int4_gemm.py(backends/mlx/custom_kernel_ops/test/) is a standalone MLX benchmark covering every Muse Glimmer GEMM shape:qkv,o,gate_up,down(52 occurrences each) andlm_head(1).quantized_matmul, 4-bit affine, group size 32.--m 1 2for decode.How the benchmark measures
x = x0 + 0 * y[:, :1]), then evaluates the chain once. Calls cannot overlap on the GPU, just as consecutive layers cannot in a model. The same chain without the GEMM is timed and subtracted. Timing independent calls instead lets the GPU overlap them, which hides per-call latency and badly overstates small-M results.--time-budget-ms(50 ms) of GPU time, estimated from one call, between--min-iters(3) and a per-round cap. Short kernels (all of M <= 16) run the full cap; large-M timings stay short, which limits heat.--splits 1 2 4 8) tiles with BM <= max(8, M), plus matvec at M <= 2.--iters, 10).--refine-top(3) fastest configs of each lowering (full-K, split-K, matvec) are re-timed (cap--refine-iters, 20), and the fastest is kept.--final-repeats(5) times (cap--final-iters, 100); the medians are reported. MLX is only timed in this round, so both sides are measured under the same GPU conditions.--rotate-mb(512), so weights are not served from cache. int8_int4 times include activation quantization.How to run
Requirements:
executorch.backends.mlximports).pip install mlx. The M4 run below used MLX 0.32.4.Please post
int8_int4_bench.log. Close other GPU-heavy apps while it runs. Drop--dtypes fp16to also run bf16.Useful flags:
--m 1 2 4 8 16 256: token batch sizes.--shapes qkv o gate_up down lm_head--dtypes fp16 bf16--pad-m: zero-pad activations to whole BM tiles in the quantize kernel, so partial tiles read zeros.--bm 8 16 32 64 128,--bn 16 32 64 128: output tile rows and columns. MPP requires BM to be a multiple of 8.--splits 1 2 4 8: split-K factors.--matvec-max-m 2: largest M that also tries the matvec kernel (0 disables it).--bk -1: K tile size;-1means full K. Add e.g.32 64 128to also sweep the staged kernel.--group-size 32: MLX baseline group size.--time-budget-ms 50,--min-iters 3,--iters 10,--refine-top 3,--refine-iters 20,--final-iters 100,--final-repeats 5: see "How the benchmark measures".--warmup 1: untimed runs before each timing, so kernels are compiled first.--dtype-pause 60,--shape-pause 10: seconds to pause between dtypes and between shapes.--verbose: print every config and the refine ranking.How to read a progress line:
[fp16] qkv M=8 split_k8 8x128 int8_int4 124.1 us | mlx 204.0 us | 1.64x (refine 123.0 us, 6 candidates).full_k BMxBNis the full-K kernel with that output tile,split_kS BMxBNis split-K with S splits, andmatvec mrRis the matvec kernel with R activation rows per threadgroup.refineis the winner's time in the refine round, for reference.M4 Max results (fp16 activations, MLX 0.32.4,
--m 1 2 4 8 16 256 --dtypes fp16 --pad-m)Estimated total Muse Glimmer GEMM time per forward:
Summary:
gate_up,downandlm_headstill win (1.13-1.18x), whileqkvandolose (0.85-0.90x). Matvec does the dot product on the ALUs, so its cost grows with M; past M = 2, MPP split-K is faster.gate_up,downandlm_headgain 1.08-1.09x;qkvandoare about even (1.00-1.03x). int8_int4 sustains 11.6-14.1 TOPS against MLX's 10.8-13.2. Split-K by 2 or 4 still wins 4 of 5 shapes.lm_headat M = 256 is at 5%; the shortoM = 2 timing is the outlier). The official round alternates int8_int4 and MLX, so the speedups stay comparable. Treat differences of a few percent as noise.mlx_metal_small_batch_quantized_mlp.patch, which is tuned for these shapes on M4 Max at M = 4 and M = 16. This benchmark uses pip MLX.Full fp16 output