Repository navigation
Implied-bias affine quantized_matmul / dequantize (0-d biases factor) - #6
Conversation
…-tech/gemma4-qat#190) An affine checkpoint whose per-group bias is a fixed multiple of its scale (bias = f * scale; f = -2/-8/-128 for the Gemma 4 mobile-QAT 2/4/8-bit tensors at group size 64, -0.5 for Bonsai-27B 1-bit and -1 for Ternary Bonsai 2 2-bit at group size 128) can now pass the factor instead of a materialized bias array: a 0-d `biases` argument to `quantized_matmul` and `dequantize` in mode "affine" selects kernels that form the bias in place where the per-group bias used to be loaded. - ops: validation accepts a 0-d floating `biases` (dtype promotion from `scales`); `quantized_matmul` keeps the scalar out of `broadcast_arrays`; `dequantize` casts it to the scales dtype before the GPU kernel (which reads biases as the scales type) and multiplies it out in the CPU fallback; `gather_qmm` throws rather than broadcasting a constant bias. - Metal: every affine qmv/qvm/qmm/dequantize kernel (quad, fast, generic, wide, split-k, qmm_t/qmm_n and their NAX variants, the block loader) gets a trailing `implied_bias` template bool; the host appends `_ib` to the kernel name and the .metal files instantiate the 1/2/4/8-bit variants for the non-JIT library (3/5/6-bit implied kernels exist in JIT builds only). The bias is formed as `T(f) * T(scale)` and then widened, i.e. rounded through the scales dtype exactly as a materialized bias would be, so the output is bit-identical to the rebuilt-bias path even for fp16 subnormal scales where `f * scale` is not representable (Bonsai-27B-1bit carries a handful). - CPU: `QuantizedMatmul::eval_cpu` materializes `f * scales` as a temporary and dispatches the existing kernels. - vjp: scales gradient uses the factor as the derived-bias constant; a gradient wrt the factor itself is an error. jvp is unchanged. Tests: `tests/ops_tests.cpp` "test implied-bias affine quantized_matmul and dequantize" (all five factors, fp16/fp32, K in {128, 512, 1536}, M in {1,2,3,5,8,33,128}, transposed and not, CPU and GPU, bitwise equality against the materialized bias; unsupported paths throw) and `python/tests/test_quantized.py::test_implied_bias_*` (matmul, every GPU dispatch path incl. split-k and batched, dequantize incl. gathered rows, subnormal fp16 scales, grads, error cases).
…re the kernels cannot honor it (numen-tech/gemma4-qat#190) Codex review of the implied-bias op found three parity breaks and one dispatch gap, all from the bias being formed somewhere other than the scales' storage dtype: - `quantized_matmul` promoted the scales to the activations' dtype and the `_ib` kernel then formed the bias in the promoted dtype (fp32 x over fp16 scales holding 2^-24 with f = -0.5: materialized 0, implied -2^-25). - the CPU `dequantize` fallback multiplied `scales * f` in the factor's dtype (an fp32 factor promotes it), skipping the fp16 rounding. - an arbitrary floating factor was accepted but `T(f) * scale` is not `T(f * scale)` (fp16 scale 0.10009765625, f = 1.0004). - non-JIT builds (a prebuilt metallib, as mlx-swift ships) have no 3/5/6-bit `_ib` instantiations, so a 3-bit implied dequantize threw on kernel lookup. Contract (ops.h, the Python docstrings): a 0-d `biases` factor f stands for the materialized bias `scales * T(f)` formed in the scales' dtype T, f cast to T first; a factor not exactly representable in T is outside the contract (no runtime value check -- the op is lazy). One predicate, `implied_bias_kernels_exact`, says when the `_ib` kernels honor that exactly: the kernel sees the scales in T (no promotion) and bits is in {1, 2, 4, 8}. Otherwise `materialize_implied_bias` forms `scales * T(f)` in T up front and the ordinary per-group-bias path runs; the CPU dequantize fallback uses the same helper. The 3/5/6-bit kernels stay JIT-only so the metallib does not grow. Tests: `tests/ops_tests.cpp` "test implied-bias factor rounding contract" (the four repros on CPU and GPU, exact values asserted) and `python/tests/test_quantized.py::test_implied_bias_{dtype_promotion, cpu_dequantize_rounding,non_representable_factor,non_power_of_two_bits}`; the existing tests' reference is now spelled `scales * T(f)`.
The CUDA quantized kernels have no implied-bias variant: dequantize would index the 0-d factor as a per-group bias. implied_bias_kernels_exact now also requires a CPU stream or Metal, so CUDA takes the materialized-bias path (numen-tech/gemma4-qat#190).
- One `is_implied_bias` predicate (mlx/backend/common/quantized.h) shared by ops.cpp, primitives.cpp and the Metal/CPU backends; `implied_bias_suffix` derives the "_ib" kernel-name suffix and every launch site passes the flag computed once as the trailing template arg. - `affine_bias` is the single METAL_FUNC bias selector for the affine kernels (bias-free / implied / stored), reused by `QuantizedBlockLoader::load_bias`. Rounding order is unchanged: the implied bias is still `f * scale` in T. - `implied_bias_kernels_exact` is Metal-only; the CPU QuantizedMatmul rejects a 0-d bias instead of carrying its own dequant branch (quantized_matmul materializes `scales * T(f)` before the op reaches the CPU). - `gather_qmm` materializes the implied bias instead of throwing. - Tests: shared helpers in ops_tests.cpp and test_quantized.py; new gather_qmm bitwise-equality tests on CPU and GPU. - Reverted reformat-only hunks; kernel-name set of the built metallib is identical before/after.
Collapsing the two get_quantized_kernel_wrapped calls into one that always appends `implied` fed a seventh template argument to kernels that have no implied_bias parameter (affine_sym_*, fp_*). Non-JIT builds never saw it; the JIT build in mlx-swift failed to compile affine_sym_qmv at pipeline creation. Restore the per-branch dispatch.
Codex review of the cleanup pass: routing CPU streams through materialize_implied_bias changed behaviour the cleanup was meant to keep. - Gradients: the materialized bias is a differentiable multiply, so on CPU a gradient wrt the factor no longer threw and the scales gradient summed in a different order than on Metal. - Exported graphs: a CPU QuantizedMatmul with a 0-d bias exported before the cleanup re-imports the primitive directly, and eval_cpu now threw. - Validation: restricting the floating-dtype check to implied biases let int32 per-group biases through quantized_matmul/dequantize/gather_qmm. Restore QuantizedMatmul::eval_cpu's factor expansion and the CPU clause of implied_bias_kernels_exact, and check every biases dtype. Test: int32 per-group biases raise. DEVICE=cpu test_quantized -k implied: 11 passed.
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. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 499fc3e27e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Implied bias (a 0-d `biases` array holding a factor f): the contract is that | ||
| // it stands for the materialized per-group bias `scales * T(f)`, formed in the | ||
| // scales' storage dtype T -- f is cast to T first, so a factor that is not | ||
| // exactly representable in T is outside the contract (the reference is still | ||
| // `scales * T(f)`, not `T(f * scales)`). The `_ib` kernels form `T(f) * |
There was a problem hiding this comment.
Condense this 17-line block to the short reason for the predicate and keep the detailed rounding contract in the public documentation or tests. The block duplicates implementation details and exceeds the repository's explicit limit of concise, usually 1–2 line comments.
AGENTS.md reference: AGENTS.md:L31-L33
Useful? React with 👍 / 👎.
| auto b = (is_implied_bias(*biases) && | ||
| !implied_bias_kernels_exact(scales, dtype, bits, to_stream(s))) | ||
| ? materialize_implied_bias(scales, *biases, s) |
There was a problem hiding this comment.
When this fallback runs—for example with fp32 activations and fp16 scales, 3/5/6-bit weights, or CUDA—it replaces the scalar factor with a differentiable multiply(scales, factor). QuantizedMatmul::vjp then sees a full bias array and propagates a factor gradient through that multiply, whereas configurations using the scalar kernel throw no gradient wrt an implied-bias factor; the same API therefore has device-, dtype-, and bit-width-dependent autodiff behavior and can silently train a value intended to be constant. Preserve the scalar factor's no-gradient semantics when materializing it.
Useful? React with 👍 / 👎.
| // The gather kernels have no implied-bias variant: materialize the factor. | ||
| std::optional<array> biases = biases_; | ||
| if (is_implied_bias(biases_)) { | ||
| biases = materialize_implied_bias(scales, *biases_, s); |
There was a problem hiding this comment.
This adds implied-bias-factor support to gather_qmm, but its public Python documentation at python/src/ops.cpp:4845-4853 still says that biases must have the same batch dimensions as w and describes only per-group bias arrays. Unlike the updated quantized_matmul and dequantize docs, it therefore tells users that the newly accepted scalar form is invalid and leaves the memory-saving API undiscoverable; update the binding documentation and the C++ API comment with the scalar contract.
Useful? React with 👍 / 👎.
|
@codex review |
|
Codex Review: Didn't find any major issues. 🎉 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
For numen-tech/gemma4-qat#190 (epic numen-tech/gemma4-qat#178, Phase 2).
What
quantized_matmul,dequantizeandgather_qmmin mode"affine"now accept a 0-dbiasesarray: an implied-bias factorf. Some checkpoints store a per-group bias that is always a fixed multiple of the group's scale:For those, the factor can stand in for the whole bias array, so the bias never takes memory. On E2B/E4B that saves 131.6 / 211.4 MB resident.
Contract (
ops.h, Python docstrings): a factorfmeans the materialized biasscales * T(f), formed in the scales' dtypeTwithfcast toTfirst. A factor that is not exactly representable inTis outside the contract. The op is lazy, so this isn't checked at runtime.Dispatch.
implied_bias_kernels_exactdecides when the kernels can honor the contract bit-for-bit. All of these must hold:QuantizedMatmul::eval_cpuexpands the factor into a temporary) or a Metal GPU stream;T(no promotion from the activations);bits ∈ {1, 2, 4, 8}.When it holds on Metal, an
_ibkernel variant formsT(f) * T(scale)where the per-group bias used to be loaded. This is bit-identical to a materialized bias, including fp16 subnormal scales, which Bonsai-27B carries.In every other case,
materialize_implied_biasformsscales * T(f)inTup front and the ordinary path runs. That covers CUDA, promoted scales, 3/5/6-bit,gather_qmm, and CPUdequantize.Commits
0e7a5b0 Implied-bias affine qmv/qvm/qmm/dequantize kernels.
implied_biastemplate bool, and the kernel name takes an_ibsuffix.5a7aeac Pins the rounding contract after Codex round 1. It fixes three parity breaks, each from the bias being formed outside
T:It also fixes one dispatch gap: no 3/5/6-bit
_ibin a prebuilt metallib. All four now take the materialize path.6fa3f0c CUDA streams materialize (no
_ibCUDA kernels).477f91a
/simplifypass:is_implied_biaspredicate inbackend/common/quantized.h;affine_biasMETAL_FUNC as the single bias selector;gather_qmmmaterializes instead of throwing;The kernel-name set of the built metallib is identical before and after.
a84b8d0 Passes the
implied_biastemplate arg only on the implied branch.affine_sym_*andfp_*have no such parameter. The JIT build, which is what mlx-swift ships, failedaffine_sym_qmvpipeline creation without this.499fc3e Codex round 3 on the cleanup. It undoes two cleanup changes that weren't behaviour-preserving:
QuantizedMatmulholds a 0-d bias. The CPU factor expansion ineval_cpuis restored.Verification
C++
testsat 499fc3e:MLX_METAL_JIT=ON:-tc="*quantiz*"3/3 (609 assertions).New cases:
Python:
test_quantized+test_fast, 85 passed and 1 skipped. WithDEVICE=cpu,test_quantized -k impliedpasses 11/11. The other CPU failures are the existing Metal-only bias-free affine path from the Thank you very much for using CMAKE! ml-explore/mlx#86 work. The newtest_implied_bias_*cover:gather_qmm;Downstream (gemma4-qat):
See the gemma4-qat PR.
Metallib size: unchanged (2,401,016 B). The JIT build keeps quantized kernels out of the prebuilt library.
Review
Codex adversarial review ran three rounds:
Then a
/simplifypass.Downstream: numen-tech/mlx-swift
prism-190-implied-biasregenerates from this branch.