[Metal] Skip empty NAX GEMM output groups

[XXXXRT666]

Proposed changes

NAX GEMM can dispatch SIMD groups whose output extent is zero for partial M or N tiles. These groups previously continued through the K loop, including tile loads and NAX matrix operations, even though they could not produce or store any output.

This change:

• Detects SIMD groups with no output rows or columns.
• Skips tile loads and matrix operations for empty groups.
• Skips the epilogue and output store when the output extent is empty.

Performance

Measured on an iPhone 17 Pro Max running iOS 27.0 with Metal 4.1. The benchmarks use FP16, NN layout, batch size 1, 10 warmup iterations, and 20 timed iterations.

M x N x K	Before	After	Speedup	Notes
129 x 2305 x 768	0.6692 ms	0.4315 ms	1.55x	Synthetic double-tail stress case. M and N are both one element past 128-aligned tile boundaries
197 x 3072 x 768	0.7415 ms	0.6290 ms	1.18x	ViT-B/16 MLP expansion: 197 tokens projected from 768 to 3072.

Each shape was measured for seven repeats with a three-second pause between repeats. The first result from each build was excluded, and the table reports the median GPU completion time of the remaining six results.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Can you do some benchmarking on aligned shapes to check if there are performance regressions? The added branching could have bad affect.

[XXXXRT666]

I benchmarked two output-aligned regular-NAX shapes on the same
iPhone 17 Pro Max setup (FP16, NN, batch size 1, W10/T20):

Shape	Before	After	Change
256x3072x768	1758.19 GFLOP/s	1754.45 GFLOP/s	-0.21%
1024x1024x1024	1941.77 GFLOP/s	1947.61 GFLOP/s	+0.30%

The first shape used R7 and the second used R5; the first repeat was excluded
and the median of the remaining repeats is reported.

For aligned M/N specializations, the checks in the K loop are also excluded by
if constexpr (!kAlignedM || !kAlignedN), and the epilogue predicate
specializes to true.

For reference, I also forced the NAX path on an M4 MacBook by bypassing MLX's GPU-generation check. I observed similar speedups, but mx.matmul produced incorrect results with both the baseline and PR builds, so these numbers are performance-only:

Shape	Before	After	Speedup
129x2305x768	0.5115 ms	0.3075 ms	1.663x
197x3072x768	0.6319 ms	0.5535 ms	1.142x

Forcing the path uses a NAX lane layout intended for supported newer GPU generations(maybe), which does not match the M4 GPU. Since both builds fail numerically in the same way, the incorrect results are unrelated to this PR.

[zcbenz]

Did a simple benchmarking with:

import mlx.core as mx
import time
import statistics

def benchmark(M, N, K, iters=20, dtype=mx.bfloat16):
    A = mx.random.uniform(shape=(M, K)).astype(dtype)
    B = mx.random.uniform(shape=(K, N)).astype(dtype)
    mx.eval(A)
    mx.eval(B)

    elapsed = []
    for _ in range(iters):
        C = mx.matmul(A, B)
        start = time.perf_counter()
        mx.eval(C)
        elapsed.append(time.perf_counter() - start)

    print(
        f"{M}x{N}x{K} ms: "
        f"p50={statistics.median(elapsed)*1000:.2f}, "
        f"p95={sorted(elapsed)[int(0.95*len(elapsed))]*1000:.2f}, "
        f"max={max(elapsed)*1000:.2f}"
    )

benchmark(4096, 32000, 4096)
benchmark(4096, 4096, 11008)
benchmark(4096, 4096, 12288)
benchmark(4096, 4096, 4096)
benchmark(1024, 2048, 2048)
benchmark(128, 2048, 768)

benchmark(5000, 33000, 4096)
benchmark(5000, 5000, 11008)
benchmark(5000, 5000, 12288)
benchmark(5000, 5000, 4096)
benchmark(1100, 2100, 2048)
benchmark(130, 2000, 768)

The main branch:

4096x32000x4096 ms: p50=16.96, p95=22.06, max=22.06
4096x4096x11008 ms: p50=7.00, p95=7.49, max=7.49
4096x4096x12288 ms: p50=7.27, p95=10.50, max=10.50
4096x4096x4096 ms: p50=2.46, p95=2.95, max=2.95
1024x2048x2048 ms: p50=0.40, p95=0.51, max=0.51
128x2048x768 ms: p50=0.26, p95=0.30, max=0.30
5000x33000x4096 ms: p50=22.40, p95=27.92, max=27.92
5000x5000x11008 ms: p50=11.21, p95=17.71, max=17.71
5000x5000x12288 ms: p50=12.94, p95=13.34, max=13.34
5000x5000x4096 ms: p50=3.74, p95=3.89, max=3.89
1100x2100x2048 ms: p50=0.48, p95=0.53, max=0.53
130x2000x768 ms: p50=0.25, p95=0.31, max=0.31

This branch:

4096x32000x4096 ms: p50=16.99, p95=22.14, max=22.14
4096x4096x11008 ms: p50=7.04, p95=7.41, max=7.41
4096x4096x12288 ms: p50=7.29, p95=10.55, max=10.55
4096x4096x4096 ms: p50=2.46, p95=2.92, max=2.92
1024x2048x2048 ms: p50=0.40, p95=0.50, max=0.50
128x2048x768 ms: p50=0.26, p95=0.32, max=0.32
5000x33000x4096 ms: p50=22.28, p95=27.94, max=27.94
5000x5000x11008 ms: p50=11.12, p95=11.88, max=11.88
5000x5000x12288 ms: p50=12.90, p95=13.19, max=13.19
5000x5000x4096 ms: p50=3.74, p95=3.82, max=3.82
1100x2100x2048 ms: p50=0.46, p95=0.56, max=0.56
130x2000x768 ms: p50=0.26, p95=0.33, max=0.33

So I think it does improve a lot for certain unaligned shapes, while does not seem to regress other shapes.

[zcbenz]

This is only skipping the epilogue, should the whole gemm_loop be skipped too?

[XXXXRT666]

The GEMM work is already skipped inside gemm_loop, but the call itself cannot be skipped per SIMD group because the loop contains a threadgroup_barrier for every aligned-K iteration. An edge threadgroup can contain both useful and empty SIMD groups, so returning before the call would leave the useful groups waiting at the barrier.
Empty groups therefore still enter gemm_loop and participate in each barrier, but immediately continue before tile loads and matmad. The epilogue and store have no threadgroup-wide synchronization, so they can be skipped entirely.

[zcbenz]

With float16 the difference is not significant:

Main branch

4096x32000x4096 ms: p50=16.96, p95=21.91, max=21.91
4096x4096x11008 ms: p50=7.03, p95=7.64, max=7.64
4096x4096x12288 ms: p50=7.28, p95=10.53, max=10.53
4096x4096x4096 ms: p50=2.47, p95=2.96, max=2.96
1024x2048x2048 ms: p50=0.41, p95=0.52, max=0.52
128x2048x768 ms: p50=0.26, p95=0.33, max=0.33
5000x33000x4096 ms: p50=22.40, p95=28.10, max=28.10
5000x5000x11008 ms: p50=11.21, p95=12.17, max=12.17
5000x5000x12288 ms: p50=12.92, p95=13.28, max=13.28
5000x5000x4096 ms: p50=3.76, p95=3.83, max=3.83
1100x2100x2048 ms: p50=0.46, p95=0.56, max=0.56
130x2000x768 ms: p50=0.27, p95=0.45, max=0.45

This branch

4096x32000x4096 ms: p50=16.96, p95=22.26, max=22.26
4096x4096x11008 ms: p50=7.01, p95=7.48, max=7.48
4096x4096x12288 ms: p50=7.27, p95=10.35, max=10.35
4096x4096x4096 ms: p50=2.46, p95=2.92, max=2.92
1024x2048x2048 ms: p50=0.41, p95=0.56, max=0.56
128x2048x768 ms: p50=0.27, p95=0.38, max=0.38
5000x33000x4096 ms: p50=22.25, p95=27.70, max=27.70
5000x5000x11008 ms: p50=11.06, p95=11.83, max=11.83
5000x5000x12288 ms: p50=12.89, p95=13.27, max=13.27
5000x5000x4096 ms: p50=3.72, p95=3.85, max=3.85
1100x2100x2048 ms: p50=0.47, p95=0.61, max=0.61
130x2000x768 ms: p50=0.26, p95=0.40, max=0.40

I'm running on a M5 Max so the saved SIMD groups probably do not make a change.

[XXXXRT666]

The M5 Max result makes sense with the current routing. I reran the two smaller unaligned shapes on an iPhone 17 Pro Max:

Shape / dtype	GPU interval: baseline → PR	Change	Wall: baseline → PR	Change	iPhone work bound
1100x2100x2048 FP16	2.1161 → 2.0292 ms	+4.1%	2.5578 → 2.4380 ms	+4.7%	5.64%
1100x2100x2048 FP32	2.0938 → 2.1273 ms	-1.6%	2.5839 → 2.5117 ms	+2.8%	5.64%
130x2000x768 FP16	0.5643 → 0.3849 ms	+31.8%	1.1672 → 0.9840 ms	+15.7%	38.48%
130x2000x768 FP32	0.7266 → 0.5490 ms	+24.4%	1.3293 → 1.1374 ms	+14.4%*	38.48%

The GPU interval uses command-buffer GPUStartTime/GPUEndTime, not pure pipeline duration. The FP32 wall baseline was noisy, so that result is directional.

I estimated the work-count upper bound as:

G_before = ceil(M/BM) * ceil(N/BN) * WM * WN

G_useful = ceil(M/(BM/WM)) * ceil(N/(BN/WN))

max reduction = 1 - G_useful/G_before

max speedup = G_before/G_useful

For the M5 Max routing (BM=64, BN=128, WM=2, WN=4):

Shape	Max reduction	Max speedup
5000x33000x4096	0.63%	1.006x
5000x5000x11008	2.50%	1.026x
5000x5000x12288	2.50%	1.026x
5000x5000x4096	2.50%	1.026x
1100x2100x2048	5.64%	1.060x
130x2000x768	17.97%	1.219x

The first four M5 Max bounds are only 0.63–2.50%, so no measurable wall-time change is expected. The iPhone routing gives 130x2000x768 a much larger 38.48% bound, consistent with its larger GPU improvement.

[zcbenz]

I'm good with PR since it is not introducing regressions and can help edge devices, but I would need another view before merging.