[BUG] fast.metal_kernel silently truncates the Metal dispatch: group_dims clamped against grid dimensions in CustomKernel::eval_gpu

[pseudobacon]

Summary

mx.fast.metal_kernel custom kernels are silently dispatched over a truncated threadgrid whenever a grid dimension is smaller than the matching requested threadgroup dimension — which is the common case (e.g. one threadgroup per query row with a 256-wide threadgroup).

Root cause

In CustomKernel::eval_gpu (mlx/backend/metal/custom_kernel.cpp):

MTL::Size group_dims =
    MTL::Size(std::min(tx, gx), std::min(ty, gy), std::min(tz, gz));
MTL::Size grid_dims = MTL::Size(gx, gy, gz);
compute_encoder.dispatch_threads(grid_dims, group_dims);

dispatch_threads interprets grid_dims as thread counts, so with group_dims clamped to min(tx, gx) etc. the effective threadgroup count per dimension is ceil(min(tx,gx)/tx_per_group) — a small prefix of the requested grid. For a per-row kernel with requested grid (24, Q, 1) and threadgroup (256, 1, 1), the effective dispatch covers 24 threads of x per row (1 threadgroup of 24 threads) instead of 256 — i.e. ~1/256 of the requested work executes.

Behavior

• The truncation is stable across repeated same-geometry dispatches (verified: 20+ dispatches all truncated at a fresh geometry; re-dispatch does not clear it).
• It is invisible to every check except full coverage of the requested grid (e.g. a probe kernel writing each thread's unique index and comparing the output to {1..N}).
• Benchmark impact: a custom kernel's timing measured through this path can appear orders of magnitude faster than the true full-work cost. In our case a flash-SDPA kernel measured 3.5–4.8× faster than the incumbent with truncated dispatch, and 57.7× slower at true full work. Any custom-kernel performance claim made before this fix is invalid without per-dispatch fullness evidence.

Repro (minimal, no MLXLM, no attention code)

import MLX

// probe kernel: out[idx] = idx for each executing thread
let src = """
kernel void full_probe(device const int2& in [[buffer(0)]],
                       device int& out [[buffer(1)]]) {
  long idx = (long)(in.y * in.x + threadgroup_position_in_grid.x) * 256L
      + position_in_threadgroup.x;
  out[idx] = (int)idx;
}
"""
let fn = MLXFast.metalKernel(
    name: "full_probe",
    inputNames: ["in"],
    outputNames: ["out"],
    source: src)

let inArr = MLXFast.metalKernel(/* ... */ ) // unused values, shape {gy*gx, 2} int32
let out = MLX.zeros(shape: [24 * 64 * 256], dtype: MLX.DType.int32)!
// grid (24, 64, 1), threadgroup (256, 1, 1)
let r = try fn(inputs: [inArr], outputShapes: [out.shape], outputDTypes: [.int32],
               threadsPerThreadgroup: (256, 1, 1), threadsPerGrid: (24, 64, 1))
MLX.eval(r.0)
// Full coverage: out == arange(1, N+1). Before the fix: only a small
// non-zero prefix (truncated dispatch); stable across repeated calls.

C++ equivalent: fast::metal_kernel(...) in tests/gpu_tests.cpp (added in the linked PR).

Evidence from our investigation (Qwen3.8 MTP server, Qwen3.8-27B, M5 Pro)

1. Trivial probe (no attention code): dispatch alternates between RECOMBINED (24 groups × 64 threads) and FULL (1536 groups × 256 threads) on every other dispatch; pure same-geometry sequences stay truncated for all dispatches (20/20).
2. Q sweep: deterministic across 3 runs; truncation fires for all Q (not just small Q), first dispatch of a fresh geometry included.
3. Production geometry: Q=2048, prefix=8192 (our prefill chunk): truncated dispatch writes 49,152 elements (= 24 × 2048) instead of 12,582,912 (= 24 × 2048 × 256).
4. Pins at time of report: mlx-swift 2bebe4e9ad127758ebcd76c6ad45a1740d0d2852, mlx C++ 1f8e74e3f12f31365464a6867c6579f0e9b29d85 (2026-08-25).
5. Full-work re-bench after local fix: flash kernel is bit-exact deterministic and correct (max|diff| 6.1e-5 vs bf16 reference) at true full work, confirming the truncation was the sole cause of the earlier measurements.

Fix

Dispatch the requested threadgroup dimensions unclamped through dispatch_threadgroups:

MTL::Size group_dims = MTL::Size(tx, ty, tz);
MTL::Size grid_dims = MTL::Size(gx, gy, gz);
compute_encoder.dispatch_threadgroups(grid_dims, group_dims);

(See linked PR for the patch + a full-coverage regression test in tests/gpu_tests.cpp.)

Versions

• MLX: 1f8e74e3f12f31365464a6867c6579f0e9b29d85 (2026-08-25)
• mlx-swift: 2bebe4e9ad127758ebcd76c6ad45a1740d0d2852
• macOS 26.5.2 (25F84), Apple M5 Pro, 48 GB

[haoxu0]

grid in MLX is a thread count, not a threadgroup count, so min(tx, gx) is the correct group size for dispatch_threads and I don't think anything is truncated.

docs/src/dev/custom_metal_kernels.rst: "we will launch mx.prod(grid) threads, subdivided into threadgroup size threadgroups. For optimal performance, each thread group dimension should be less than or equal to the corresponding grid dimension." The CUDA docstrings say it too (fast.cpp:580, :742).

Coverage on current main, mlx 0.32.2 — every executing thread writes 1:

grid            threadgroup    covered      requested
(24, 1, 1)      (256, 1, 1)    24           24
(257, 1, 1)     (256, 1, 1)    257          257
(13141739,1,1)  (256, 1, 1)    13,141,739   13,141,739
(24, 64, 1)     (256, 1, 1)    1,536        1,536
(24, 512, 1)    (256, 1, 1)    12,288       12,288

The last two are your geometry — grid.x (24) below the threadgroup width (256) — fully covered.

Your kernel is running 1/256 of the threads you intended, but because grid is in threads: grid=(24, Q, 1) requests 24 threads per row, not 24 groups of 256. For 24 groups, pass grid=(24*256, Q, 1). That also explains the benchmark — the 3.5-4.8x was measured at 1/256 of the work, and 57.7x slower is the real number.

On #4535: dispatch_threadgroups changes what grid means for every existing caller. Kernels launched as grid=(n,1,1) would get n groups of 256 — 256x the threads, and out-of-bounds writes for any kernel that sizes grid exactly and omits a guard, which is reasonable today since the grid is the thread count.

repro

import mlx.core as mx, numpy as np
k = mx.fast.metal_kernel(
    name="probe", input_names=["nx","ny"], output_names=["out"],
    source="""
        uint x = thread_position_in_grid.x, y = thread_position_in_grid.y;
        if (x >= (uint)nx[0] || y >= (uint)ny[0]) return;
        out[y * (uint)nx[0] + x] = 1;
    """)
for nx, ny in [(24,64), (24,512), (13_141_739,1)]:
    out, = k(inputs=[mx.array([nx], mx.uint32), mx.array([ny], mx.uint32)],
             output_shapes=[(nx*ny,)], output_dtypes=[mx.uint32],
             grid=(nx,ny,1), threadgroup=(256,1,1), init_value=0)
    mx.eval(out)
    print(nx, ny, int(np.array(out).sum()), nx*ny)