Repository navigation
Add gather_qqmm - #3757
Merged
Merged
Add gather_qqmm#3757
Conversation
zcbenz
force-pushed
the
qmm-global-scale
branch
3 times, most recently
from
June 30, 2026 07:39
ca36a2d to
8180ce5
Compare
zcbenz
force-pushed
the
qmm-global-scale
branch
3 times, most recently
from
July 1, 2026 01:19
b396742 to
6a43d36
Compare
nastya236
self-requested a review
July 23, 2026 14:37
Collaborator
|
Looks great, thank you! |
nastya236
approved these changes
Jul 23, 2026
zcbenz
force-pushed
the
qmm-global-scale
branch
from
August 4, 2026 07:38
6811cde to
1d8f3b5
Compare
zcbenz
force-pushed
the
qmm-global-scale
branch
from
August 4, 2026 07:40
1d8f3b5 to
d85ac75
Compare
3 of 4 tasks
7 tasks done
inureyes
added a commit
to lablup/mlxcel
that referenced
this pull request
Aug 7, 2026
## Summary `loading::vlm::gemma_unified::tests::unified_sanitize_quantized_split_dequant_equivalence` was the only failing test on `main`. The root cause is a real defect, but not in the gemma4_unified sanitize split: `mlx::core::dequantize` silently returns wrong values whenever its `biases` input is not row-contiguous. The test expectation (exact equality, `assert_eq!(max_abs_diff(...), 0.0)`) was sound and is unchanged. ## Root cause The sanitize split is correct. Slicing the fused quantized `gate_up_proj` on the output axis selects exactly the right `weight` / `scales` / `biases` bytes: dequantizing a row-contiguous copy of the sanitize output reproduces `split(dequantize(W))` bit for bit on both legs. Isolating the three inputs one at a time shows the result is wrong if and only if `biases` is strided, regardless of the layout of `weight` and `scales`. `quantize_impl` on the MLX Metal backend (<https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp>) binds `w` (buffer 0) and `scales` (buffer 1) on the shared compute encoder and only then calls `ensure_row_contiguous` on `biases`. For a strided `biases`, that call enqueues a copy kernel on the same encoder, and the copy rebinds buffers 0 and 1 to its own operands. The dequantize dispatch that follows reads `w` and `scales` from the copy's operands instead of its own, so the affine kernel computes `biases * (q + 1)` with `q` decoded from the raw bias bytes. That model matches the observed numbers exactly: the gate leg's minimum was `-3.344`, which is `min(biases) * 16`, and its maximum was `-0.1793`, which is `max(biases) * 1`. This is an upstream ordering defect, not a contract on the caller. `QuantizedMatmul::eval_gpu` in the same file hoists every `ensure_row_contiguous*` above its first binding, and so did `fast::Quantize::eval_gpu` until ml-explore/mlx#3757 folded it into the shared `quantize_impl`. ## Regression window The test has been red since `94a8608a` (#1046, 2026-08-06), which bumped the MLX pin from `b7c3dd6d` (2026-07-17) to `2c46b953` (2026-08-05). The upstream commit that reordered the bindings, `a79332de6` (ml-explore/mlx#3757), is an ancestor of `2c46b953` and is not an ancestor of `b7c3dd6d`; the pre-bump tree still has all three contiguity copies hoisted above the first `set_input_array`. The test passed from the day it was added in #156 until that pin bump, so it was red for two days, not silently red for months. The window was established by reading `fast::Quantize::eval_gpu` in both pinned MLX trees (both are present in the local FetchContent clone) rather than by rebuilding mlxcel at the old pin, which would have cost a cold MLX C++ build. ## What changed - `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp`: the `dequantize` shim forces `biases` row-contiguous before calling `mlx::core::dequantize`. `biases` is the only input whose contiguity copy is enqueued after a binding, so it is the only one guarded, and `mlx::core::contiguous` shares the buffer instead of copying when the input is already row-contiguous, so the common path is unchanged. The comment records the mechanism and the upstream reference for the next MLX bump. - `src/lib/mlxcel-core/src/ffi_tests.rs`: adds `dequantize_commutes_with_output_axis_slice_on_strided_inputs`, which pins the guard at the FFI boundary where the defect actually lives, covering both halves of the split. - `src/loading/vlm_gemma_unified.rs`: doc-comment only. `split_fused_gate_up_component` keeps emitting strided views rather than materializing copies, because `gather_qmm` and `quantized_matmul` receive per-input batch strides explicitly and consume the views zero-copy; materializing would add a second copy of every fused expert tensor at load time for no correctness gain. The comment now says consumers must not assume row-contiguity there. ## Notes - Both legs confirmed. The `up_proj` assertion that the earlier failure masked also passes; before the fix it was wrong by 4.0477 against an expected 0. - `mistral4` and `qwen3_5_moe` split fused MoE experts the same way and therefore also hand strided quantized triplets downstream. They are not changed here: the guard is at the `dequantize` boundary, so it covers every caller. - The other in-tree `mlx::core::dequantize` call sites (`mlx_cxx_nemotron.cpp`, the fused-MoE parity reference in `mlx_cxx_kernels.cpp`, and `quantized_embedding`) feed `biases` from `take` or `astype`, both of which produce row-contiguous output, so none of them were exposed. ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --lib unified_sanitize` (6 passed, 0 failed; was 5 passed, 1 failed) - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel-core --lib dequantize_commutes_with_output_axis_slice_on_strided_inputs` (passed) - [x] Negative control: the same test fails with a 3.71 max absolute difference when the guard is reverted, so it is not vacuous - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --test granite4_vision_parity` (3 passed, real 4-bit checkpoint) - [x] `cargo fmt --check` - [x] `cargo clippy --profile test-fast --features metal,accelerate -p mlxcel-core --lib --tests -- -D warnings` - [x] `make verify-test` Closes #1074 ## Review follow-ups A second commit adds comments only, no behavior change: the three call sites that reach `mlx::core::dequantize` directly (`quantized_embedding` in the bridge, the nemotron embedding path, and the fused-MoE parity reference) now record inline that they are safe only because `biases` comes from `take` / `astype`; the shim's cost claim is corrected (`mlx::core::contiguous` always adds a `Contiguous` node, and the elision happens at eval and only for a standalone row-contiguous input); the shim records why the guard is unconditional rather than gated on `flags().row_contiguous` (flags are not meaningful before eval, so the obvious short-circuit would skip the copy exactly when it is needed); and the Rust declaration states the normalization where callers look. ## Not covered - **`gather_qmm` on strided expert weights is not pinned by a test.** This PR deliberately keeps `split_fused_gate_up_component` emitting strided views, and that decision rests on `QuantizedMatmul::eval_gpu` and `GatherQMM::eval_gpu` hoisting every `ensure_row_contiguous_matrix` above their first binding and passing per-input batch strides explicitly. That was verified by reading the pinned MLX source, not by a test. A future MLX bump that regresses it the way ml-explore/mlx#3757 regressed `dequantize` would corrupt quantized MoE inference for gemma4_unified, mistral4, and qwen3_5_moe silently. A companion FFI test comparing `gather_qmm` on sliced expert weights against `gather_qmm` on materialized copies would close that gap and is worth a follow-up issue. - **The MLX bump checklist was not updated.** `CLAUDE.md` carries the post-bump re-validation list and would be the right home for "re-check the `dequantize` biases guard, regression test `dequantize_commutes_with_output_axis_slice_on_strided_inputs`". It is gitignored in this repository, so the entry cannot land through this PR and is left to the maintainer. ## Observed while verifying `make verify-test` was run twice on this branch. Both runs were green except for one real-model VLM parity test asserting `text-only logits must be finite`, and it was a **different** test each time (`granite4_vision_parity` in the first run, `hunyuan_vl_parity` in the second). Each passes in isolation, including its full test binary. That looks like load-dependent flakiness in the real-checkpoint parity tests when the whole workspace runs concurrently, not a regression from this change; worth its own issue.
1 task done
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
## Summary `loading::vlm::gemma_unified::tests::unified_sanitize_quantized_split_dequant_equivalence` was the only failing test on `main`. The root cause is a real defect, but not in the gemma4_unified sanitize split: `mlx::core::dequantize` silently returns wrong values whenever its `biases` input is not row-contiguous. The test expectation (exact equality, `assert_eq!(max_abs_diff(...), 0.0)`) was sound and is unchanged. ## Root cause The sanitize split is correct. Slicing the fused quantized `gate_up_proj` on the output axis selects exactly the right `weight` / `scales` / `biases` bytes: dequantizing a row-contiguous copy of the sanitize output reproduces `split(dequantize(W))` bit for bit on both legs. Isolating the three inputs one at a time shows the result is wrong if and only if `biases` is strided, regardless of the layout of `weight` and `scales`. `quantize_impl` on the MLX Metal backend (<https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp>) binds `w` (buffer 0) and `scales` (buffer 1) on the shared compute encoder and only then calls `ensure_row_contiguous` on `biases`. For a strided `biases`, that call enqueues a copy kernel on the same encoder, and the copy rebinds buffers 0 and 1 to its own operands. The dequantize dispatch that follows reads `w` and `scales` from the copy's operands instead of its own, so the affine kernel computes `biases * (q + 1)` with `q` decoded from the raw bias bytes. That model matches the observed numbers exactly: the gate leg's minimum was `-3.344`, which is `min(biases) * 16`, and its maximum was `-0.1793`, which is `max(biases) * 1`. This is an upstream ordering defect, not a contract on the caller. `QuantizedMatmul::eval_gpu` in the same file hoists every `ensure_row_contiguous*` above its first binding, and so did `fast::Quantize::eval_gpu` until ml-explore/mlx#3757 folded it into the shared `quantize_impl`. ## Regression window The test has been red since `94a8608a` (lablup#1046, 2026-08-06), which bumped the MLX pin from `b7c3dd6d` (2026-07-17) to `2c46b953` (2026-08-05). The upstream commit that reordered the bindings, `a79332de6` (ml-explore/mlx#3757), is an ancestor of `2c46b953` and is not an ancestor of `b7c3dd6d`; the pre-bump tree still has all three contiguity copies hoisted above the first `set_input_array`. The test passed from the day it was added in lablup#156 until that pin bump, so it was red for two days, not silently red for months. The window was established by reading `fast::Quantize::eval_gpu` in both pinned MLX trees (both are present in the local FetchContent clone) rather than by rebuilding mlxcel at the old pin, which would have cost a cold MLX C++ build. ## What changed - `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp`: the `dequantize` shim forces `biases` row-contiguous before calling `mlx::core::dequantize`. `biases` is the only input whose contiguity copy is enqueued after a binding, so it is the only one guarded, and `mlx::core::contiguous` shares the buffer instead of copying when the input is already row-contiguous, so the common path is unchanged. The comment records the mechanism and the upstream reference for the next MLX bump. - `src/lib/mlxcel-core/src/ffi_tests.rs`: adds `dequantize_commutes_with_output_axis_slice_on_strided_inputs`, which pins the guard at the FFI boundary where the defect actually lives, covering both halves of the split. - `src/loading/vlm_gemma_unified.rs`: doc-comment only. `split_fused_gate_up_component` keeps emitting strided views rather than materializing copies, because `gather_qmm` and `quantized_matmul` receive per-input batch strides explicitly and consume the views zero-copy; materializing would add a second copy of every fused expert tensor at load time for no correctness gain. The comment now says consumers must not assume row-contiguity there. ## Notes - Both legs confirmed. The `up_proj` assertion that the earlier failure masked also passes; before the fix it was wrong by 4.0477 against an expected 0. - `mistral4` and `qwen3_5_moe` split fused MoE experts the same way and therefore also hand strided quantized triplets downstream. They are not changed here: the guard is at the `dequantize` boundary, so it covers every caller. - The other in-tree `mlx::core::dequantize` call sites (`mlx_cxx_nemotron.cpp`, the fused-MoE parity reference in `mlx_cxx_kernels.cpp`, and `quantized_embedding`) feed `biases` from `take` or `astype`, both of which produce row-contiguous output, so none of them were exposed. ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --lib unified_sanitize` (6 passed, 0 failed; was 5 passed, 1 failed) - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel-core --lib dequantize_commutes_with_output_axis_slice_on_strided_inputs` (passed) - [x] Negative control: the same test fails with a 3.71 max absolute difference when the guard is reverted, so it is not vacuous - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --test granite4_vision_parity` (3 passed, real 4-bit checkpoint) - [x] `cargo fmt --check` - [x] `cargo clippy --profile test-fast --features metal,accelerate -p mlxcel-core --lib --tests -- -D warnings` - [x] `make verify-test` Closes lablup#1074 ## Review follow-ups A second commit adds comments only, no behavior change: the three call sites that reach `mlx::core::dequantize` directly (`quantized_embedding` in the bridge, the nemotron embedding path, and the fused-MoE parity reference) now record inline that they are safe only because `biases` comes from `take` / `astype`; the shim's cost claim is corrected (`mlx::core::contiguous` always adds a `Contiguous` node, and the elision happens at eval and only for a standalone row-contiguous input); the shim records why the guard is unconditional rather than gated on `flags().row_contiguous` (flags are not meaningful before eval, so the obvious short-circuit would skip the copy exactly when it is needed); and the Rust declaration states the normalization where callers look. ## Not covered - **`gather_qmm` on strided expert weights is not pinned by a test.** This PR deliberately keeps `split_fused_gate_up_component` emitting strided views, and that decision rests on `QuantizedMatmul::eval_gpu` and `GatherQMM::eval_gpu` hoisting every `ensure_row_contiguous_matrix` above their first binding and passing per-input batch strides explicitly. That was verified by reading the pinned MLX source, not by a test. A future MLX bump that regresses it the way ml-explore/mlx#3757 regressed `dequantize` would corrupt quantized MoE inference for gemma4_unified, mistral4, and qwen3_5_moe silently. A companion FFI test comparing `gather_qmm` on sliced expert weights against `gather_qmm` on materialized copies would close that gap and is worth a follow-up issue. - **The MLX bump checklist was not updated.** `CLAUDE.md` carries the post-bump re-validation list and would be the right home for "re-check the `dequantize` biases guard, regression test `dequantize_commutes_with_output_axis_slice_on_strided_inputs`". It is gitignored in this repository, so the entry cannot land through this PR and is left to the maintainer. ## Observed while verifying `make verify-test` was run twice on this branch. Both runs were green except for one real-model VLM parity test asserting `text-only logits must be finite`, and it was a **different** test each time (`granite4_vision_parity` in the first run, `hunyuan_vl_parity` in the second). Each passes in isolation, including its full test binary. That looks like load-dependent flakiness in the real-checkpoint parity tests when the whole workspace runs concurrently, not a regression from this change; worth its own issue.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add a simple
gather_qqmmimplementation, by supporting global scale inqmmkernels and use them as fallback.The implementation is bare minimum, for CUDA it is rerouted to
qmm_naiveand for Metal it isqmv. We will add global scale support to moreqmmkernels in followup PRs, and will have specialized fast kernels in future. This PR mostly serves as a way for users to useqqmmbefore we have the native kernels ready.Support for CPU backend is lacked at the moment..