Skip to content

fix(mlx-cpp): stop casting comparison inputs to bool in compiled kernels #1792

Description

@inureyes

Problem

The mlxcel overlay of MLX's Metal compiled-kernel generator casts each input whose type differs from the op's output type to the output type. It exists for mixed bf16/f32 arithmetic (commit 1923da3: Divide(bfloat16_t, float) failed template deduction on macOS 26.4). A comparison (Equal, NotEqual, Less, LessEqual, Greater, GreaterEqual) outputs bool, so every non-bool input is cast to bool first, same-dtype inputs included: Less(0.2, 0.7) becomes Less(true, true), which is false. isnan and isinf lower to NotEqual / Equal and break too. The result is silently wrong, never an error. CUDA has no compiled.cpp overlay (none under patches/mlx/backend/cuda/ or patches-cuda/) and is unaffected.

Latent today: no compiled function in the tree contains a comparison (every mlx::core::compile and compile_shapeless_audited site in mlx_cxx_bridge.cpp and mlx_cxx_kernels.cpp checked), and the #1392 audit found all 20 shapeless sites equal to eager on Metal. The compiled min-p filter removed in PR #1391 (#1379) had this shape and returned its input unfiltered on Metal, mechanism undiagnosed (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5483-5493); this cast reproduces that symptom. Found during review of PR #1772.

Evidence

  • src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp:2 (overlay header) and :220-244 (the cast keyed on out_type, which is the overlay's whole delta against upstream 81ba1c6a).
  • Probe against the in-tree libmlx.a (81ba1c6a, M1 Ultra), compiled vs eager: where(less(a, b), a, b) with f32 a = [0.2, -0.5, 0.7, 0, 3], b = [0.7, 0.3, 0.2, 1, 3] gives [0.7, 0.3, 0.2, 0, 3] instead of [0.2, -0.5, 0.2, 0, 3]; greater(2a, b) is all false; an f16/f32 comparison diverges; where(isnan(x), 7, x) passes NaN through; the min-p graph (softmax, max * 0.3, greater_equal, where) returns its input unfiltered. A bf16/f32 arithmetic chain still matches, and upstream MLX 0.32.2 matches eager on every graph.

Proposed fix

Exclude bool-output primitives from the output-type cast, or cast mismatched inputs to the promoted type of the op's inputs instead of the output type. Either keeps the arithmetic case the cast exists for. Keep the overlay delta small, since every MLX pin bump three-way merges it.

Acceptance criteria

  • Compiled comparisons on same-dtype f32 inputs and on mixed f16/f32 inputs, and a compiled isnan, match the uncompiled result.
  • The mixed bf16/f32 arithmetic case still compiles and matches eager.

Verification

A new mlxcel-core test that fails before the fix (through a test-only bridge entry, since no generic compile call is exposed to Rust), then cargo test --workspace --profile test-fast --features metal,accelerate on a Metal host.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layerspriority:mediumMedium prioritystatus:readyReady to be worked ontype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions