Cortex-M: fold constant attention-score scale into the softmax-input quantize#21027
Cortex-M: fold constant attention-score scale into the softmax-input quantize#21027rascani wants to merge 5 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21027
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ No FailuresAs of commit 7bde4b7 with merge base 5ed466e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
AdrianLundell
left a comment
AdrianLundell
left a comment
There was a problem hiding this comment.
What happens here if the user node here is a SharedQspec operator which does not do any rescaling, e.g. a view? I think it would be better to do this after qparam-folding and check that either the input or output is an operator which is known to do rescaling and has quantization parameters and fold it into those.
For un-symmetric qspecs it should be possible to fold constant add/sub ops into the zp as well.
…x-input quantize A constant elementwise scale between a quantized producer and a per-tensor quantize -- e.g. the attention scores/sqrt(d) before softmax -- otherwise stays an fp32 aten.div/mul with dq/q boundaries. Since quantize(x/c, S) == quantize(x, S*c) (and *c divides S), the constant folds into the adjacent quantize with no numerical change. Add a cortex_m-scoped FoldScaleIntoQuantizePass (run before FoldAndAnnotateQParamsPass, while the softmax-input quantize is still explicit) that folds a constant scalar div/mul whose sole consumer is a quantize_per_tensor into that quantize's scale and drops the div/mul. Only constant (param, numel==1) scales fold, so data-dependent normalization (e.g. LayerNorm's divide by sqrt(var)) is correctly left untouched. Verified test-first (RED->GREEN) with div and mul scaled-attention cases (numerics at qtol=1 pin the fold direction); full cortex_m op dialect suite still passes (179). On a SAM mask decoder this removes the 7 attention /sqrt(d) divs so the attention chain is int8 through softmax. Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This branch adds backends/cortex_m/passes/fold_scale_into_quantize_pass.py and imports it from cortex_m_pass_manager. Add the file to the cortex_passes python_library srcs so the internal buck build can resolve the import; OSS pytest imports by path and does not catch a missing buck src. The pass's imports (arm _passes, exir, exir.dialects._ops, exir.pass_base, torch.fx) are already covered by the target's existing deps and sibling srcs, so no new dep is needed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Address AdrianLundell's review of pytorch#21027: run FoldScaleIntoQuantizePass after FoldAndAnnotateQParamsPass and fold the constant scale into a rescaling op that owns quantization parameters, rather than into an explicit quantize node before annotation. Post-annotation the constant scale (e.g. attention /sqrt(d)) is already absorbed into the consuming op's input scale: the quantizer's observer saw the post-scale range, so the softmax-input scale equals the bmm-output scale divided by the constant (multiplied, for mul) with matching zero-points. The div/mul therefore survives inside a redundant dequantize -> op -> quantize sandwich that is the identity on the int8 values. The pass recognises that relation -- checking the scale, zero-point, and clamp range/dtype all line up -- and deletes the sandwich, wiring the producer's int8 straight into the consumer. It rewrites no quantization parameter, so a SharedQspec consumer that shares the scale (a view/reshape/pool) is never disturbed, the case Adrian raised. If the relation does not hold the fold is skipped. Tests cover the div and mul folds, bit-exactness over the int8 range, a constant scale feeding a SharedQspec pool (shared scale unchanged), a constant-first operand (must not crash), and the skip guards. Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The div/mul cases (and test_attention_scale.py) use a bare rank-3 torch.bmm, whose output the quantizer quantizes directly, so they do not exercise the matmul->bmm dependency the fold actually relies on. Add a rank-4 multi-head q @ k^T attention like SAM's mask decoder: MatmulToBmmPass rewrites the matmul to a quantizable bmm at annotation, so the /sqrt(head_dim) scale lands in the requantize sandwich the fold removes. The test fails if matmul->bmm regresses (the scale would stay fp32 and never fold). Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
e5cfa0c to
d90e409
Compare
Add coverage for the two skip guards that had none -- a dtype/clamp-range mismatch (args 3..5, the int16-activation path) and a non-param runtime divisor; both are load-bearing (without them the pass would wrongly fold, or crash reading a param tensor for a non-param node). Remove the dead ArmPass _passes_required_after boilerplate (CortexMPassManager never reads it) and its now-unused typing imports. Correct the add/sub docstring: an additive shift moves the affine qparams by a generally non-integer amount, so no integer-preserving requantize identity exists. Reword two test comments that named a reviewer. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
I've updated to move it after qparam-folding, but I'm not particularly happy with the result. I will probably keep iterating on this one. |
2 participants
Summary
A constant elementwise scale between a quantized producer and a per-tensor quantize -- e.g. the attention scores/sqrt(d) before softmax -- otherwise stays an fp32 aten.div/mul with dq/q boundaries. Since quantize(x/c, S) == quantize(x, S*c) (and *c divides S), the constant folds into the adjacent quantize with no numerical change.
Add a cortex_m-scoped FoldScaleIntoQuantizePass (run before FoldAndAnnotateQParamsPass, while the softmax-input quantize is still explicit) that folds a constant scalar div/mul whose sole consumer is a quantize_per_tensor into that quantize's scale and drops the div/mul. Only constant (param, numel==1) scales fold, so data-dependent normalization (e.g. LayerNorm's divide by sqrt(var)) is correctly left untouched.
Test plan
Verified test-first (RED->GREEN) with div and mul scaled-attention cases (numerics at qtol=1 pin the fold direction); full cortex_m op dialect suite still passes (179). On a SAM mask decoder this removes the 7 attention /sqrt(d) divs so the attention chain is int8 through softmax.
Authored with Claude Code.