[Metal] Add additional template type S for scales and biases - #4473
Closed
chillosu wants to merge 1 commit into
Closed
[Metal] Add additional template type S for scales and biases#4473chillosu wants to merge 1 commit into
chillosu wants to merge 1 commit into
Conversation
chillosu
force-pushed
the
up/mixed-qscales
branch
2 times, most recently
from
September 8, 2026 04:25
dd55177 to
1f782ee
Compare
…atvec kernels Affine quantized_matmul and gather_qmm used one template type T for the activations, the scales, the biases and the output. With bfloat16 activations and float16 scales (what an FP8-sourced conversion that dequantizes to float16 produces) no kernel existed, so ops.cpp promoted to float32 and astype()'d the whole scales and biases tables on every call, and the float32 result propagated through the rest of the model. The matvec kernels (qmv, qmv_fast, gather_qmv, gather_qmv_fast) get a second template type S for scales and biases, defaulting to T so existing instantiations are unchanged, and convert the scale in-register per group where it was already going to the float accumulator. Both 16-bit pairs are instantiated across all bit widths and group sizes; the dispatch tags the kernel name with the scales type in affine mode only. ops.cpp no longer promotes a 16-bit/16-bit mismatch on a Metal stream; every other Metal path (qmm, qvm, split-k, quad/wide, sorted gather) runs the float32 computation with the casts as temporaries inside the primitive and casts the result back, so the output dtype is the activation dtype on every path. CPU and other backends keep the promotion. The behaviour change: a mismatched call now returns the activation dtype instead of float32. Measured on M3 Ultra (256 experts, 3072x1024, 8 distinct experts per call): mismatched gather_qmm 0.465 ms/call as float32 before, 0.245-0.269 ms as bfloat16 after, against 0.254-0.274 for a consistent checkpoint; qmv 3072->8192 0.0022 -> 0.0017 ms. With native float16 scales the result is within bfloat16 output rounding of the float32 reference. Adds test_mixed_scale_dtype. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BK9PxwPiviSTjacBvknBzs
chillosu
force-pushed
the
up/mixed-qscales
branch
from
September 8, 2026 23:12
1f782ee to
66f5a9c
Compare
zcbenz
requested changes
Sep 9, 2026
zcbenz
left a comment
zcbenz
left a comment
Member
There was a problem hiding this comment.
The kernels do require activations and scales to be the same type.
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
2 participants
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.
The c++ template assumes a single typed parameter, T, across activations, scales, biases, and the output. If activations are bfloat16 and scales are float16, no kernel exists. The common type conversion penalty is especially expensive for scales and biases which could convert on every call, which comes back and poisons the rest of the network. Proposing giving scales and biases their own type parameter, S, with default to T.
all 39 quantized tests pass, test_ops passes. Mismatched gather_qmm at decode: 0.245 to 0.269 ms against 0.254 to 0.274 for a consistent checkpoint, down from 0.465 as float32. Prefill-shaped fallback at M=512: 1.92 ms against 1.47 consistent, the one-time cast now amortized over 512 rows. CPU stream: float32 output, exact, unchanged.