Skip to content

floor_divide on unsigned integers now builds seven ops (#4515); a MoE decode graph hit a Metal GPU timeout #4636

Closed

Description

@CodeAndCanvas728

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

Activity

  1. aaishwarymishra commented on Oct 6, 2026

    @aaishwarymishra
    Contributor

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

    auto step = logical_and(
        not_equal(b, zero, s),
        logical_and(
            not_equal(rem, zero, s),
            not_equal(less(a, zero, s), less(b, zero, s), s),
            s),
        s);

    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?

  2. zcbenz commented on Oct 7, 2026

    @zcbenz
    Member

    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.

  3. aaishwarymishra commented on Oct 7, 2026

    @aaishwarymishra
    Contributor

    @CodeAndCanvas728 Hi, do you want to give it a try?

  4. CodeAndCanvas728 commented on Oct 7, 2026

    @CodeAndCanvas728
    ContributorAuthor

    Yes, I can pick this up!

  5. added a commit that references this issue on Oct 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions