feat(compile): CustomKernel and GatherQMM implement output_shapes for shapeless compile

[dexwritescode]

Problem

mx::compile(shapeless=true) calls Primitive::output_shapes() on every
node when re-tracing a compiled function after input shapes change. Two
primitives were missing this override, causing compiled functions that
contain them to throw at runtime:

[Primitive::output_shapes] CustomKernel cannot infer output shapes
[Primitive::output_shapes] GatherQMM cannot infer output shapes

This makes it impossible to use mx::compile on models that combine custom
Metal kernels with gather-quantized-matmul — for example, hybrid SSM+attention
MoE models (like Qwen3 MoE) where the SSM step uses a custom Metal kernel and
the MoE routing uses gather_qmm.

Fix

CustomKernel (mlx/fast_primitives.h, mlx/backend/metal/custom_kernel.cpp, mlx/backend/cuda/custom_kernel.cpp)

The output shapes are already provided by the caller at creation time via
metal_kernel()(inputs, output_shapes, ...) and passed to array::make_arrays.
They just weren't stored on the primitive.

Add an optional output_shapes constructor parameter (default {} — backward
compatible), store in output_shapes_ member, override output_shapes() to
return it. Falls through to the base-class throw when empty (legacy path).

GatherQMM (mlx/primitives.h)

The output shape is fully inferrable from the stored fields and input shapes:

out_shape = lhs_indices.shape() + [x.shape(-2), w_outer_dims]

where w_outer_dims = transpose ? w.shape(-2) : w.shape(-1) * 32 / bits.

Input layout differs by quantization mode: Affine mode has biases at index 3,
pushing lhs_indices to index 4; other modes have lhs_indices at index 3.

Testing

Verified by enabling mx::compile(shapeless=true) on a 94-layer hybrid
SSM+attention MoE model (Qwen3.6-35B-A3B-4bit) where the GatedDeltaNet SSM
step uses a custom Metal kernel and the MoE routing uses gather_qmm.
Previously crashed on re-trace; with this fix the compiled graph is reused
correctly across decode steps.

[zcbenz]

The change overall looks good to me, can you add some simple tests?

[dexwritescode]

Added two tests to python/tests/test_compile.py:

• test_shapeless_compile_custom_kernel — compiles a metal_kernel passthrough with shapeless=True, then calls it with a larger shape. Fails before the fix with CustomKernel cannot infer output shapes.
• test_shapeless_compile_gather_qmm — compiles a gather_qmm call with shapeless=True, then calls it with a different M dimension. Fails before the fix with GatherQMM cannot infer output shapes.

Also addressed the two inline comments (removed std::move on const ref in both Metal and CUDA backends, fixed comment alignment in primitives.h).

[dexwritescode]

Hey @zcbenz, I've pushed a change to fix a test that failed. Would appreciate if you can trigger the CI run. Thanks!