Pass the floating point activation to et_vk.linear_q8ta_q8csw - #22776
Pass the floating point activation to et_vk.linear_q8ta_q8csw#22776giuliocorradi wants to merge 1 commit into
Conversation
The Vulkan implementation of linear_q8ta_q8csw quantizes the activation tensor itself: QuantizedLinear.cpp names the argument fp_input, allocates a temporary packed int8 tensor for it, and passes it to add_quantize_and_pack_4h4w_node. The AOT fusion was passing the quantized tensor instead. In QuantizedLinearMatch, pattern_input_node is set to the int8 output of the input quantize node, and only the dynamic quantization branch steps back past it to the floating point tensor. Statically quantized inputs therefore reached the op as int8, and the delegate failed at the first inference looking for a shader variant that takes an already quantized input: Could not find ShaderInfo with name clone_buffer_to_image_int8_int32 Add QuantizedLinearMatch.get_fp_input_node() to resolve the floating point activation, and use it when building the op. Where the activation is only available as a quantized tensor - it is a graph input, or it is produced by a preceding op that was already replaced with q8ta_linear - insert a dequantize node so the op still receives what it expects. q8ta_linear is unaffected: it takes packed_int8_input and continues to receive pattern_input_node. Note that no bundled quantizer produces this op today, which is why the mismatch went unnoticed. XNNPACKQuantizer also quantizes the linear's output, which selects q8ta_linear, and VulkanQuantizer has no static activation quantization mode. The added test therefore builds the quantize/dequantize pattern directly. Test Plan: python -m unittest backends.vulkan.test.test_vulkan_passes -v 12 tests pass. The new test, test_linear_q8ta_q8csw_takes_floating_point_input, fails without this change with "linear_q8ta_q8csw was given the quantized activation; it expects the floating point one". Also verified end to end on an AMD Radeon 8060S (RADV GFX1151), running the lowered .pte and comparing against the eager module: case before after single static-quant linear shader error max|d| 2.3e-05 2 chained quantized linears shader error max|d| 1.2e-04 3 chained quantized linears shader error max|d| 1.1e-04 int8 activation as a graph input shader error max|d| 1.8e-05 backends.vulkan.test.test_vulkan_delegate shows an identical set of results before and after the change.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22776
Note: Links to docs will display an error until the docs builds have been completed.
|
|
Hi @giuliocorradi! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
This PR needs a
|
There was a problem hiding this comment.
🟡 Changes recommended
The new fallback path in make_linear_q8ta_q8csw_custom_op hard-codes a dequantize op (qmin/qmax/dtype) instead of reusing the existing dequantize node, which can diverge from the graph’s original dequantization parameters.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR fixes Vulkan AOT fusion for et_vk.linear_q8ta_q8csw so the custom op receives the floating-point activation (which the runtime quantizes internally) instead of incorrectly receiving the already-quantized int8 tensor in the statically-quantized activation case.
Changes:
- Add
QuantizedLinearMatch.get_fp_input_node()to resolve the correct floating-point activation input forlinear_q8ta_q8csw. - Update
make_linear_q8ta_q8csw_custom_opto use the resolved fp activation when constructing the op. - Add a unit test that asserts
linear_q8ta_q8cswis not fed by an activation quantize node (and is not int8-typed).
File summaries
| File | Description |
|---|---|
| backends/vulkan/patterns/quantized_linear.py | Resolve and pass the floating-point activation to et_vk.linear_q8ta_q8csw during pattern replacement. |
| backends/vulkan/test/test_vulkan_passes.py | Add regression test ensuring linear_q8ta_q8csw takes a floating-point activation input. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # This op quantizes the activation tensor itself, so it takes the floating | ||
| # point input rather than the quantized one. | ||
| fp_input_node = match.get_fp_input_node() | ||
| if fp_input_node is None: | ||
| # The input is only available as a quantized tensor, i.e. it is produced | ||
| # by a preceding quantized op. Dequantize it so that the op has the | ||
| # floating point tensor it expects. | ||
| with graph_module.graph.inserting_before(match.output_node): | ||
| fp_input_node = graph_module.graph.create_node( | ||
| "call_function", | ||
| exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, | ||
| args=( | ||
| match.pattern_input_node, | ||
| match.input_scales_node, | ||
| match.input_zeros_node, | ||
| -128, | ||
| 127, | ||
| torch.int8, | ||
| ), | ||
| ) | ||
| fp_input_node.meta["val"] = match.dequantize_input_node.meta["val"] | ||
|
|
![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: 793fad1bf4
ℹ️ 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".
| if utils.is_quant_node(self.pattern_input_node): | ||
| return self.pattern_input_node.args[0] # pyre-ignore[7] |
There was a problem hiding this comment.
When the matched static quantizer uses a valid non-signed-int8 configuration, such as torch.uint8 with range [0, 255], this unconditionally bypasses the quantize node without inspecting its dtype or bounds. The fused op then requantizes the float input using quantize_and_pack in linear_int8_input_block.glslh, which hard-codes [-128, 127]; the fallback below likewise forces those bounds and torch.int8. Because is_input_static_per_tensor_quantized() checks only that the scale is a float, these patterns are currently accepted and can produce severely incorrect results rather than preserving the original quantize/dequantize semantics. Validate that the matched q/dq pair uses the supported [-128, 127] int8 scheme, or skip the fusion.
Useful? React with 👍 / 👎.
Fixes #22775
Summary
The Vulkan implementation of
linear_q8ta_q8cswquantizes the activation tensor itself.QuantizedLinear.cppnames the argumentfp_input, allocates a temporary packed int8 tensor for it, and hands it toadd_quantize_and_pack_4h4w_node:The AOT fusion was passing the quantized tensor instead. In
QuantizedLinearMatch.__init__,pattern_input_nodeis set to the int8 output of the input quantize node, and only the dynamic-quantization branch steps back past it to the float:So a statically quantized activation reached the op as int8. The program lowered and serialized fine, then failed at the first inference asking for a shader variant whose name carries an input dtype that cannot exist:
The change
Add
QuantizedLinearMatch.get_fp_input_node()to resolve the floating point activation, and use it when building the op.Where the activation is genuinely only available quantized, there is nothing to step back to, and the op still needs a float — this happens when the int8 tensor is a graph input, and when it is produced by a preceding linear that pattern replacement already turned into
q8ta_linear. In that case adequantize_per_tensornode is inserted so the op receives what it expects. That case aborts today, so the extra dequantize is strictly an improvement; a follow-up could avoid the round trip with an int8-input variant of the op.q8ta_linearis deliberately untouched: it takespacked_int8_inputand keeps receivingpattern_input_node.Why this was not caught
No bundled quantizer produces this op.
XNNPACKQuantizeralso quantizes the linear's output, which routes toq8ta_linear, andVulkanQuantizer.get_symmetric_quantization_confighas only weight-only (is_dynamic=False→act_quantization_spec = None) and dynamic modes — no static per-tensor activation mode, which is the one thingis_input_static_per_tensor_quantized()matches on. The added test therefore builds the quantize/dequantize pattern directly.Test plan
12 tests pass. The new test,
test_linear_q8ta_q8csw_takes_floating_point_input, fails without the source change:Also verified end to end on an AMD Radeon 8060S (RADV GFX1151, RDNA 3.5), running the lowered
.ptethrough_load_for_executorchand comparing against the eager module:backends.vulkan.test.test_vulkan_delegateproduces an identical set of results before and after (the failures in that suite in my environment are pre-existing and unrelated).Why it matters
On RDNA 3.5 this op is the path to
v_wmma_i32_16x16x16_iu8, the INT8 matrix instruction. On this part llama.cpp's Vulkan backend reaches 13.71 TFLOPS with int8 weights on the same shape where this backend reaches 6.5 in fp32, while the fp32 paths are close (6.50 vs 7.40) — so the quantized path is where the remaining performance is.Found while lowering openpi's π₀.₅ to the Vulkan delegate.