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. ❌ 1 New Failure, 1 Unclassified FailureAs of commit c0c3322 with merge base 822cb60 ( NEW FAILURE - The following job has failed:
UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:
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.
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. It uses Metal Performance Primitives (Metal 4) for INT8 x INT4 -> INT32 accumulation.int4b_format.full_k=False("staged"): one fused kernel that quantizes activations and unpacks weights into threadgroup memory one K tile at a time, then calls MPPmatmul2d.full_k=True: a per-token quantize kernel, then one full-K MPPint8 x int4b_formatGEMM that reads packed weights straight from device memory.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).BK=-1means the full-K lowering. Every config is checked against a reference, and the fastest one that passes is reported.quantized_matmul, 4-bit affine, group size 32.How to run
Requirements:
executorch.backends.mlximports).pip install mlx. The M4 run below used MLX 0.32.4.This takes about 11-16 minutes on M4 Max, including a 60 s pause between fp16 and bf16. Please post
int8_int4_bench.log.Useful flags:
--m 4 8 16 256 512: token batch sizes.--shapes qkv o gate_up down lm_head--dtypes fp16 bf16--bm 8 16 32 64 128: output tile rows. MPP requires a multiple of 8.--bn 32 64 128: output tile columns.--bk 32 64 128 -1: K tile size;-1means full K.--group-size 32: MLX baseline group size.--warmup 1: untimed runs per config/shape/M before timing, so kernels are compiled first.--iters 10: timed runs per config/shape/M.--dtype-pause 60: seconds to pause between fp16 and bf16.--verbose: print every tile config.How to read a progress line:
[bf16] qkv M=256 mlx 2355.2 us | best full_k 32x128 2166.4 us, 1.09x vs mlx.full_k BMxBNmeans the full-K kernel with that output tile.staged BMxBNxBKmeans the K-tiled kernel.M4 Max results (fp16 activations, MLX 0.32.4, default flags)
Estimated total Muse Glimmer GEMM time per forward, fp16:
Summary:
oat M=4 has measured anywhere from 64 to 231 us on the same M4 Max. Treat the small-M MLX baselines, and the speedups computed from them, as approximate.lm_headat M=8 takes 11.3 ms, against 4.1 ms at M=4.mlx_metal_small_batch_quantized.patch) may behave differently from pip MLX at these sizes.Full fp16 output