[Metal] Expose kernel_mul_mm_id wrapper and parallelize rowids fan-out

[fiorelorenzo]

The single-pass MoE-fused matmul shader kernel_mul_mm_id is vendored at candle-metal-kernels/src/metal_src/quantized.metal:7267 but is not callable from Rust today. This PR:

1. Adds call_quantized_matmul_mm_id in candle-metal-kernels mirroring call_quantized_matmul_mm_t. Public dispatch function for callers that own their own MoE routing (CustomOp users, candle-transformers MoE models). Not auto-wired into QMatMul::forward; that is a follow-up once the wrapper has callers.

2. Parallelizes the rowids fan-out at quantized.metal:7301. The vendored shader carries a TODO: parallelize this loop here, inherited from the older single-pass kernel_mul_mm_id shape. The serial scan runs in every thread of every threadgroup, redundantly scanning the full M * T ids buffer; for Qwen3-MoE shapes (num_experts=128, T=8) at M=256 this dominates the kernel runtime. The replacement uses a stride-128 scan plus a threadgroup atomic counter. The downstream matmul reads rowids by index and writes into dst[jid[0], jid[1], :], so the resulting non-deterministic rowid order does not affect output. Upstream ggml-org/llama.cpp has since refactored to a two-pass kernel_mul_mm_id_map0 + kernel_mul_mm_id design (see ggml/src/ggml-metal/ggml-metal.metal:9721-9784) where the first pass is already parallel, so there is no longer anything to upstream this change to.

3. Adds a Criterion microbench at candle-core/benches/benchmarks/qmatmul_id.rs driving the wrapper at Q4_K and F16, N=2048, K=4096, num_experts=128, T=8. Q4_K sweeps M in {1, 32, 256}; F16 sweeps M in {1, 32} only (the F16 weight stack at these shapes is ~2.15 GB and crowds a 16 GB host at M=256). Plus a QMatMul::forward baseline at the equivalent per-expert load for comparison.

4. Adds an F32 correctness test (src/tests.rs::qmatmul_mm_id_f32_correctness) that validates the wrapper output against a naive CPU reference. Writing this test surfaced a wrapper layout bug (commit 0223447): ne1 was set to nei0 * nei1 instead of nei0, and src1 was wired with ne11 = num_tokens instead of ne11 = 1, so the kernel wrote dst at +token * (n * num_tokens * T) and read src1 at row slot % num_tokens. Microbenches never read their output and the prior iterations silently wrote out of bounds for M > 1. The fix and the correctness test landed together; the rewritten wrapper now produces a contiguous [num_tokens, experts_per_token, n] output the test asserts byte-for-byte against the CPU reference.

Benchmarks on MacBook Air M4, 16 GB unified memory, macOS 26.3.1 (post wrapper-layout fix):

Q4_K (the realistic MoE quantization):

Case	this PR (fused)	mm_t * num_experts baseline
M=1 (decode, 8 experts hit)	1.94 ms / 69 GFLOPs/s	8 * 40.2 us = 0.32 ms
M=32 (small prefill, all 128)	27.9 ms / 154 GFLOPs/s	128 * 273 us = 35.0 ms
M=256 (chunked, all 128)	52.9 ms / 649 GFLOPs/s	128 * 273 us = 35.0 ms

F16 (skips the per-row dequant the Q4_K path pays):

Case	this PR (fused)
M=1	1.71 ms / 78 GFLOPs/s
M=32	25.3 ms / 170 GFLOPs/s

Effect of the rowids fan-out parallelization in isolation (with vs without point 2 above): Q4_K M=256 goes from 800 ms (43 GFLOPs/s) to 52.9 ms (649 GFLOPs/s), a 15.1x speedup attributable to the scan rewrite.

The fused dispatch wins over the per-expert mm_t loop at M=32 (1.25x faster on Q4_K) and is within 1.5x of the loop at M=256. Both extremes are still limited by per-threadgroup overhead from launching all num_experts expert-threadgroups regardless of routing density; closing that is a follow-up (the M=1 microbench at num_experts=8 simulating an ideal active-expert pre-pack only gave ~9% over the current 128-launch path, so the structural fix would be a matvec kernel variant).

Bench command:

cd candle-core && cargo bench --bench bench_main --features metal -- qmatmul_id

References:

• kernel_mul_mm_id upstream: ggml-org/llama.cpp. Llama.cpp has since refactored to a two-pass _id_map0 + mm_id design (see ggml-metal-ops.cpp:2329); candle vendored only the older single-pass kernel.
• Related: Metal: fused MoE kernels, QTensor buffer access, encoder reuse, Q4K tuning #3444 (bounced for scope); this PR addresses only the MoE-fused matmul wrapper, no fused-quant cross-products, no encoder reuse.