Implement output_shapes for GatherMM and GatherQMM by dexwritescode · Pull Request #3485 · ml-explore/mlx

dexwritescode · 2026-05-05T19:26:19Z

Problem

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

[Primitive::output_shapes] GatherQMM cannot infer output shapes

Fix

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

Added test_shapeless_compile_gather_qmm to python/tests/test_compile.py:
compiles a gather_qmm call with shapeless=True and verifies it reuses the
graph correctly when the M dimension changes between calls.

`mx::compile(shapeless=true)` calls `Primitive::output_shapes()` on every node when re-tracing a compiled function with changed input shapes. `CustomKernel` never implemented this override, so any compiled function containing a `metal_kernel` / `custom_kernel` call would throw: [Primitive::output_shapes] CustomKernel cannot infer output shapes 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. Fix: add an optional `output_shapes` parameter to the `CustomKernel` constructor (default `{}` for backward compatibility), store it in a new `output_shapes_` member, and override `output_shapes()` to return it. If the field is empty (legacy construction path), fall through to the base-class throw as before. Update both Metal and CUDA call sites to copy the shapes before `std::move`-ing them into `array::make_arrays` and pass the copy to the constructor.

output_shapes() is called on every primitive during shapeless=true retracing. GatherQMM was missing this override, causing compile to throw when any graph containing gather_qmm was retraced. The output shape is fully inferrable from inputs and stored fields: 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 mode: Affine has biases at index 3 (pushing indices to 4/5); other modes have indices at 3/4.

zcbenz

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

Address review feedback from zcbenz: - output_shapes is a const& in the lambda parameter, so std::move(output_shapes) compiles but silently copies rather than moves. Remove the misleading std::move in both metal and cuda backends — make_arrays receives a plain copy. - Fix one extra space in the GatherQMM input layout comment to correctly align lhs_idx under the Affine layout line.

…rQMM Verify that mx.compile(shapeless=True) correctly re-traces functions containing mx.fast.metal_kernel (CustomKernel) and mx.gather_qmm (GatherQMM) when input shapes change between calls. Both tests fail before the fix with the respective 'cannot infer output shapes' error and pass after output_shapes() is implemented.

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).

Remove the intermediate output_shapes_copy and pass output_shapes directly to the CustomKernel constructor, which takes it by value.

…:async_eval encoder Adds INFER_CPP_PHASE_TIMING=1 stderr probes around the two C++ FFI hot paths so we can split "Rust async_kick = 23ms" into its components: - crates/mlx-sys/src/mlx_qwen35_model.cpp:2541 — `forward_build_us` around `m->forward(inputs)` (lazy graph build). - crates/mlx-sys/src/mlx_bridge.cpp:2072 — `async_eval_call_us` around the actual `mx::async_eval(arrs)` call (encoder + commit work). Cached env probe (one atomic read after first call); zero prod cost when env unset. file-static helper in each TU to avoid header churn. Bench (Qwen3.6 35B-A3B-4bit, c=4 + c=8): forward_build_us c=4 p50 = 1509μs ← lazy graph build is FAST forward_build_us c=8 p50 = 1793μs async_eval_call_us count=82 p50 = 24992μs ← here's the 25ms (count=82 = logits + new_sampled + 80 packed_kv_flat slabs) → Hypothesis confirmed (per MLX async_eval research subagent this date): mx::async_eval does graph traversal + Metal command-buffer encoding SYNCHRONOUSLY on the calling thread. Only GPU completion is async. For a 40-layer MoE forward (~600-1000 primitives at c=4-8), the ~25ms is real CPU encoder work — NOT GPU compute. Confirmed by mlx/transforms.cpp eval_impl(... async=true) which only skips the final wait, never offloads encoding. Erratum: AGENTS.md narrows the previous "MLX_MAX_OPS_PER_BUFFER=200 recommended for any Metal bench at c≥8" recommendation. That was Qwen3.5-dense-specific and benched as wash-or-loss on Qwen3.6 MoE per docs/experience/wins/2026-05-07-bench-qwen36-baseline.md. Removed from default guidance; downgraded to "per-workload matched-A/B tunable". Auto-wired-limit (default since 180e48b) is the canonical Metal serving knob. Wins entry: docs/experience/wins/2026-05-07-bench-qwen36-encode-bottleneck.md captures the localization, the implications (mx::compile blocked on ml-explore/mlx#3485, multi-thread encode blocked by #3078), and four viable next levers ranked S/M. 644 infer tests pass; clippy --features metal -- -D warnings clean. Bench: doc + instrumentation only; no hot-path behavior change when env unset (default).

lhs_indices shape (8,) cannot broadcast with the auto-generated rhs_indices arange(num_experts) shape (4,), causing the second shapeless compile call to fail during graph update. Fix by keeping both indices fixed at shape (num_experts,) and varying only the M dimension via x.shape = (num_experts, M, K).

grid=(x.size, 1, 1) is captured as a fixed tuple at trace time. On the second shapeless-compile call (x.size=8) the primitive still holds grid=(4,1,1), so only 4 of 8 output elements are written and array_equal fails. The test goal is to verify output_shapes prevents a throw and returns the correct shape — not value correctness, which would require the grid to be updated dynamically (out of scope).

dexwritescode · 2026-05-08T01:43:33Z

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

…ompile test output_shapes_ stores the trace-time shape and is returned as-is — it does not recompute from inputs. A kernel where output_shapes=[x.shape] causes shapeless compile to reuse a stale (4,) shape for (8,) inputs. Replace with a kernel whose output is always (1,) regardless of input size, so output_shapes_ is correct across all shapeless compile reuses.

zcbenz

Looks good to me, thanks!

angeloskath

Sorry but this is not correct. You can't simply pass the output shapes and then compile in shapeless mode the output shape will most likely change based on the input shape.

Take for instance the following code to see what the problem is

import mlx.core as mx
from functools import partial
# A "double" kernel: output has the same shape as input
double = mx.fast.metal_kernel(
    name="double",
    input_names=["inp"],
    output_names=["out"],
    source="uint i = thread_position_in_grid.x; out[i] = inp[i] * 2;",
)
def fn(x):
    return double(
        inputs=[x],
        grid=(x.size, 1, 1),
        threadgroup=(256, 1, 1),
        output_shapes=[x.shape],  # output shape == input shape
        output_dtypes=[x.dtype],
        stream=mx.gpu,
    )[0]
cfn = mx.compile(fn, shapeless=True)
x = mx.array([1.0, 2.0, 3.0])
print(cfn(x))  # [2, 4, 6] — shape (3,), correct
x = mx.array([1.0, 2.0, 3.0, 4.0, 5.0])
print(cfn(x))  # BUG: output_shapes() returns (3,) from the first trace,
               # but the correct output shape should be (5,)

feel free to remove the custom kernel and I will look into the gather qmm one.

The fix was incorrect for the general case: output_shapes_ stored the shapes from the first trace and returned them unconditionally, which produces stale shapes whenever output dimensions depend on input shapes in shapeless compile mode. The test was written around the bug rather than exposing it, using a fixed-output-shape kernel so the stale shapes were always correct. GatherQMM output_shapes() is retained — its shape is a deterministic function of input shapes and is correct in all cases.

dexwritescode · 2026-05-21T22:22:33Z

Sorry but this is not correct. You can't simply pass the output shapes and then compile in shapeless mode the output shape will most likely change based on the input shape.

Take for instance the following code to see what the problem is

import mlx.core as mx
from functools import partial
# A "double" kernel: output has the same shape as input
double = mx.fast.metal_kernel(
    name="double",
    input_names=["inp"],
    output_names=["out"],
    source="uint i = thread_position_in_grid.x; out[i] = inp[i] * 2;",
)
def fn(x):
    return double(
        inputs=[x],
        grid=(x.size, 1, 1),
        threadgroup=(256, 1, 1),
        output_shapes=[x.shape],  # output shape == input shape
        output_dtypes=[x.dtype],
        stream=mx.gpu,
    )[0]
cfn = mx.compile(fn, shapeless=True)
x = mx.array([1.0, 2.0, 3.0])
print(cfn(x))  # [2, 4, 6] — shape (3,), correct
x = mx.array([1.0, 2.0, 3.0, 4.0, 5.0])
print(cfn(x))  # BUG: output_shapes() returns (3,) from the first trace,
               # but the correct output shape should be (5,)

feel free to remove the custom kernel and I will look into the gather qmm one.

Hey @angeloskath, I've removed the custom kernel and kept the qmm.

angeloskath

Added the GatherMM output_shapes as well. Looks good.

Co-authored-by: Angelos Katharopoulos <[email protected]> (cherry picked from commit 0ab2176)

…:async_eval encoder Adds INFER_CPP_PHASE_TIMING=1 stderr probes around the two C++ FFI hot paths so we can split "Rust async_kick = 23ms" into its components: - crates/mlx-sys/src/mlx_qwen35_model.cpp:2541 — `forward_build_us` around `m->forward(inputs)` (lazy graph build). - crates/mlx-sys/src/mlx_bridge.cpp:2072 — `async_eval_call_us` around the actual `mx::async_eval(arrs)` call (encoder + commit work). Cached env probe (one atomic read after first call); zero prod cost when env unset. file-static helper in each TU to avoid header churn. Bench (Qwen3.6 35B-A3B-4bit, c=4 + c=8): forward_build_us c=4 p50 = 1509μs ← lazy graph build is FAST forward_build_us c=8 p50 = 1793μs async_eval_call_us count=82 p50 = 24992μs ← here's the 25ms (count=82 = logits + new_sampled + 80 packed_kv_flat slabs) → Hypothesis confirmed (per MLX async_eval research subagent this date): mx::async_eval does graph traversal + Metal command-buffer encoding SYNCHRONOUSLY on the calling thread. Only GPU completion is async. For a 40-layer MoE forward (~600-1000 primitives at c=4-8), the ~25ms is real CPU encoder work — NOT GPU compute. Confirmed by mlx/transforms.cpp eval_impl(... async=true) which only skips the final wait, never offloads encoding. Erratum: AGENTS.md narrows the previous "MLX_MAX_OPS_PER_BUFFER=200 recommended for any Metal bench at c≥8" recommendation. That was Qwen3.5-dense-specific and benched as wash-or-loss on Qwen3.6 MoE per docs/experience/wins/2026-05-07-bench-qwen36-baseline.md. Removed from default guidance; downgraded to "per-workload matched-A/B tunable". Auto-wired-limit (default since 180e48b) is the canonical Metal serving knob. Wins entry: docs/experience/wins/2026-05-07-bench-qwen36-encode-bottleneck.md captures the localization, the implications (mx::compile blocked on ml-explore/mlx#3485, multi-thread encode blocked by #3078), and four viable next levers ranked S/M. 644 infer tests pass; clippy --features metal -- -D warnings clean. Bench: doc + instrumentation only; no hot-path behavior change when env unset (default).

dexwritescode added 2 commits May 5, 2026 15:24

dexwritescode changed the title ~~feat(compile): CustomKernel stores and returns output shapes~~ feat(compile): CustomKernel and GatherQMM implement output_shapes for shapeless compile May 5, 2026

dexwritescode mentioned this pull request May 5, 2026

Phase O: Backend-specific model architecture dexwritescode/neurons#20

Merged

zcbenz reviewed May 6, 2026

View reviewed changes

Comment thread mlx/backend/cuda/custom_kernel.cpp Resolved

Comment thread mlx/primitives.h Resolved

dexwritescode added 2 commits May 6, 2026 19:32

dexwritescode commented May 6, 2026

View reviewed changes

dexwritescode requested a review from zcbenz May 6, 2026 23:45

refactor(compile): simplify output_shapes copy in CustomKernel

1109c8c

Remove the intermediate output_shapes_copy and pass output_shapes directly to the CustomKernel constructor, which takes it by value.

dexwritescode added 2 commits May 7, 2026 19:38

zcbenz approved these changes May 9, 2026

View reviewed changes

XXXXRT666 mentioned this pull request May 11, 2026

Add Metal DLPack zero-copy sharing #3531

Merged

4 tasks done

angeloskath requested changes May 19, 2026

View reviewed changes

dexwritescode added 2 commits May 21, 2026 18:18

style(compile): restore blank line removed by CustomKernel revert

7396860

dexwritescode changed the title ~~feat(compile): CustomKernel and GatherQMM implement output_shapes for shapeless compile~~ fix(compile): GatherQMM implements output_shapes for shapeless compile May 21, 2026

dexwritescode requested a review from angeloskath May 23, 2026 17:57

angeloskath added 2 commits May 28, 2026 14:23

Refactor the output_shapes implementation

c30b9f4

Add a test for GatherMM

64c8802

angeloskath changed the title ~~fix(compile): GatherQMM implements output_shapes for shapeless compile~~ Implement output_shapes for GatherMM and GatherQMM May 28, 2026

Remove blank lines

7af002f

angeloskath approved these changes May 28, 2026

View reviewed changes

zcbenz merged commit 0ab2176 into ml-explore:main May 29, 2026
16 checks passed

megacpp pushed a commit to DatasunriseOU/mlx that referenced this pull request Jun 3, 2026

Implement output_shapes for GatherMM and GatherQMM (ml-explore#3485)

a400bae

Co-authored-by: Angelos Katharopoulos <[email protected]> (cherry picked from commit 0ab2176)

BrewTestBot mentioned this pull request Jul 7, 2026

mlx 0.32.0 Homebrew/homebrew-core#291849

Merged

1 task done

Navigation Menu

Repository navigation

Implement output_shapes for GatherMM and GatherQMM - #3485

Implement output_shapes for GatherMM and GatherQMM#3485
zcbenz merged 13 commits into
ml-explore:mainfrom
dexwritescode:fix/custom-kernel-output-shapes

dexwritescode commented May 5, 2026 •

edited

Loading

Uh oh!

zcbenz left a comment

Uh oh!

Uh oh!

Uh oh!

dexwritescode left a comment

Uh oh!

dexwritescode commented May 8, 2026

Uh oh!

zcbenz left a comment

Uh oh!

angeloskath left a comment

Uh oh!

dexwritescode commented May 21, 2026

Uh oh!

angeloskath left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

dexwritescode commented May 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Testing

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dexwritescode left a comment

Choose a reason for hiding this comment

Uh oh!

dexwritescode commented May 8, 2026

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

angeloskath left a comment

Choose a reason for hiding this comment

Uh oh!

dexwritescode commented May 21, 2026

Uh oh!

angeloskath left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dexwritescode commented May 5, 2026 •

edited

Loading