fix(metal): build the GPT-OSS MXFP4 prefill tile into the metallib

[davidtai]

Proposed changes

This PR adds two gather_qmm_rhs kernels for the GPT-OSS MXFP4 prefill tile to mlx/backend/metal/kernels/fp_quantized.metal:

• mxfp4_gather_qmm_rhs_nt_float_gs_32_b_4_bm_32_bn_32_bk_32_wm_2_wn_2
• mxfp4_gather_qmm_rhs_nt_bfloat16_t_gs_32_b_4_bm_32_bn_32_bk_32_wm_2_wn_2

MLX_GPTOSS_MXFP4_PREFILL_TILE=m32n32k32 selects the tile bm 32, bn 32, bk 32, wm 2, wn 2 (gptoss_mxfp4_policy.h). quantized.cpp selects this tile only for mxfp4, group size 32, 4 bits, transposed weights, and float32 or bfloat16 inputs. These two kernels cover all of these cases.

Before this PR, only builds that compile jit_kernels.cpp could make this kernel. These builds make the kernel from source at run time and do not read mlx.metallib. A CMake build with the default MLX_METAL_JIT=OFF (CMakeLists.txt:44) compiles nojit_kernels.cpp instead. It loads kernels only from mlx.metallib, so it failed with "Unable to load function". The C++ test gptoss mxfp4 prefill tiles preserve expert boundaries and partial rows (tests/gptoss_mxfp4_tests.cpp) showed the failure in #22 (run https://github.com/Layr-Labs/mlx/actions/runs/36172106158).

The fix goes into the metallib because every CMake build with the default MLX_METAL_JIT=OFF fails without it: the Python wheel, the C++ test binary and any other CMake consumer. d-inference does not fail without this fix. On Apple platforms its SwiftPM build compiles jit_kernels.cpp: libs/mlx-swift/Package.swift in d-inference excludes nojit_kernels.cpp on every platform (line 284) and excludes jit_kernels.cpp only on platforms without Metal (line 25). Its JIT loader builds this kernel from source at run time.

This PR does not change the test. It does not change host code, because quantized.cpp already asks for these kernel names. It adds one sentence about these kernels to the GPT-OSS section of fork.yaml.

Runner probe result

Run: https://github.com/Layr-Labs/mlx/actions/runs/36174742752 (success, 22.0 minutes). This is the Fork Tests job of #22 with the same kernel diff, before this branch moved onto main.

• Runner label: macos-26 (image macos-26-arm64, version 20260907.0351.1). macOS 26.6.2 on an Apple M1 (Virtual) with 3 cores.
• Xcode 26.6 (17F113). Swift 6.3.3.
• Build minutes: 13.5 with an empty compile cache (wheel 6.8, C++ 6.7).
• Metal C++ tests: 290 of 290 test cases passed. 3865 of 3865 assertions passed.
• Ring distributed test: 8 processes ran 13 tests each. All passed.
• Python tests: 875 tests ran in 292 seconds and passed. 52 tests were skipped by skip conditions in the test sources: 45 need PyTorch, which the action does not install; 3 are too slow for CI; 1 uses too much memory for CI; 2 are for CUDA only; 1 needs different head and value dimensions.
• Example extension: it built, and its test passed on the CPU and the GPU.

Test plan

• git diff of the kernel commit before and after the rebase onto main: identical. It adds 4 lines to fp_quantized.metal.
• The instantiation macro gives the kernel name mxfp4_gather_qmm_rhs_nt_float_gs_32_b_4_bm_32_bn_32_bk_32_wm_2_wn_2. This is the name in the failed test log of ci: build with Metal and run the tests on pull requests #22.
• /usr/bin/python3 scripts/check_forkdiff.py --upstream-ref refs/remotes/upstream/main on the head commit: 172 of 172 changed files described.
• git merge-tree --write-tree of this branch with ci/fork-tests and with ci/drop-upstream-workflows: no conflicts.
• Fork Tests without the fix, on main plus the CI files of ci: build with Metal and run the tests on pull requests #22: https://github.com/Layr-Labs/mlx/actions/runs/36172106158. 289 of 290 C++ test cases pass. The prefill tile test fails.
• Fork Tests with the same kernel diff: https://github.com/Layr-Labs/mlx/actions/runs/36174742752. C++ 290 of 290, ring test 8 × 13, Python 875 passed and 52 skipped, example extension passed.
• Build and Test (https://github.com/Layr-Labs/mlx/actions/runs/36187633589), Fork Diff Check (https://github.com/Layr-Labs/mlx/actions/runs/36187633609) and CodeQL (https://github.com/Layr-Labs/mlx/actions/runs/36187631812) pass on the head commit, with main as the base.
• gh api repos/Layr-Labs/mlx/commits/<sha> --jq .commit.verification.verified gives true for both commits.

CONTRIBUTING checklist

• Formatting: pre-commit ran on the head commit. The Check Lint job passed in run https://github.com/Layr-Labs/mlx/actions/runs/36187633589. pre-commit does not format .metal files.
• Tests: the existing test gptoss mxfp4 prefill tiles preserve expert boundaries and partial rows in tests/gptoss_mxfp4_tests.cpp fails on ci/fork-tests without the fix (run https://github.com/Layr-Labs/mlx/actions/runs/36172106158) and passes with it (run https://github.com/Layr-Labs/mlx/actions/runs/36174742752).
• Benchmarks: this PR adds two kernel instantiations that the default path never selects. Only MLX_GPTOSS_MXFP4_PREFILL_TILE=m32n32k32 selects them. No benchmark is needed.

Dependencies

The base branch is main. This PR merges before #22. Its evidence is run https://github.com/Layr-Labs/mlx/actions/runs/36172106158 (red without the fix, on main plus the CI files of #22 only) and run https://github.com/Layr-Labs/mlx/actions/runs/36174742752 (green with the identical kernel diff). main has no Fork Tests workflow yet, so the checks of this PR do not build with Metal. The next Fork Tests run of #22, after #22 is updated from main, tests this change on main.

• ☐ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: this change and this description were written by Claude (Opus 5.5) under David Tai's direction and review.

🤖 Generated with Claude Code