[dhkts1]

Metal: 64-token tile for the quantized matrix-matrix kernel

What changes

The Metal quantized matmul for more than one token (kernel_mul_mm) now multiplies each tile of
dequantized weights by 64 tokens at a time instead of 32, wherever that pads no more tokens than the
32-token tile (see "Where the 64-token tile is used", below). A single token is unchanged: it still
goes through fwd_mv and the matrix-vector kernel.

• kernel_mul_mm takes the token tile as a template parameter. It defaults to 32, so every existing
  instantiation compiles to what it was.
• New _n64 instantiations exist for every type the Rust dispatch serves: F32, F16, BF16, Q4_0, Q4_1,
  Q5_0, Q5_1, Q8_0, Q2_K, Q3_K, Q4_K, Q5_K and Q6_K.
• call_quantized_matmul_mm_t keeps its signature and runs the tile QuantizedMmTile::for_tokens
  picks. candle-core's QMetalStorage::fwd uses the same function.
• call_quantized_matmul_mm_t_with_tile and a QuantizedMmTile enum let a caller pick either tile.
• The threadgroup memory is set per tile: 8192 bytes at 32 tokens, as before, and 16384 bytes at 64.

Why the tile sets the weight reuse

A threadgroup owns 64 weight rows. On every step along K it reads those rows' quantized blocks,
dequantizes them into threadgroup memory, and multiplies them by the tokens of its tile. Nothing
outside the tile shares that work. So each weight is fetched and dequantized once per tile of tokens:
16 times per 512-token sequence at a 32-token tile, 8 times at 64. A batch of sequences does not
change this, because the batch is only more tiles. Doubling the tile halves the weight traffic and
the dequantize work per token. The cost is the threadgroup memory for the wider activation tile,
and twice the accumulators per simdgroup.

Where the 64-token tile is used

The 64-token tile is used only when it pads no more tokens than the 32-token tile: when the token
count is a multiple of 64, or leaves more than 32 over a multiple of 64. Otherwise its last tile is
mostly padding, and that wasted work costs more than the weight reuse saves. The measurements below
show that pattern: it was slower at every count that pads more (up to 32, and at 96, 136, 160, 200,
224 and 288 tokens). At every count that pads no more it was faster or even, with one exception: a
20-call run at 64 tokens read 1.09, while the other run there read 0.94. So the rule says why the
tile helps, and it is not a cut-off tuned to one machine.

Same results, bit for bit

Each output element is accumulated in exactly the same order at either tile: the same K steps and
the same 8x8 multiply-accumulates, only grouped into larger threadgroups. The new test
quantized::metal::tests::mm_tile64_is_byte_identical_to_tile32 runs both tiles and requires every
output to be equal bit for bit:

• all 13 types;
• five shapes: 7 tokens (a partial tile at either size), 32 (one full 32-token tile), 33 (between the
  two tiles), 130 (several tiles and a partial one, with 96 weight rows), and a batch of 2 at 65
  tokens with 128 rows and K = 768;
• a guard that the result is not all zeros.

To check the test catches a real defect, I made every token row of a 64-token tile load the first
row's activations. The test failed ("Q4_0 1x33x200x512: 200 of 6600 outputs differ between the
tiles"). With the change restored it passes again.

A second test, mm_tile_for_tokens_follows_padding, pins the choice at 2, 32, 33, 64, 96, 100, 128,
136 and 2048 tokens. With the rule's <= changed to < it fails ("33 tokens: Tokens32 vs Tokens64").

Checks run

On an Apple M4 Pro with macOS's Metal, in this branch's own target directory:

• cargo fmt --all -- --check: clean.
• cargo clippy -p candle-metal-kernels -p candle-core --features candle-core/metal -- -D warnings:
  clean.
• cargo clippy -p candle-core --features metal --lib --tests -- -D warnings -A clippy::items_after_test_module:
  clean. The one allowed lint is reported on candle-core/src/metal_backend/device.rs, which this
  change does not touch.
• candle-metal-kernels' own tests do not build on main at present: src/tests.rs uses rand and
  rand_distr, which are not among its dependencies. This change does not touch that file.
• candle-core with --features metal:
  • both new tests pass;
  • quantized_tests: 83 passed;
  • the 27 lib tests and the matmul, bilinear, layout, pool, indexing, gguf, custom_op, display,
    metal_concurrent, pth and serialization tests pass. Run the last two from candle-core/, because
    they read fixture files by a relative path.
  • Three Metal tests fail, none in the quantized path, and they fail the same way on main
    (12bb5d1f, built in a separate worktree):
    • conv2d_grad_noncontiguous_kernel_metal;
    • index_select_strided_ids_grad_metal ("Metal index_select requires contiguous ids");
    • meshgrid_with_empty_axis_metal.

Numbers

All on one Apple M4 Pro (20 GPU cores), shared with other jobs (load average 160 to 220 during the
candle runs).

This branch, kernel level

fwd_with_tile was timed at each tile on the seven Qwen3-Embedding-4B Q4_K_M projection shapes:

projection	shape	type
attn_q	4096 x 2560	Q4_K
attn_k	1024 x 2560	Q4_K
attn_v	1024 x 2560	Q6_K
attn_output	2560 x 4096	Q4_K
ffn_gate	9728 x 2560	Q4_K
ffn_up	9728 x 2560	Q4_K
ffn_down	2560 x 9728	Q6_K

The weights are random, quantized to those types. Each sample is 20 or 50 calls between two
synchronizes (column "calls"). Each round runs both tiles, and the order alternates every round.
There were six rounds per run.

"Ratio" is the 64-token tile's time divided by the 32-token tile's, each the median ms per
projection summed over the seven shapes (one layer's multiplies). A ratio below 1 means the 64-token
tile is faster. "Rounds faster" counts the rounds in which it was. Two values mean two separate
runs. "Chosen" is the tile the dispatch now picks.

tokens	tokens mod 64	chosen	ratio	rounds faster	calls
2	2	32	1.41	0/6	20
4	4	32	1.37	0/6	20
8	8	32	1.49	0/6	20
16	16	32	1.48	0/6	20
32	32	32	1.41	0/6	20
48	48	64	1.01	3/6	20
64	0	64	0.94, 1.09	4/6, 2/6	20
96	32	32	1.14	1/6	50
100	36	64	0.95	5/6	50
128	0	64	0.97, 0.95	5/6, 5/6	50
136	8	32	1.16	0/6	50
160	32	32	1.15, 1.15	0/6, 0/6	50
176	48	64	1.00	2/6	50
192	0	64	0.99, 1.02	4/6, 1/6	50
200	8	32	1.03	0/6	50
224	32	32	1.04, 1.08	0/6, 0/6	50
240	48	64	0.95	6/6	50
256	0	64	0.93, 0.95	6/6, 6/6	50
288	32	32	1.04	0/6	50
300	44	64	0.97	6/6	50
320	0	64	0.96	6/6	50
352	32	32	1.00	2/6	50
384	0	64	0.93	6/6	50
400	16	32	1.00	2/6	50
512	0	64	0.93	4/6	20
1024	0	64	0.96	4/6	20
2048	0	64	0.89	5/6	20

The full model, from a standalone crate

qkernels-bench carries the same 64-token tiling of this kernel, and its 64-token result is
byte-identical to candle's. There, the full 36-layer forward of Qwen3-Embedding-4B Q4_K_M ran with
every projection at the 64-token tile, over three rounds, each against a same-round rerun of
candle's path:

• 512 tokens: 0.92, 0.94 and 0.92 of the time.
• 2048 tokens: 0.94, 0.92 and 0.93.

The pooled embeddings were byte-identical to candle's.

Two commits

1. The 64-token tile and the byte-identical test.
2. Use the 64-token tile only when it pads no more than the 32-token tile.
   QuantizedMmTile::for_tokens holds the rule, and both call_quantized_matmul_mm_t and
   candle-core's QMetalStorage::fwd use it. Up to 32 tokens it is always the 32-token tile, so a
   short prefill runs exactly the kernel it ran before.
   • The byte-identical test gains a 32-token shape.
   • The new mm_tile_for_tokens_follows_padding test pins the rule.

https://claude.ai/code/session_01Lodcb7ZaRA8icgxCC8cudg
🍄