[davidtai]

Proposed changes

This PR computes float16 and bfloat16 floor_divide, and the divmod quotient, in float32. The result is then rounded once to the output type. David asked for this as a separate PR after the review of PR #38.

The defect

Upstream computes these operations as floor(a / b) in the half type. The quotient is rounded to float16 or bfloat16 before the floor. When the true quotient is less than half an ulp below an integer n, it rounds to n, and the result is n instead of n - 1.

numpy's float16 divmod converts to float32 and uses npy_divmodf, which works from fmod, so it gives n - 1. Upstream's test_divmod draws unseeded float16 inputs from uniform(1, 100) and compares with numpy. My estimate is that it fails in about 4% of runs per device. It failed once in the GPU shard of PR #38 (run 39966062214), and a re-run passed.

The fix

A helper exact_floor_divide(float a, float b) uses the method of npy_divmodf:

1. mod = fmod(a, b)
2. div = (a - mod) / b
3. If mod and b have different signs, subtract 1 from div.
4. Floor div. Add 1 if div - floor(div) > 0.5. Keep the sign of zero. Return a / b when b == 0.

Where the fix applies:

File	Change
mlx/backend/metal/kernels/binary_ops.h	FloorDivide uses the helper for half and bfloat16_t. DivMod uses FloorDivide, so its quotient is fixed too.
mlx/backend/cpu/simd/base_simd.h	The helper, and the scalar floor_divide for float16 and bfloat16.
mlx/backend/cpu/simd/accelerate_simd.h	The SIMD floor_divide for float16 and bfloat16 (one lane at a time through the helper).
mlx/backend/cpu/binary.cpp	The CPU DivMod computes the half-type quotient and remainder in float32 and rounds each once.

What does not change:

• float32, float64 and integer paths keep their code and their speed. Only the half-type branches change.
• Remainder does not change. fmod is exact, so the remainder agrees with the new quotient: a == q * b + r within the dtype rounding. The new test checks this.
• The CUDA kernels do not change. No workflow in this fork builds CUDA.

Tests

• test_divmod does not change, and it stays unseeded.
• New: test_floor_divide_half_rounding_window in python/tests/test_ops.py.
  • It builds float16 and bfloat16 pairs whose true quotient is less than a quarter ulp below an integer, so the half-type quotient certainly rounds up.
  • It checks floor_divide and the divmod quotient against the floor of the exact quotient. For float16 it also checks against numpy's divmod, both the quotient and the remainder.
  • It checks that 0 <= r < b and q * b + r == a.
  • It runs on both the CPU and the GPU stream, so the GPU and CPU test passes each cover both devices.
• Before the fix (head 7e044d98, test only, run 39971647170): python-gpu (0) and python-cpu (0) fail in all 4 subtests (float16 and bfloat16, on the CPU and the GPU). Every quotient is 1 too high, for example [5.0, 9.0, 10.0, ...] against [4.0, 8.0, 9.0, ...].
• After the fix (head 790f01b2, run 39974672718): every check passes, including the new test and test_divmod on both passes. On attempt 1, python-gpu (0) passed all its tests but failed later in "Build Metal extension example" on a DNS error while fetching nanobind from PyPI. A re-run of that job passed.

fork.yaml

New section "half and bfloat16 floor divide in float32". It names the upstream behavior and why the fork changes it.

Draft upstream issue (not opened)

Title: floor_divide and divmod for float16/bfloat16 round the quotient before the floor

mx.floor_divide and the quotient of mx.divmod compute floor(a / b) in the input dtype for float16 and bfloat16. The division result is rounded to the half type first. If the true quotient is less than half an ulp below an integer n, it rounds to n, and the result is n instead of n - 1. numpy's float16 divmod computes in float32 from fmod (npy_divmodf) and returns n - 1.

Example: a = 3.001953125 (float16), b = 1.0009765625 (float16). The true quotient is about 2.99902, so floor is 2. In float16 the quotient rounds to 3.0, and mx.floor_divide(a, b) returns 3.

This makes python/tests/test_ops.py::test_divmod flaky. It draws unseeded float16 inputs from uniform(1, 100). The rounding window near a quotient q is about q / 2950 of the range, so about 0.07% of pairs hit it. With about 56 float16 pairs per call, the test fails in about 4% of runs.

Suggested fix: compute float16 and bfloat16 floor division in float32 with the npy_divmodf method (mod = fmod(a, b), div = (a - mod) / b, floor with the sign correction), then round once to the output dtype. This applies to the Metal FloorDivide (and DivMod), the CPU scalar and SIMD floor_divide, the CPU DivMod, and probably the CUDA FloorDivide.

Test plan

There are no local or ai-server test runs. A compile-only Release build on ai-server (-DMLX_BUILD_TESTS=ON -DMLX_METAL_JIT=OFF) of the fix finished with 0 errors, and I removed the folder afterwards. PR CI is the check:

• Fork Tests: build-cpp, build-wheel, the C++ tests, python-gpu (0..1), python-cpu (0..2), and the aggregate metal-build-and-test
• Build and Test: Check Lint, Linux Sanitizer (ASAN, UBSAN), Linux Fedora (x86_64, aarch64)
• Fork Diff Check
• CodeQL does not run on this PR, because its base is not main.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes (I ran the same hook versions on the changed files: clang-format 21.1.8, ruff 0.16.10.)
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed) (fork.yaml)
• I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: An AI agent (Claude Opus 5.5 in Claude Code) wrote the fix, the test, the commits and this description. David directed the work. The fork inherits upstream's rule against AI-written PR text, so the box above is not ticked. David decides.

🤖 Generated with Claude Code