Skip to content

prepare for nax jit / mlx v0.30.1#319

Merged
davidkoski merged 7 commits intomainfrom
nax-jit
Jan 8, 2026
Merged

prepare for nax jit / mlx v0.30.1#319
davidkoski merged 7 commits intomainfrom
nax-jit

Conversation

@davidkoski
Copy link
Collaborator

@davidkoski davidkoski commented Dec 17, 2025

Proposed changes

mlx v0.30.1 which includes the NAX JIT work.

This still requires a tag on mlx-c and updating the CMakeLists.txt to reflect that.

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)

@davidkoski davidkoski changed the title prepare for nax jit prepare for nax jit / mlx v0.30.1 Dec 19, 2025
@davidkoski davidkoski requested a review from awni December 19, 2025 16:33
.target(
name: "MLXNN",
dependencies: ["MLX", "MLXRandom", "MLXFast"],
dependencies: ["MLX"],
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

These dependencies are obsolete, clean up.

///
/// ### See Also
/// - ``depends(inputs:dependencies:)``
public func depends(input: MLXArray, dependencies: [MLXArray]) -> MLXArray {
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Noticed this was missing

/// - bits: Number of bits used to represent each element of `x` and `w`
/// - mode: The quantization mode. Default is `.affine`
/// - stream: Stream or device to evaluate on
public func qqMatmul(
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

New in 0.30.1

@davidkoski
Copy link
Collaborator Author

Note: cmake builds will fail until mlx-c is tagged and we pick up the tag.

Comment on lines +2313 to +2317
/// - w: weight matrix. If quantized, it is packed in unsigned integers.
/// - scales: The scales to use per `groupSize` elements of `w` if `w` is quantized
/// - groupSize: Number of elements in `x` and `w` that share a scale
/// - bits: Number of bits used to represent each element of `x` and `w`
/// - mode: The quantization mode. Default is `.affine`
Copy link
Member

Choose a reason for hiding this comment

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

The default here should be "nvfp4" (affine doesn't make sense as it's not supported in this op).

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Ah, I missed the new quantization modes, glad this showed up!

@@ -215,6 +215,7 @@ Note: the `-` and `/` operators are not able to be linked here.
- ``addMM(_:_:_:alpha:beta:stream:)``
- ``quantizedMatmul(_:_:scales:biases:transpose:groupSize:bits:mode:stream:)``
- ``gatherQuantizedMatmul(_:_:scales:biases:lhsIndices:rhsIndices:transpose:groupSize:bits:mode:sortedIndices:stream:)``
- ``qqMatmul(_:_:scales:groupSize:bits:mode:stream:)``
Copy link
Member

@awni awni Jan 8, 2026

Choose a reason for hiding this comment

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

This is a bit of a nit, but the naming here is inconsistent with itself (and inconsistent with MLX core which is also inconsistent with itself 😓)

  • In some cases Quantized is expanded in some it's just q (should it be Q?)
  • In some cases we have MM and in others Matmul

Unclear why it's different from MLX core in the conventions and the pattern is also somewhat random.

Copy link
Member

Choose a reason for hiding this comment

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

We don't have to fix this now but I want to flag it as I could see it becoming messier as we add more ops.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Per offline discussion we will do:

  • matmul
  • gatherMM
  • blockMaskedMM
  • addMM
  • quantizedMM
  • gatherQuantizedMM
  • quantizedQuantizedMM

Copy link
Member

@awni awni left a comment

Choose a reason for hiding this comment

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

Looks awesome, minor comments. O/w good to go!

@davidkoski davidkoski merged commit 07407d9 into main Jan 8, 2026
13 of 14 checks passed
@davidkoski davidkoski deleted the nax-jit branch January 8, 2026 21:26
@@ -1006,6 +1006,11 @@ public protocol UnaryLayer: Module {
func callAsFunction(_ x: MLXArray) -> MLXArray
}

/// A `Layer` (``Module`` subclass) that can be evaluated with an array and offset.

Choose a reason for hiding this comment

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

@davidkoski this declaration is from an earlier version and has made its way back.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Sorry, must have messed it up in a rebase. Let me get a PR out for it.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Ugh, I worked up OffsetLayer as well somehow ... oh shoot, I had it in both PRs and moved it in #322 . Which is #329 !!

davidkoski added a commit that referenced this pull request Jan 8, 2026
- build error as OffsetLayer was added in both #322 and #319
- #319 had it there "temporarily" for development but it should have been removed
- and #322 moved it to a new location so now we have two
- each PR built but the merge does not
@davidkoski davidkoski mentioned this pull request Jan 8, 2026
4 tasks
davidkoski added a commit that referenced this pull request Jan 9, 2026
- build error as OffsetLayer was added in both #322 and #319
- #319 had it there "temporarily" for development but it should have been removed
- and #322 moved it to a new location so now we have two
- each PR built but the merge does not
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.

3 participants