[Satchitananda]

Fixes missing stream argument in some ops. Ops call helper op skipping the stream argument, in that case op is running on the default stream, causing a crash.

This PR is fixing found sites, adds regression tests which were failing before and passing after, adds a guard and commit-hook to check stream args passing.

Found when working on my own app and got crashes with:

MLX/ErrorHandler.swift:345: Fatal error: There is no Stream(gpu, 6) in current thread.

• I understand it is strictly prohibited to use AI to write PR description

AI usage disclosure: Fix written with Claude Code Opus 5.5, verified and adjusted by Fable 5.1 and Astra 6

[zcbenz]

Thanks for a thorough fix!

[ahuzmeza]

Thanks for this PR, glad I found it! While running Qwen3-TTS voice cloning in Swift (mlx-audio-swift), I hit a 2x slowdown compared to the Python version of the same model, and it turns out the quantized_matmul hunk here fixes it too.

When x and the scales have different dtypes, the astype calls run on the default stream, so a call on an explicit stream waits on the default stream every time. mlx-swift always passes an explicit stream, so every Swift caller with a dtype mismatch pays this. Python code under with mx.stream(...) doesn't, because the context changes the default stream.

So on top of the crash, this PR also fixes that slowdown:

	before	with the quantized_matmul hunk
Qwen3-TTS voice cloning, mlx-swift (frames/s)	11.8-12.3	23.9-25.8
quantized_matmul, float32 x, bf16 scales, on mx.new_stream(mx.gpu) (ms per 75 calls)	7.56-7.58	0.92-1.02

Mac mini M2 Pro. Swift: mlx-community/Qwen3-TTS-12Hz-1.7B-Base-8bit, mlx-swift's bundled core v0.32.2 with the hunk applied. Python: mlx main 56c2bf9 vs. main + the hunk. With matching dtypes, or on the default stream, nothing changes.

Python repro

import time, mlx.core as mx
mx.random.seed(0)
w, s, b = mx.quantize((mx.random.normal((2048, 1024)) * 0.02).astype(mx.bfloat16), group_size=64, bits=8)
mx.eval(w, s, b)
def bench(x, stream):
    eps = [mx.array(i * 1e-3, x.dtype) for i in range(75)]; mx.eval(eps); x = mx.array(x); mx.eval(x)
    f = lambda: mx.eval([mx.quantized_matmul(mx.add(x, e, stream=stream), w, s, b, transpose=True, group_size=64, bits=8, stream=stream) for e in eps])
    for _ in range(10): f()
    t0 = time.perf_counter()
    for _ in range(50): f()
    return (time.perf_counter() - t0) / 50 * 1e3
print("mlx", mx.__version__)
other = mx.new_stream(mx.gpu)
for name, st in [("default", mx.default_stream(mx.gpu)), ("new_stream", other)]:
    for dt in [mx.bfloat16, mx.float32]:
        x = mx.random.normal((1, 1, 1024)).astype(dt); mx.eval(x)
        print(f"stream={name} x={'f32' if dt == mx.float32 else 'bf16'} scales=bf16 ms_per_75={bench(x, st):.2f}")

Investigated with AI: it traced the slowdown to this hunk, wrote the same fix with a test, and ran the benchmarks above. I then used it to search the open and closed PRs for an existing fix, which is how I found this one.

Leaving this here so the speed improvement is on record for anyone who runs into the same slowdown.