Repository navigation
Regenerate from numen-tech/mlx implied-bias quantized matmul; QuantizedEmbedding passes a 0-d bias through - #8
Merged
Merged
Conversation
…l, gemma4-qat ml-explore#190) Bumps Source/Cmlx/mlx to numen-tech/mlx fork/190-implied-bias head 0e7a5b09: a 0-d `biases` array passed to `quantized_matmul`/`dequantize` in mode "affine" is an implied-bias factor f, and the kernels form `bias = f * scale` in place of loading a per-group bias (bit-identical to a materialized bias, including fp16 subnormal scales). tools/update-mlx.sh (run with LC_ALL=C) regenerates the `_ib` kernel variants into mlx-generated/metal/quantized{,_nax}.h and mlx-generated/quantized{,_nax}.cpp; the Metal.hpp commit stamp moves. MLXNN: `QuantizedEmbedding.callAsFunction` passes a 0-d `biases` through to `dequantized` unindexed (there are no rows to gather) instead of `biases[x]`. `QuantizedLinear` already hands `biases` straight to `quantizedMM`, and `Module.update` compares shapes, so a loader that installs the factor also builds its placeholder with shape `[]`.
…ct, gemma4-qat ml-explore#190) Bumps Source/Cmlx/mlx to fork/190-implied-bias head 5a7aeace: a 0-d `biases` factor f now stands, by contract, for the materialized bias `scales * T(f)` formed in the scales' dtype T, and `quantized_matmul`/`dequantize` materialize that bias up front whenever the `_ib` kernels cannot honor it exactly -- the activations promoted the scales, or the bit width (3/5/6) has no `_ib` instantiation in the prebuilt metallib this package ships. tools/update-mlx.sh (run with LC_ALL=C) regenerates nothing but the mlx-c header copy (the new doc comment on quantized_matmul/dequantize) and the Metal.hpp commit stamp: no kernel source changed.
… CUDA, gemma4-qat ml-explore#190) ops.cpp only, so tools/update-mlx.sh leaves mlx-generated/ unchanged. Its glob emitted cuda_jit_sources.h in a different order (the same entries, sorted content identical), so the committed ordering is kept.
…e implied branch) The previous regeneration always appended the implied_bias template arg, which the JIT build could not compile for affine_sym_*/fp_* kernels (ParoQuantMLXTests crashed at affine_sym_qmv pipeline creation). No generated kernel source changes; only the submodule pin moves.
…emma4-qat ml-explore#190) ops.cpp and the CPU backend only: tools/update-mlx.sh (LC_ALL=C) regenerates nothing but the Metal.hpp commit stamp.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
…ero implied-bias factor gradient, gemma4-qat ml-explore#190)
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
1 participant
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.
For numen-tech/gemma4-qat#190 (epic numen-tech/gemma4-qat#178, Phase 2). Its fork dependency has landed: numen-tech/mlx#6 and #7 are merged into
prism-0.32.2-fixes, andSource/Cmlx/mlxnow points at that merged head, a7ed9286. Ready to merge.What
Source/Cmlx/mlxmoves to numen-tech/mlxprism-0.32.2-fixesa7ed9286 (merged Regenerate from numen-tech/mlx 2f0e44f9: re-land the counter fixes (gemma4-qat #181) #6 + Bump mlx to b2c93bd1: build and hash each custom kernel variant once, bounded memo (gemma4-qat #189) #7). That fork change makes a 0-dbiasesan implied-bias factorf: the kernels formscales * T(f)in the scales' dtype, bit-identical to a materialized bias.tools/update-mlx.sh, run withLC_ALL=C._ibkernel variants tomlx-generated/metal/quantized{,_nax}.handmlx-generated/quantized{,_nax}.cpp.cuda_jit_sources.h, with identical content. That reorder was discarded, so the committed order is kept.QuantizedEmbedding.callAsFunctionpasses a 0-dbiasesstraight through todequantized, instead ofbiases[x](a 0-d array has no rows to gather).QuantizedLinearalready handsbiasesstraight toquantizedMM.Commits
implied_biastemplate arg only on the implied branch; without itaffine_sym_qmvfailed at pipeline creationstop_gradientinmaterialize_implied_bias;QuantizedMatmulvjp/jvp no longer throw) and condenses comments. Regenerates comment-only changes inmlx-generated/metal/quantized.handquantized.cpp, plus the mlx-ops.h docs and the Metal.hpp stamp. No forward-path change.Verification
The tests ran in numen-tech/gemma4-qat (
perf/190-implied-bias, pinned to 89b625c; re-checked at dd0cf30 where noted):ParoQuantMLXTests: 134 tests in 20 suites. This includes the newImpliedBiasKernelTests, which check bitwise equality against a materialized bias over the shipped bit widths and factors,QuantizedEmbeddingincluded.NumenCoreTests: 1549 executed, 0 failures.NumenDispatchTests: ×3, 22/22.