[Metal] Gather mm improvement

[nastya236]

Description:

In the current gather_mm kernel, tiles an assigned to thread groups independently from experts assignment. Example:

BM = 16
Expert 0: 13, expert 1: 8, expert 2: 2, expert 3: 10
tiles: [13,3][5,2,9][1] 

This means that for a tile [5,2,9] we do "as many matmuls as possible":
matmul with expert 0, matmul with expert 1, matmul with expert 2

mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h

Line 68 in 073d225

// Do as many matmuls as necessary

It seems suboptimal to do extra K passes, especially for shorter sequences.

Proposed changes:

Instead of tiling independently from the assignment I propose to do similar as in https://github.com/Dao-AILab/quack/blob/main/quack/tile_scheduler.py and change the way thread groups load tiles. Inside gather_mm we run schedule_row_tile function that given precomputed offsets from rhs_indices fills the row to start from, expert to load, and the number of valid rows in the tile.

Then in the above example instead of 6 passes we do 4.

Note: it seems like a big diff, but it is just the tiling that differs

Results: MoE layer timing on M5Ultra

Model = shape is taken from the models config.

Model	seq len	after (ms)	before (ms)	Speedup
glm-5.3-flash	512	21.30320	33.16050	1.56×
glm-5.3-flash	713	24.83171	37.49358	1.51×
glm-5.3-flash	1024	22.27508	35.77488	1.61×
glm-5.3-flash	2048	25.16435	45.07096	1.79×
glm-5.3-flash	4123	40.18762	59.67581	1.48×
glm-5.3-flash	8192	65.52309	85.18234	1.30×
glm-5.3-flash	16384	107.03660	123.66846	1.16×
glm-5.3	512	28.13403	44.23767	1.57×
glm-5.3	713	29.82409	49.86723	1.67×
glm-5.3	1024	32.78672	48.42596	1.48×
glm-5.3	2048	36.94385	56.73476	1.54×
glm-5.3	4123	58.45475	81.96201	1.40×
glm-5.3	8192	96.37147	120.72586	1.25×
glm-5.3	16384	163.47526	186.35038	1.14×
qwen3.5-35b-a3b	512	2.70105	3.68130	1.36×
qwen3.5-35b-a3b	713	2.76690	3.78845	1.37×
qwen3.5-35b-a3b	1024	3.00206	3.93652	1.31×
qwen3.5-35b-a3b	2048	3.38963	4.21988	1.24×
qwen3.5-35b-a3b	4123	4.96703	6.15534	1.24×
qwen3.5-35b-a3b	8192	7.73033	8.91376	1.15×
qwen3.5-35b-a3b	16384	13.18756	14.19697	1.08×
qwen3.5-122b-a10b	512	6.73582	10.37210	1.54×
qwen3.5-122b-a10b	713	7.21023	11.04632	1.53×
qwen3.5-122b-a10b	1024	7.17348	11.06095	1.54×
qwen3.5-122b-a10b	2048	8.96438	13.41360	1.50×
qwen3.5-122b-a10b	4123	13.45173	18.19452	1.35×
qwen3.5-122b-a10b	8192	21.42444	26.36917	1.23×
qwen3.5-122b-a10b	16384	35.37630	39.33245	1.11×
qwen3.5-397b-a17b	512	17.36702	26.20703	1.51×
qwen3.5-397b-a17b	713	19.09090	28.52625	1.49×
qwen3.5-397b-a17b	1024	21.88976	31.99872	1.46×
qwen3.5-397b-a17b	2048	27.62238	36.44672	1.32×
qwen3.5-397b-a17b	4123	32.78854	45.61754	1.39×
qwen3.5-397b-a17b	8192	49.80969	64.99580	1.30×
qwen3.5-397b-a17b	16384	81.45321	97.39703	1.20×

Todo:

Do the same for gather_qmm.

[zcbenz]

It is amazing that changing scheduling can make such an improvement!