Skip to content

fix(metal): size gather_qmm_rhs by the broadcast x rows - #2138

Open
inureyes wants to merge 11 commits into
mainfrom
fix/issue-1599-metal-gather-qmm-rhs-rows
Open

inureyes wants to merge 11 commits into
mainfrom
fix/issue-1599-metal-gather-qmm-rhs-rows

Conversation

@inureyes

@inureyes inureyes commented Oct 6, 2026 •

Copy link
Copy Markdown
Member

Fix Metal gather_qmm_rhs to use the broadcast output row count, so shared activations gathered across sorted slots write every output row. Make two split-MTP tests respect canonical macOS paths and filesystems that reject invalid-byte names; production path and symlink safeguards are unchanged.

The branch includes main 4b0c4519 and MLX pin 59d600b5e. The #1599 quantized.cpp overlay preserves the pin's global-scale changes and the B * M row-count fix.

Candidate Metal nightly 38430123889 failed while compiling the pinned MLX sources at source head 4323eb39, before the #1599 numerical tests ran. Xcode 27 rejected the thread address-space qualifiers carried by decltype(ct_a) and decltype(ct_b) into three gated-delta MPP destination cooperative-tensor factories. The follow-up overlay strips those qualifiers with metal::remove_addrspace_t, matching the pinned tree's working steel NAX helpers and the compatibility issue addressed by upstream MLX #4594. It does not change the cooperative-tensor dtype, descriptor, mode, or algorithm. Independent source review of the functional fix at 1b0ec371 found no correctness, security, or performance issue; head 7d545981 adds the preserved upstream Apple copyright required by the repository license contract. A new candidate Metal nightly is required; no current-head Metal pass is claimed yet.

The earlier candidate Metal nightly 38418557399 passed at source head 2dbffe6d: formatting, Clippy, 12,283 tests (390 ignored), and 51 video tests. That evidence predates the material MLX pin update from 81ba1c6a to 59d600b5e.

Linux focused portability tests pass 2/2, with 123 Python contract tests and 49 WebUI fixtures. These runtime results also predate the material MLX pin update. Independent source review of the rebased row-count overlay at 4323eb39 and the gated-delta compatibility overlay at 1b0ec371 found no source or security issue. The full Python tooling contract passes 123/123 at 7d545981, including the source-license check.

Historical checkpoint evidence remains separate: at exact head 25202198 on M1 Ultra, all five focused Metal mxfp tests passed and Qwen3-30B-A3B-4bit greedy 64-token output matched base and repeated base byte for byte (SHA-256 d2f40b2d5831c3d7790d467fb51123f0d1f615fb7300aa2f92d86a9d8831a0de). The checkpoint smoke was not rerun at the current head.

Keep #1599 open through merge. Close it only after the final integrated main Metal nightly passes.

Refs #1599

@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:high High priority platform:macos macOS (Apple Silicon) specific area:core mlxcel-core: MLX FFI, primitives, KV cache, layers labels Oct 6, 2026
inureyes added a commit that referenced this pull request Oct 6, 2026
`GatherQMM::eval_gpu` passes `x.size() / K` to `gather_qmm_rhs` as its row count. That is the row count of x before `gather_qmm_rhs` broadcasts it against the indices, so when one activation row is gathered against B sorted slots (M == 1, B >= 16, B / E >= 4) the grid and the kernel row bound cover one row and rows 1..B-1 of the freshly allocated output are never written. That is the nightly failure on the M1 Ultra runner since 2026-09-30: the `SortedSharedActivation` case of `mxfp4/mxfp8_gather_qmm_matches_host_reference` reads uninitialized memory (relative L2 7.4e32, non-finite values), while the same case passes on MLX's CPU backend in the same run.

The overlay now passes `B * M`, which equals the old value whenever x is already expanded (every `gather_sort` caller, so production MoE paths are unchanged) and covers the NAX variant too. Upstream ml-explore/mlx main carries the same line.

Not verified on Metal: this host is CUDA, where the overlay is not compiled.

Refs #1599
@inureyes
inureyes force-pushed the fix/issue-1599-metal-gather-qmm-rhs-rows branch from 838936a to 2520219 Compare October 7, 2026 14:34
@inureyes

inureyes commented Oct 7, 2026

Copy link
Copy Markdown
Member Author

Metal validation passed on Apple M1 Ultra (128 GB), macOS 27.0 (26A428), full Xcode 27 RC, after rebasing this PR onto main b34d517bc422dd0465a6c404cbb2b84d4181124d. Tested head: 25202198c3ea7444cf9aa6b14c059f79683b75ab.

  • Main reproduces both SortedSharedActivation non-finite failures: mxfp4 and mxfp8 gather, 3 passed / 2 failed. This is the existing [nightly-verify] main is red #1599 defect.
  • PR: all five models::switch_layers::mxfp_tests pass with --profile test-fast --features metal,accelerate -- --test-threads=1. The Metal quantized.cpp object was rebuilt.
  • Release build passes. A transient sccache file-descriptor exhaustion was bypassed with RUSTC_WRAPPER=; no source change was needed.
  • Qwen3-30B-A3B-4bit, greedy 64 tokens, rainbow explanation prompt: main, main-repeat, and PR generated text are byte-identical via scripts/ab_output_equality.sh; SHA-256 d2f40b2d5831c3d7790d467fb51123f0d1f615fb7300aa2f92d86a9d8831a0de.

This satisfies #2158 item 4's local gate. Keep #1599 open until a post-merge Metal nightly make verify-test passes. M5 was not tested.

@inureyes
inureyes marked this pull request as ready for review October 7, 2026 14:39
Bring PR #2138 onto current main so its Metal kernel correction and the macOS nightly portability repairs are reviewed and validated against the complete post-October integration surface.

Refs #1599
Keep the broadcast-row Metal gather fix and make split-MTP staging tests follow canonical macOS paths and the host filesystem's invalid-byte filename support. Filesystems that reject the name must now return the same I/O error without leaving staging residue.

Update the bilingual report with the exact historical M1 evidence and current Linux build, focused-test, contract, and independent-review results. A fresh Metal nightly remains required before #1599 closes.

Refs #1599
Integrate the completed Kimi Linear chunked-prefill correction before publishing the refreshed #1599 candidate. The change is disjoint from the reviewed Metal overlay and split-MTP portability tests and will be covered by PR CI and the final integrated gate.

Refs #1599
Clarify that the #1599 Linux build, focused tests, contracts, and independent review preceded the disjoint merge of main's Kimi Linear correction. PR CI and the final integrated gate cover that last sync while the historical M1 evidence remains tied to its exact head.

Refs #1599
Record the successful candidate nightly at source head 2dbffe6, including the 12,283 core/workspace tests and 51 video tests. Keep historical checkpoint smoke evidence separate and retain the post-merge main nightly requirement for #1599.
Include the merged Qwen3 Next prefill contract update before validating the Metal correction against the current default branch. The tested Metal kernel and portability test sources remain unchanged.
…gather-qmm-rhs-rows

# Conflicts:
#	src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp
Normalize local cooperative-tensor operand types before passing them to the MPP destination factory. Register the pinned-source overlay so the new MLX gated-delta kernel compiles under Xcode 27.
Carry the upstream Apple copyright in the whole-file overlay so the repository license contract recognizes its source provenance.

This branch has not been deployed

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

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:macos macOS (Apple Silicon) specific priority:high High priority status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant