Skip to content

Fix kernel cache collision in Compiled constructor - #3427

Merged
angeloskath merged 1 commit into
ml-explore:mainfrom
dogukanveziroglu:fix-compile-cast-kernel-cache
Apr 21, 2026
Merged

angeloskath merged 1 commit into
ml-explore:mainfrom
dogukanveziroglu:fix-compile-cast-kernel-cache

Conversation

@dogukanveziroglu

Copy link
Copy Markdown
Contributor

The constructor moves inputs and constant_ids into members but the loop building kernel_lib_ continued to iterate the moved-from parameters. After the move those parameters are empty, so the kernel name omits the per-input dtype suffix and a kernel compiled for one input type can be incorrectly reused for another. With ops fused around the cast (e.g. chained astype feeding matmul or einsum), this surfaces as byte-level corruption of the output.

Reproducer:

import mlx.core as mx

def f(x_u16, y_i8):
    a = x_u16.astype(mx.int32)
    b = y_i8.astype(mx.int32)
    af = a.astype(mx.float32)
    bf = b.astype(mx.float32)
    return mx.matmul(bf.reshape(2, 1), af.reshape(1, 2))

x = mx.array([10, 20], dtype=mx.uint16)
y = mx.array([1, 2], dtype=mx.int8)
print(f(x, y))                 # [[10, 20], [20, 40]]
print(mx.compile(f)(x, y))     # [[5130, 10260], [0, 0]]  <-- wrong

5130 = 0x140A is the little-endian byte pattern of the two uint16 inputs packed into one int32, confirming the kernel was specialized for the wrong dtype.

Switching the loop to use the inputs_ / constant_ids_ members restores the dtype suffix and the bug goes away.

Verified locally against the issue reproducer, exhaustive dtype matrix, and fresh-process isolation tests (44/44 + 9/9 pass). I can add an inline regression test if you guys would like a specific format / location.

Fixes #3338.

Checklist

Put an x in the boxes that apply.

  • 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 have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

The constructor moves `inputs` and `constant_ids` into members but the
loop building `kernel_lib_` continued to iterate the moved-from parameters.
After the move those parameters are empty, so the kernel name omits the
per-input dtype suffix and a kernel compiled for one input type can be
incorrectly reused for another. With ops fused around the cast (e.g.
chained `astype` feeding `matmul` or `einsum`), this surfaces as
byte-level corruption of the output.

Reproducer:

```python
import mlx.core as mx

def f(x_u16, y_i8):
    a = x_u16.astype(mx.int32)
    b = y_i8.astype(mx.int32)
    af = a.astype(mx.float32)
    bf = b.astype(mx.float32)
    return mx.matmul(bf.reshape(2, 1), af.reshape(1, 2))

x = mx.array([10, 20], dtype=mx.uint16)
y = mx.array([1, 2], dtype=mx.int8)
print(f(x, y))                 # [[10, 20], [20, 40]]
print(mx.compile(f)(x, y))     # [[5130, 10260], [0, 0]]  <-- wrong
```

`5130 = 0x140A` is the little-endian byte pattern of the two `uint16`
inputs packed into one `int32`, confirming the kernel was specialized
for the wrong dtype.

Switching the loop to use the `inputs_` / `constant_ids_` members
restores the dtype suffix and the bug goes away.

Fixes ml-explore#3338.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great found, thanks for fixing this!

@angeloskath
angeloskath merged commit c594e6e into ml-explore:main Apr 21, 2026
16 checks passed
@BrewTestBot BrewTestBot mentioned this pull request Apr 22, 2026
1 task done
@dogukanveziroglu
dogukanveziroglu deleted the fix-compile-cast-kernel-cache branch April 22, 2026 12:37
rockiestar-com pushed a commit to RockieStar-Inc/mlx that referenced this pull request Aug 25, 2026
eejd pushed a commit to eejd/mlx-swift that referenced this pull request Oct 18, 2026
Apply ml-explore/mlx@c594e6ec38 (ml-explore/mlx#3427) to the vendored
MLX. Compiled::Compiled built the kernel name from the moved-from
inputs, so graphs that differ only in input dtype got the same name and
shared one Metal library.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] mx.compile produces wrong results when chaining unsigned->signed astype with matmul

3 participants