Add gather_qmm for mixture-of-experts layers

[Codcore]

mx::gather_qmm is a quantized matmul where each row selects which weight
matrix to use. It is the primitive a mixture-of-experts layer needs: the expert
weights stay in one stacked {experts, out, in} tensor and the gather happens
inside the kernel, so no per-token copy of the weights is ever materialised.

Without it, an MoE block has to loop over experts and call quantized_matmul
once per expert. That scales with the expert count, which is 128 in Gemma 4 and
256 in the current Qwen MoE checkpoints — far past the point where a loop is
usable.

What this adds

• NIF(gather_qmm) in emlx_nif.cpp, written against the existing
  quantized_matmul binding. Same parameters plus lhs_indices, rhs_indices
  and sorted_indices; registered at arity 13 through ASYNC_NIF.
• EMLX.gather_qmm/11, following the shape of EMLX.quantized_matmul/8 —
  optional tensors unwrapped the same way, device merged across all six inputs,
  dispatched through the same command queue.
• test/emlx/gather_qmm_test.exs.

Nothing existing is touched: 96 added lines, no deletions.

Correctness

Checked against mx.gather_qmm through the Python bindings on the same machine
and the same libmlx build: maximum absolute difference 6e-8 over 4 experts,
6 rows, 2 experts per row, 4-bit affine weights with group size 64.

The tests here cover it independently of Python:

• each row matches a plain quantized_matmul against the expert it names
• repeated indices produce identical rows
• sorted_indices: true and false agree
• output shape follows the index shape

The full suite passes: 2672 tests, 825 doctests.

Notes

The shape protocol is the one mlx_lm's SwitchGLU uses — input expanded to
{..., 1, 1, in}, indices {..., experts_per_row}, output
{..., experts_per_row, 1, out} — so an MoE layer built on this can follow the
reference implementation directly.

EMLX.quantize/2 is rank-2 only, so stacked expert weights have to go through
EMLX.quantize/4. That is a separate matter and is left alone here.

[polvalente]

Thanks! Let's also expose the function in the quantization module as gather_quantized_matmul

emlx/emlx/lib/emlx/quantization.ex

Line 286 in 09dcaae

deftransform quantized_matmul(activation, qw) do

There's a pattern there for proper fusion in the compiler

[Codcore]

Pushed gather_quantized_matmul/4 in EMLX.Quantization, following the quantized_matmul/2 pattern you pointed at: a deftransform that computes the output shape and type, builds an Nx.template and goes through Nx.runtime_call, plus the EMLX.gather_quantized_matmul/4 counterpart that unpacks the config off the tensor.

One question about the intended path for building the input, rather than guessing at it.

gather_qmm needs the expert weights stacked on a leading axis, so qw is rank 3. EMLX.quantize/2 rejects anything but rank 2, so the tests build the stacked tensor through the device-level EMLX.quantize/4 and quantized_tensor/5 — which reads like what that function is documented for ("when you already have packed weights from a checkpoint"), and is how the weights arrive in practice anyway.

Is that the path you have in mind, or should quantize/2 accept rank > 2 and quantize along the last axis? I tried the latter first and backed it out: mx::quantize handles the leading axes fine, but widening quantize/2 alone lets a rank-3 quantized tensor reach quantized_matmul/2, where {out_features, _} = Nx.shape(qw) then fails with a MatchError instead of a clear message. That seemed like a separate change with its own consequences, so I left it out of this PR — happy to open one if it is worth having.

[polvalente]

You can increase the shape support to quantize and dequantize! The shape restriction is likely not intentional