Skip to content

JIT quantized-kernel fixes; adjacent minor underscore fix - #4372

Open
dwijenpatel wants to merge 1 commit into
ml-explore:mainfrom
dwijenpatel:fix-qmv-jit-template
Open

JIT quantized-kernel fixes; adjacent minor underscore fix#4372
dwijenpatel wants to merge 1 commit into
ml-explore:mainfrom
dwijenpatel:fix-qmv-jit-template

Conversation

@dwijenpatel

@dwijenpatel dwijenpatel commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Small PR fixing #4350

  • the added results_per_simdgroup param is unused, but needed to fix the compilation issue, and the underlying implementations hardcode 4.
  • fp_qmv_fast was mentioned in the issue, but has the correct param count (and actually uses results_per_simdgroup). The other 3 mentioned are fixed in this PR.
  • Running the full test_quantized suite under JIT surfaced two pre-existing, unrelated errors - removing underscores in two get_qmm_nax_kernel_wrapped calls to reference correct function names, re-ran tests to verify all fixes (37/37, AOT unchanged).
  • Adding a MLX_METAL_JIT=ON job to CI might be worth doing at some point, to proactively surface this class of bugs.

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure:
    Pair-programming with Claude, attempting to do useful work as a side-effect of deep-diving into GPU kernel optimization.

@dwijenpatel
dwijenpatel marked this pull request as ready for review August 21, 2026 23:47
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.

1 participant