feat(quant): thread a mode argument through gather_qmm - #358
Open
sergey-scherbina wants to merge 1 commit into
Open
feat(quant): thread a mode argument through gather_qmm#358sergey-scherbina wants to merge 1 commit into
mode argument through gather_qmm#358sergey-scherbina wants to merge 1 commit into
Conversation
`gather_qmm` hardcoded `DEFAULT_MODE` ("affine"), so it could not gather over
weights stored in a microscaling format such as "mxfp4" (scales only, no
zero-point biases), even though the underlying `mlx_gather_qmm` C API and
`quantize`/`dequantize` already accept a `mode`. Add an optional `mode`
parameter (default "affine", so existing behavior is unchanged), mirroring the
existing `mode` plumbing on the nvfp4/mxfp8 quantization helpers.
Existing `test_gather_qmm` callers updated to pass the new optional argument.
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.
Summary
gather_qmmhardcodesDEFAULT_MODE("affine") when callingmlx_gather_qmm, so it cannot gather over weights stored in a microscalingformat such as
"mxfp4"(scales only, no zero-point biases) — even though theunderlying C API accepts a
mode, andquantize/dequantize(and thenvfp4/mxfp8 helpers) already expose one.
Change
Add an optional
mode: impl Into<Option<&str>>parameter togather_qmm,defaulting to
"affine"so existing behavior is unchanged. The plumbing mirrorsthe existing
modehandling on the nvfp4/mxfp8 quantization helpers (CStringmode_cstr.as_ptr()).Existing
test_gather_qmmcallers are updated to pass the new optionalargument; they continue to exercise the default affine path.