diff --git a/TECHNICAL_REPORTS/pr-2138-metal-gather-qmm-rhs-rows.en.md b/TECHNICAL_REPORTS/pr-2138-metal-gather-qmm-rhs-rows.en.md new file mode 100644 index 000000000..539fab26e --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2138-metal-gather-qmm-rhs-rows.en.md @@ -0,0 +1,49 @@ +# PR #2138: Size Metal `gather_qmm_rhs` by the broadcast x rows + +**Date**: 2026-10-14 +**Status**: The kernel fix passed focused Metal and model-smoke validation on M1 Ultra at exact head `25202198`; the refreshed #1599 change passes Linux build, focused tests, contracts, and independent review and the refreshed candidate Metal nightly +**Risk**: Low (no-op for every production caller; changes behavior only for an x that `gather_qmm_rhs` has to broadcast) + +## Summary + +The Metal nightly on the M1 Ultra runner has failed since 2026-09-30 on `models::switch_layers::mxfp_tests::{mxfp4,mxfp8}_gather_qmm_matches_host_reference`, both in the `SortedSharedActivation` case. `GatherQMM::eval_gpu` in MLX's Metal backend passes `x.size() / K` to `gather_qmm_rhs` as the row count. That is the row count of x before `gather_qmm_rhs` broadcasts x against the indices. For one activation row gathered against 32 sorted slots it is 1, so the grid and the kernel's `M` bound cover one row and output rows 1..31 are never written. The overlay now passes `B * M`. The PR also makes two split-MTP tests portable across macOS path canonicalization and filesystems that reject invalid-byte names; the production resolver and its symlink protections are unchanged. + +Refs #1599. + +## 1. Nightly history + +| Dates | Failure | Status | +|---|---|---| +| 2026-09-28 | `verify-clippy`: unused `Result` in `fused_norm_parity_tests.rs:158` and `fused_rope_parity_tests.rs:133` | fixed by #2029 (d8d34e2b) | +| 2026-09-29 | `tests::family_order_is_exhaustive`: `FAMILY_ORDER` missing `Speech` | fixed by #2080 (5486e404) | +| 2026-09-30 to 10-11 | the two mxfp gather tests | fixed by this PR; focused Metal validation passed at `25202198` | +| 2026-10-12 (`1e8c7022`) | the same mxfp tests; split-MTP `/var` canonicalization and invalid-byte filename assumptions; three YAML `err_expect` lints | kernel and split-MTP tests fixed by this PR; YAML lint already fixed on current main | + +The mxfp tests arrived with #2071 (2cee9cf4) on 2026-09-30, so the last signature is a new test exposing a latent upstream edge case, not a regression. + +## 2. Root cause + +- The case uses x `[1, 1, K]`, indices `[32]` over 8 experts, `sorted = true`. With M == 1, B = 32, B / E = 4 and `right_sorted_`, `eval_gpu` takes `gather_qmm_rhs`. +- `gather_qmm_rhs` broadcasts x to 32 rows (`broadcast_with_indices`), but its `M` argument was computed by the caller from the original x: 1. `grid_dims.y = ceil(M / 16)` and the kernel's `tgp_bm = min(BM, M - y_row)` then write row 0 only. `out` comes from `allocator::malloc`, so the other rows hold whatever was there: 7.4e32 relative error (mxfp4) and non-finite values (mxfp8). +- The same runs pass the three earlier gather cases in bf16, f16 and f32, and `mxfp_matmuls_match_host_reference_on_cpu_device`, which runs `SortedSharedActivation` on MLX's CPU backend. The reference and the test are correct; "bf16" in the message is just the first dtype tried. +- The defect is mode-independent (affine would hit it too) and is still present on ml-explore/mlx main. + +## 3. Change + +`B * M` (the output row count) replaces `x.size() / K` at the one call site. It equals the old value whenever x is already expanded, which is every production caller (`SwitchLinear::forward(.., true)` always receives `gather_sort` output), and it also fixes the NAX variant, which receives the same argument. The overlay header and the CMake overlay notes list the new hunk. + +The split-MTP test for `.` and `..` outputs now compares the staging parent with the canonical temporary-directory path, matching the resolver's intentional behavior on macOS where `/var` resolves to `/private/var`. The non-UTF-8 test first probes the temporary filesystem: filesystems that accept the byte name retain the successful round trip, while rejecting filesystems must return the same I/O error and leave no output or staging entry. + +## 4. Validation + +- Historical exact-head Metal gate: Apple M1 Ultra (128 GB), macOS 27.0 (26A428), full Xcode 27 RC, PR head `25202198` on base `b34d517b`. All five `models::switch_layers::mxfp_tests` passed with the Metal overlay rebuilt, and the release build passed. +- Historical exact-head model smoke: Qwen3-30B-A3B-4bit, greedy 64 tokens, produced byte-identical text on base, repeated base, and PR; SHA-256 `d2f40b2d5831c3d7790d467fb51123f0d1f615fb7300aa2f92d86a9d8831a0de`. +- CUDA (GB10): the original PR branch built and all 32 `models::switch_layers` tests passed. This does not validate Metal. +- Refreshed #1599 Linux candidate: the exact `mlxcel` plus `mlxcel-surgery` CUDA test graph compiled successfully in 25m54s. Both focused portability tests pass (2/2), repository contracts pass (123 Python tests and 49 WebUI fixtures), and formatting and diff checks pass. +- Independent review found no source, security, or documentation issue in the three-source-file change, canonical source SHA-256 `2e36ccb93dfe3fa49206f5741249ebb45c49d3e9e19195f94f6d1282e55155a5`. +- Main `398f9ba8`, containing the independently tested Kimi Linear chunked-prefill fix, was merged afterward without overlap. The #1599 native graph was not rebuilt for that unrelated sync; PR CI and the final integrated gate cover it. +- Candidate Metal nightly [38418557399](https://github.com/lablup/mlxcel/actions/runs/38418557399) passed at source head `2dbffe6d` on `lablup-macmini01`: formatting, Clippy, 12,283 tests (390 ignored), and 51 video tests. No current-head model smoke is claimed; the historical exact-head checkpoint evidence remains separate. + +## 5. Follow-up + +Report the call to ml-explore/mlx, and drop the hunk when a pin bump carries an upstream fix. Keep #1599 open through merge and close it only after a green post-merge Metal nightly. diff --git a/TECHNICAL_REPORTS/pr-2138-metal-gather-qmm-rhs-rows.ko.md b/TECHNICAL_REPORTS/pr-2138-metal-gather-qmm-rhs-rows.ko.md new file mode 100644 index 000000000..dd3274e96 --- /dev/null +++ b/TECHNICAL_REPORTS/pr-2138-metal-gather-qmm-rhs-rows.ko.md @@ -0,0 +1,49 @@ +# PR #2138: Metal `gather_qmm_rhs`의 행 수를 broadcast된 x 기준으로 지정 + +**작성일**: 2026-10-14 +**상태**: 커널 수정은 정확한 head `25202198`에서 M1 Ultra 집중 Metal 검증과 모델 스모크를 통과. 갱신한 #1599 변경은 Linux 빌드, 집중 테스트, 계약 검사, 독립 리뷰와 갱신한 후보 Metal nightly를 통과 +**위험도**: 낮음 (모든 프로덕션 호출에서 no-op이며, `gather_qmm_rhs`가 x를 broadcast해야 하는 경우에만 동작이 바뀜) + +## 요약 + +M1 Ultra 러너의 Metal nightly는 2026-09-30부터 `models::switch_layers::mxfp_tests::{mxfp4,mxfp8}_gather_qmm_matches_host_reference`의 `SortedSharedActivation` 케이스에서 실패했다. MLX Metal 백엔드의 `GatherQMM::eval_gpu`는 `gather_qmm_rhs`에 행 수로 `x.size() / K`를 넘긴다. 이 값은 `gather_qmm_rhs`가 x를 indices에 맞춰 broadcast하기 전의 행 수다. 활성 행 하나를 정렬된 슬롯 32개에 gather하면 이 값은 1이 되어, grid와 커널의 `M` 경계가 한 행만 덮고 출력의 1..31행은 기록되지 않는다. 이 PR은 overlay에서 `B * M`을 넘긴다. 또한 macOS 경로 정규화와 유효하지 않은 바이트 이름을 거부하는 파일시스템에서 split-MTP 테스트 두 개가 이식성 있게 동작하도록 고친다. 프로덕션 resolver와 심볼릭 링크 보호는 변경하지 않는다. + +Refs #1599. + +## 1. Nightly 이력 + +| 날짜 | 실패 | 상태 | +|---|---|---| +| 2026-09-28 | `verify-clippy`: `fused_norm_parity_tests.rs:158`, `fused_rope_parity_tests.rs:133`의 미사용 `Result` | #2029 (d8d34e2b)에서 수정 | +| 2026-09-29 | `tests::family_order_is_exhaustive`: `FAMILY_ORDER`에 `Speech` 누락 | #2080 (5486e404)에서 수정 | +| 2026-09-30 ~ 10-11 | mxfp gather 테스트 두 개 | 이 PR에서 수정; `25202198`에서 집중 Metal 검증 통과 | +| 2026-10-12 (`1e8c7022`) | 같은 mxfp 테스트, split-MTP의 `/var` 정규화 및 유효하지 않은 바이트 파일명 가정, YAML `err_expect` lint 세 개 | 커널과 split-MTP 테스트는 이 PR에서 수정; YAML lint는 현재 main에서 이미 수정됨 | + +mxfp 테스트는 2026-09-30 #2071 (2cee9cf4)과 함께 추가되었으므로, 마지막 실패는 회귀가 아니라 새 테스트가 upstream의 잠재된 경계 사례를 드러낸 것이다. + +## 2. 근본 원인 + +- 케이스 입력은 x `[1, 1, K]`, expert 8개에 대한 indices `[32]`, `sorted = true`다. M == 1, B = 32, B / E = 4, `right_sorted_`이므로 `eval_gpu`는 `gather_qmm_rhs`로 간다. +- `gather_qmm_rhs`는 x를 32행으로 broadcast하지만(`broadcast_with_indices`), `M` 인자는 호출자가 원래 x로 계산한 1이다. `grid_dims.y = ceil(M / 16)`과 커널의 `tgp_bm = min(BM, M - y_row)` 때문에 0행만 기록된다. `out`은 `allocator::malloc`에서 오므로 나머지 행에는 이전 메모리 값이 남는다: mxfp4는 상대 오차 7.4e32, mxfp8은 non-finite. +- 같은 실행에서 앞선 gather 케이스 세 개는 bf16, f16, f32 모두 통과했고, MLX CPU 백엔드에서 `SortedSharedActivation`을 돌리는 `mxfp_matmuls_match_host_reference_on_cpu_device`도 통과했다. reference와 테스트는 올바르며, 메시지의 "bf16"은 처음 시도한 dtype일 뿐이다. +- 결함은 모드와 무관하고(affine도 동일), ml-explore/mlx main에도 그대로 있다. + +## 3. 변경 + +유일한 호출 지점에서 `x.size() / K`를 `B * M`(출력 행 수)으로 바꿨다. x가 이미 확장된 경우 기존 값과 같으며, 모든 프로덕션 호출(`SwitchLinear::forward(.., true)`는 항상 `gather_sort` 출력을 받음)이 이에 해당한다. 같은 인자를 받는 NAX 변형도 함께 고쳐진다. overlay 헤더와 CMake overlay 주석에 새 hunk를 기록했다. + +`.` 및 `..` 출력의 split-MTP 테스트는 staging 부모를 정규화된 임시 디렉터리 경로와 비교한다. 이는 macOS에서 `/var`가 `/private/var`로 해석되는 production resolver의 의도와 일치한다. 비 UTF-8 테스트는 임시 파일시스템의 바이트 이름 지원을 먼저 검사한다. 지원하는 파일시스템에서는 성공 round trip을 그대로 검증하고, 거부하는 파일시스템에서는 같은 I/O 오류가 반환되며 출력이나 staging 항목이 남지 않아야 한다. + +## 4. 검증 + +- 이전 정확한 head의 Metal gate: Apple M1 Ultra (128 GB), macOS 27.0 (26A428), full Xcode 27 RC, base `b34d517b` 위 PR head `25202198`. Metal overlay를 다시 빌드한 뒤 `models::switch_layers::mxfp_tests` 다섯 개가 모두 통과했고 release 빌드도 통과했다. +- 이전 정확한 head의 모델 스모크: Qwen3-30B-A3B-4bit, greedy 64 tokens에서 base, base 반복, PR이 byte-identical text를 생성했다. SHA-256은 `d2f40b2d5831c3d7790d467fb51123f0d1f615fb7300aa2f92d86a9d8831a0de`다. +- CUDA (GB10): 기존 PR 브랜치가 빌드됐고 `models::switch_layers` 테스트 32개가 모두 통과했다. 이는 Metal 검증이 아니다. +- 갱신한 #1599 Linux 후보: 정확한 `mlxcel` 및 `mlxcel-surgery` CUDA 테스트 그래프가 25분 54초에 성공적으로 컴파일됐다. 이식성 집중 테스트 두 개가 모두 통과했고, 저장소 계약 검사도 Python 테스트 123개와 WebUI fixture 49개 모두 통과했다. formatting과 diff 검사도 통과했다. +- 세 소스 파일 변경에 대한 독립 소스·보안·문서 리뷰에서 발견 사항이 없었다. 정규 source SHA-256은 `2e36ccb93dfe3fa49206f5741249ebb45c49d3e9e19195f94f6d1282e55155a5`다. +- 별도로 검증된 Kimi Linear chunked-prefill 수정이 담긴 main `398f9ba8`을 그 뒤 충돌 없이 병합했다. 이 비중첩 동기화 때문에 #1599 native graph를 다시 빌드하지 않았으며 PR CI와 최종 통합 gate가 이를 검증한다. +- 후보 Metal nightly [38418557399](https://github.com/lablup/mlxcel/actions/runs/38418557399)가 소스 head `2dbffe6d`와 `lablup-macmini01`에서 통과했다. formatting, Clippy, 테스트 12,283개(390개 ignored), 영상 테스트 51개가 통과했다. 현재 head의 모델 스모크를 주장하지 않으며 이전 정확한 head의 체크포인트 증거와 구분한다. + +## 5. 후속 작업 + +ml-explore/mlx에 이 호출을 보고하고, upstream 수정이 포함된 pin bump 때 이 hunk를 제거한다. #1599는 merge 후에도 열어 두고 post-merge Metal nightly가 통과한 뒤에만 닫는다. diff --git a/src/lib/mlx-cpp/CMakeLists.txt b/src/lib/mlx-cpp/CMakeLists.txt index c55f8f14a..28f9ab7fc 100644 --- a/src/lib/mlx-cpp/CMakeLists.txt +++ b/src/lib/mlx-cpp/CMakeLists.txt @@ -29,10 +29,17 @@ function(mlx_apply_source_overlays mlx_source_dir) # The steel/gemm/mma.h overlay (issue #217: upstream safe-load fix # cherry-pick, PRs #3560/#3565) was retired when the pin moved to # 2026-06-11 upstream main, which carries the fix natively (issue #222). - # quantized.cpp carries one delta, the MLXCEL_QMV_WIDE off-switch on - # `use_qmv_wide` (issue #1187). Upstream has no equivalent knob and the - # predicate is unchanged as of pin 81ba1c6a, so it has to be overlaid here. - # Refresh the file wholesale on a bump and re-apply that one hunk. + # quantized.cpp carries two deltas: the MLXCEL_QMV_WIDE off-switch on + # `use_qmv_wide` (issue #1187; upstream has no equivalent knob), and the + # row count `GatherQMM::eval_gpu` hands `gather_qmm_rhs`, which upstream + # takes from the pre-broadcast x so a shared activation row gathered + # against sorted slots leaves most output rows unwritten (issue #1599). + # Refresh the file wholesale on a bump and re-apply both marked hunks. + # gated_delta_update_nax.h strips the local `thread` address space from the + # cooperative-tensor operand types passed to the MPP destination factory. + # The pinned source fails that constraint under Xcode 27; upstream #4594 + # carries the same compatibility fix and postdates this pin. Drop this + # overlay when a future pin includes it. # device.cpp carries a runtime override of the per-command-buffer input # budget (`mlxcel_set_mb_per_buffer_override`), which mlxcel raises around # decode steps and leaves at the device default for prefill. Upstream latches @@ -41,6 +48,7 @@ function(mlx_apply_source_overlays mlx_source_dir) set(_metal_patch_files "mlx/backend/metal/compiled.cpp" "mlx/backend/metal/device.cpp" + "mlx/backend/metal/kernels/gated_delta_update_nax.h" "mlx/backend/metal/kernels/utils.h" "mlx/backend/metal/quantized.cpp") foreach(_patch_file ${_metal_patch_files}) diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/gated_delta_update_nax.h b/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/gated_delta_update_nax.h new file mode 100644 index 000000000..b3cf0543f --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/gated_delta_update_nax.h @@ -0,0 +1,533 @@ +// Copyright © 2026 Apple Inc. +// +// Modified from upstream MLX 59d600b5e mlx/backend/metal/kernels/gated_delta_update_nax.h +// Patched by mlxcel for Xcode 27: strip local `thread` address-space +// qualifiers before forming MPP destination cooperative tensors. This +// matches MLX's existing steel NAX helpers and upstream fix #4594. +#pragma once + +#include + +#include +#include + +#include "mlx/backend/metal/kernels/steel/gemm/nax.h" + +using namespace metal; +using namespace mpp; +using namespace mpp::tensor_ops; + +// NAX MACROS I can probably do a nice template instead of doing this +// fm = base_fm + (idx >> 2) * 8; // idx>>2 = idx/4 -> 0 for idx 0-3, 1 for +// idx 4-7 fn = base_fn + (idx % 4); // 4 consecutive columns +#define AT_NAX(TILE, IDX) TILE.elems()[IDX] + +#define SUB_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) - AT_NAX(TILE2, _i); \ + } \ + } + +#define ADD_NAX(TILE0, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + AT_NAX(TILE0, _i) = AT_NAX(TILE1, _i) + AT_NAX(TILE2, _i); \ + } \ + } + +#define FMA_NAX(TILE0, S, TILE1, TILE2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < mlx::steel::BaseNAXFrag::kElemsPerFrag; _i++) { \ + (TILE0)[_i] = (S) * (TILE1)[_i] + (TILE2)[_i]; \ + } \ + } + +#define SCALE_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + AT_NAX(TILE0, _i) *= (S); \ + } \ + } + +#define SCALE_ROW_NAX(TILE0, S) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= \ + metal::fast::exp((S)[mlx::steel::BaseNAXFrag::get_coord(_w).y]); \ + } \ + } + +#define SCALE_BETA_NAX(TILE0, BETA2) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + AT_NAX(TILE0, _i) *= (BETA2)[_w >> 2]; \ + } \ + } + +#define SCALE2_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerTile; _i++) { \ + const short _w = _i % mlx::steel::BaseNAXFrag::kElemsPerFrag; \ + const short _fm = mlx::steel::BaseNAXFrag::get_coord(_w).y; \ + AT_NAX(TILE0, _i) *= metal::fast::exp((GAMMA)[(C) - 1] - (GAMMA)[_fm]); \ + } \ + } + +#define SCALE_TRI_NAX(TILE0, GAMMA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x > _c.y) \ + ? 0.f \ + : metal::fast::exp((GAMMA)[_c.y] - (GAMMA)[_c.x]); \ + } \ + } + +#define SCALE_TRIEQ_NAX1(TILE0, BETA) \ + { \ + STEEL_PRAGMA_UNROLL \ + for (short _i = 0; _i < decltype(TILE0)::kElemsPerFrag; _i++) { \ + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ \ + AT_NAX(TILE0, _i) *= (_c.x >= _c.y) ? 0.f : (BETA)[_i >> 2]; \ + } \ + } + +namespace mlx { +namespace steel { +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A0, + const thread BaseNAXFrag::dtype_frag_t& A1, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B0, + const thread BaseNAXFrag::dtype_frag_t& B1, + metal::bool_constant) { + // M=16, N=16, K=32: A and B each two K-fragments, single 16x16 C. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 16, 32, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A0[i]; + ct_a[BaseNAXFrag::kElemsPerFrag + i] = A1[i]; + ct_b[i] = B0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = B1[i]; + ct_c[i] = C[i]; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a, + bool transpose_b, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mma( + thread BaseNAXFrag::dtype_frag_t& C, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& B, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + mpp::tensor_ops::matmul2d gemm_op; + + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = B[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + ct_c[i] = C[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = 0.0; + } + + gemm_op.run(ct_a, ct_b, ct_c); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + C[i] = ct_c[i]; + } +} + +template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false, + mpp::tensor_ops::matmul2d_descriptor::mode Mode = + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate> +METAL_FUNC static constexpr void mman( + thread BaseNAXFrag::dtype_frag_t& Cn0, + thread BaseNAXFrag::dtype_frag_t& Cn1, + const thread BaseNAXFrag::dtype_frag_t& A, + metal::bool_constant, + const thread BaseNAXFrag::dtype_frag_t& Bn0, + const thread BaseNAXFrag::dtype_frag_t& Bn1, + metal::bool_constant) { + // M=16, N=32, K=16: single A (K=16), B and C two N-fragments each. + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, 32, 16, transpose_a, transpose_b, true, Mode); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op.template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + ct_a[i] = A[i]; + ct_b[i] = Bn0[i]; + ct_b[BaseNAXFrag::kElemsPerFrag + i] = Bn1[i]; + ct_c[i] = Cn0[i]; + ct_c[BaseNAXFrag::kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BaseNAXFrag::kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[BaseNAXFrag::kElemsPerFrag + i]; + } +} + +} // namespace steel +} // namespace mlx + +#define MM16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + metal::bool_constant{}); + +#define MMA16x16x32(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mma< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + A.frag_at(0, (AO)), \ + A.frag_at(0, (AO) + 1), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MM16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +#define MMA16x32x16(C, CO, A, TA, AO, B, TB, BO) \ + mlx::steel::mman< \ + float, \ + float, \ + float, \ + TA, \ + TB, \ + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate>( \ + C.frag_at(0, (CO)), \ + C.frag_at(0, (CO) + 1), \ + A.frag_at(0, (AO)), \ + metal::bool_constant{}, \ + B.frag_at(0, (BO)), \ + B.frag_at(0, (BO) + 1), \ + metal::bool_constant{}); + +template +[[kernel]] void gated_delta_fused_nax( + const device InT* q [[buffer(0)]], + const device InT* k [[buffer(1)]], + const device InT* v [[buffer(2)]], + const device float* state_in [[buffer(3)]], + const device InT* g [[buffer(4)]], + const device InT* beta [[buffer(5)]], + device InT* y [[buffer(6)]], + device float* state_out [[buffer(7)]], + constant int& T [[buffer(8)]], + uint3 thread_position_in_grid [[thread_position_in_grid]], + uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], + uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) { + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto dv_idx = thread_position_in_grid.y * 16; + const short sg_id = thread_position_in_threadgroup.y; // 0..3 + + const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort()); + const short qid = simd_lane_id >> 2; + const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)); + + // set up pointers + // g: [B, T, Hv] + auto g_ = g + b_idx * T * Hv; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk; + auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk; + + // v, y: [B, T, Hv, Dv] + y += b_idx * T * Hv * Dv + hv_idx * Dv; + auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv; + auto beta_ = beta + b_idx * T * Hv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk; + auto o_state = state_out + (n * Dv + dv_idx) * Dk; + + threadgroup float gamma_all[C * 4]; + threadgroup float* gamma = gamma_all + sg_id * C; + + float beta_fm[2]; + + mlx::steel::NAXTile S_tile; + S_tile.load(i_state, Dk); + + mlx::steel::NAXTile K_tile, Q_tile; + mlx::steel::NAXTile W_tile; // panel + + mlx::steel::NAXTile V_tile; + mlx::steel::NAXTile U_tile; + mlx::steel::NAXTile WS_tile; + mlx::steel::NAXTile delta_tile; + mlx::steel::NAXTile tmp_tile; + mlx::steel::NAXTile QKt_tile; + mlx::steel::NAXTile out_tile; + mlx::steel::NAXTile Tinv_tile, P; + + mlx::steel::NAXTile KKtK_tile, KKt_tile; + + mlx::steel::NAXTile I_tile; + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(I_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); /* {fn, fm} */ + const short _fn = _c.x; + const short _fm = _c.y; + AT_NAX(I_tile, _i) = (_fn == _fm) ? 1.0f : 0.0f; + } + mlx::steel::NAXTile TMP_tile; + + auto process_chunk = [&](const short valid_rows, + auto bounded_tag) __attribute__((always_inline)) { + constexpr bool B = decltype(bounded_tag)::value; + + auto load_seq = [&](thread auto& tile, auto src, int ld) { + if constexpr (B) { + tile.load_rows(src, ld, valid_rows); + } else { + tile.load(src, ld); + } + }; + + float g_val = (thread_index_in_simdgroup < (uint)valid_rows) + ? metal::fast::log( + metal::max( + static_cast( + g_[thread_index_in_simdgroup * Hv + hv_idx]), + 1e-6f)) + : 0.0f; + + auto gamma_val = simd_prefix_inclusive_sum(g_val); + if (thread_index_in_simdgroup < C) { + gamma[thread_index_in_simdgroup] = static_cast(gamma_val); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + + beta_fm[0] = (fm < valid_rows) ? beta_[fm * Hv + hv_idx] : 0.0f; + const short fm1 = fm + mlx::steel::BaseNAXFrag::kElemRowsJump; + beta_fm[1] = (fm1 < valid_rows) ? beta_[fm1 * Hv + hv_idx] : 0.0f; + + KKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Dk * Hk); + MMA16x16x32(KKt_tile, 0, K_tile, false, 0, K_tile, true, 0); + } + + KKtK_tile = KKt_tile; + + SCALE_TRIEQ_NAX1(KKtK_tile, beta_fm); + SUB_NAX(Tinv_tile, I_tile, KKtK_tile); + STEEL_PRAGMA_UNROLL + for (int step = 0; step < 15; step++) { + MM16x16x16(TMP_tile, 0, KKtK_tile, false, 0, Tinv_tile, false, 0); + SUB_NAX(Tinv_tile, I_tile, TMP_tile); + } + + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < Dk / 16; nn += 2) { + load_seq(K_tile, k_ + nn * 16, Dk * Hk); + SCALE_BETA_NAX(K_tile, beta_fm); + MM16x32x16(W_tile, nn, Tinv_tile, false, 0, K_tile, false, 0); + } + SCALE_ROW_NAX(W_tile, gamma) + + SCALE_TRI_NAX(Tinv_tile, gamma) + load_seq(V_tile, v_ + dv_idx, Dv * Hv); + SCALE_BETA_NAX(V_tile, beta_fm); + MM16x16x16(U_tile, 0, Tinv_tile, false, 0, V_tile, false, 0) + + WS_tile.clear(); + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < Dk / 16; kk += 2) { + MMA16x16x32(WS_tile, 0, W_tile, false, kk, S_tile, true, kk) + } + + SUB_NAX(delta_tile, U_tile, WS_tile) + + tmp_tile.clear(); + QKt_tile.clear(); + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(Q_tile, q_ + kk, Hk * Dk); + load_seq(K_tile, k_ + kk, Hk * Dk); + + MMA16x16x32(QKt_tile, 0, Q_tile, false, 0, K_tile, true, 0); + + SCALE_ROW_NAX(Q_tile, gamma); + MMA16x16x32(tmp_tile, 0, Q_tile, false, 0, S_tile, true, kk / 16); + } + + SCALE_TRI_NAX(QKt_tile, gamma) + + out_tile = tmp_tile; + MMA16x16x16(out_tile, 0, QKt_tile, false, 0, delta_tile, false, 0); + + STEEL_PRAGMA_UNROLL + for (short _i = 0; _i < decltype(out_tile)::kElemsPerFrag; _i++) { + const short2 _c = mlx::steel::BaseNAXFrag::get_coord(_i); // {fn, fm} + const short _fn = _c.x; + const short _fm = _c.y; + if (_fm < valid_rows) { + y[_fm * Hv * Dv + dv_idx + _fn] = + static_cast(AT_NAX(out_tile, _i)); + } + } + + SCALE_NAX(S_tile, metal::fast::exp(gamma[C - 1])); + + for (int kk = 0; kk < Dk; kk += 32) { + load_seq(K_tile, k_ + kk, Hk * Dk); + SCALE2_NAX(K_tile, gamma); + MMA16x32x16(S_tile, kk / 16, delta_tile, true, 0, K_tile, false, 0); + } + }; + + int t = 0; + for (; t + C <= T; t += C) { + process_chunk(C, metal::false_type{}); + q_ += C * Hk * Dk; + k_ += C * Hk * Dk; + v_ += C * Hv * Dv; + beta_ += C * Hv; + y += C * Hv * Dv; + g_ += C * Hv; + } + if (t < T) { + process_chunk(short(T - t), metal::true_type{}); + } + + S_tile.store(o_state, Dk); +} diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp index 8b16e3563..3e8295273 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp @@ -1,11 +1,14 @@ // Copyright © 2023-2026 Apple Inc. // Patched by mlxcel: `use_qmv_wide` gains an off-switch, `MLXCEL_QMV_WIDE=0`. // Modified from upstream MLX 59d600b5e mlx/backend/metal/quantized.cpp -// Synced to upstream 59d600b5e. The delta is three hunks: the includes below, -// `mlxcel_qmv_wide_flag()` with the one call it adds to `use_qmv_wide`, and the -// two bridge entry points `mlxcel_set_qmv_wide()` / `mlxcel_qmv_wide()` at the -// end of the file. Everything else is upstream verbatim, so a bump refreshes -// this file and re-applies those three hunks. +// Synced to upstream 59d600b5e. The delta is four hunks: the includes below, +// `mlxcel_qmv_wide_flag()` with the one call it adds to `use_qmv_wide`, the +// row count `GatherQMM::eval_gpu` passes to `gather_qmm_rhs` (marked +// `mlxcel:` at the call, lablup/mlxcel#1599), and the two bridge entry points +// `mlxcel_set_qmv_wide()` / `mlxcel_qmv_wide()` at the end of the file. +// Everything else is upstream verbatim, so a bump refreshes this file and +// re-applies those four hunks. The #1599 hunk can be dropped once upstream +// stops passing `x.size() / K` there (still present at 59d600b5e). // // Why it exists (lablup/mlxcel#1186, #1187). `use_qmv_wide` sends `M >= 2` // affine quantized matmuls down `qmv_wide` on GPU generation 15 and later @@ -2049,7 +2052,16 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { transpose_, group_size_, bits_, - x.size() / K, + // mlxcel: the rows of the broadcast x, not of the input x + // (lablup/mlxcel#1599). `gather_qmm_rhs` broadcasts x against the + // indices when x carries fewer batch entries than there are slots, + // but upstream passes `x.size() / K`, the pre-broadcast row count, + // which sizes the grid and the kernel's row bound. A shared + // activation row gathered against B sorted slots then writes row 0 + // and leaves rows 1..B-1 of the fresh output buffer unwritten. + // `B * M` is the same value whenever x is already expanded (every + // `gather_sort` caller) and the right one when it is not. + B * M, N, K, d, diff --git a/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs b/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs index 1b8d531df..04b45a4dd 100644 --- a/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs +++ b/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs @@ -538,6 +538,7 @@ fn staging_directory_is_private() { #[test] fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { let tmp = tempfile::tempdir().unwrap(); + let resolved_parent = std::fs::canonicalize(tmp.path()).unwrap(); let dot_target = tmp.path().join("dot-target"); let dotdot_child = tmp.path().join("dotdot-target").join("child"); std::fs::create_dir(&dot_target).unwrap(); @@ -545,7 +546,7 @@ fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { for output in [dot_target.join("."), dotdot_child.join("..")] { let staging = StagingDir::create(&output).expect("staging"); - assert_eq!(staging.path().parent(), Some(tmp.path())); + assert_eq!(staging.path().parent(), Some(resolved_parent.as_path())); std::fs::write(staging.path().join("config.json"), b"new").unwrap(); commit_staging_dir(staging, &output, true).expect("commit"); assert_eq!(std::fs::read(output.join("config.json")).unwrap(), b"new"); @@ -554,13 +555,34 @@ fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { #[cfg(unix)] #[test] -fn non_utf8_output_name_is_supported() { +fn non_utf8_output_name_is_supported_or_rejected_cleanly() { use std::os::unix::ffi::OsStringExt; let tmp = tempfile::tempdir().unwrap(); let output = tmp.path().join(std::ffi::OsString::from_vec(vec![ b'd', b'r', 0xff, b'f', b't', ])); + + if let Err(probe_error) = std::fs::create_dir(&output) { + let staging_error = match StagingDir::create(&output) { + Ok(_) => panic!("staging must honor the filesystem rejection"), + Err(error) => error, + }; + match staging_error { + crate::SurgeryError::Io(error) => { + assert_eq!(error.kind(), probe_error.kind()); + assert_eq!(error.raw_os_error(), probe_error.raw_os_error()); + } + error => panic!("expected filesystem error, got: {error}"), + } + assert!( + std::fs::read_dir(tmp.path()).unwrap().next().is_none(), + "filesystem rejection must not leave output or staging entries" + ); + return; + } + std::fs::remove_dir(&output).unwrap(); + let staging = StagingDir::create(&output).expect("staging"); std::fs::write(staging.path().join("config.json"), b"new").unwrap();