You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
floor_divide on unsigned integers now builds seven ops (#4515); a MoE decode graph hit a Metal GPU timeout #4636
#4515 made integer floor_divide round toward negative infinity by adding remainder, less, not_equal, logical_and and subtract after the divide. For unsigned dtypes that correction is always zero, but the extra ops are still built on every call.
We hit this through the usual MoE gather-sort pattern: order = argsort(indices) (uint32) followed by order // top_k, once per MoE layer. On v0.32.3 a streamed-experts decode step of a 35B-A3B model (Apple M2, 16 GB) died on every run with Caused GPU Timeout Error (kIOGPUCommandBufferCallbackErrorTimeout); v0.32.2 never failed the same request. Bisecting v0.32.2..v0.32.3 gives #4515 as the first bad commit (the commit right before it is good), and reverting only that commit makes the crash go away. Even when the server survived, decode was about 20% slower.
What I could not find: why the extra ops lead to a timeout rather than just being slower. Standalone, argsort(...) // 8 on uint32 is correct and equally fast on both versions, so I don't have a small repro; it needs the full decode graph.
Suggested change: return the plain quotient when the promoted dtype is unsigned, since the sign correction cannot apply:
if (issubdtype(dtype, unsignedinteger)) {
return quotient;
}
Signed integers keep the #4515 behaviour. We carry exactly this in a downstream fork (SharpAI/mlx-swift#22) with a test for both cases. Happy to open a PR here if that is useful.
@zcbenz I think for now we can push the unsigned early return suggested here, along with the zero-divisor check discussed in Layr-Labs#32 to preserve the existing CPU behavior.
I think we should just add a FloorDivide primitive with new kernels, binary ops like this are simple to add, there is really no need to have a workaround first.
Activity
aaishwarymishra commentedon Oct 6, 2026
@zcbenz I think for now we can push the unsigned early return suggested here, along with the zero-divisor check discussed in Layr-Labs#32 to preserve the existing CPU behavior.
Meanwhile, I can look into updating the kernels directly so we don’t need these extra graph operations for signed integers either. What do you think?
zcbenz commentedon Oct 7, 2026
I think we should just add a
FloorDivideprimitive with new kernels, binary ops like this are simple to add, there is really no need to have a workaround first.aaishwarymishra commentedon Oct 7, 2026
@CodeAndCanvas728 Hi, do you want to give it a try?
CodeAndCanvas728 commentedon Oct 7, 2026
Yes, I can pick this up!
Fix for ml-explore#4636
Fix for ml-explore#4636
Fix for ml-explore#4636