[msluszniak]

Summary

A torchao int8 per-channel weight-only linear (IntxWeightOnlyConfig(int8, PerAxis(0))) matched no route in quantized_linear.py, so its torchao.dequantize_affine was serialized into the delegate, which has no kernel for it, and the model aborted at load. Route it to the existing linear_q8csw kernel, which had no AOT definition. Per-group int8 and 8da8w still have no route.

Test plan

New test_fuse_torchao_int8_weight_only_linear fails before, passes after; test_vulkan_passes.py passes.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin @nil-is-all

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24082

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 1 Pending

As of commit 8df4f7a with merge base 7645c90 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[SS-JIA]

Requesting changes based on an independent, AI-assisted static review (Claude) of the current head, which I've gone through and agree with. Each finding is posted as an inline comment with a suggested fix; where a finding is about code outside this PR's diff, the comment sits on the closest related line and its Where line gives the actual location.

Summary

• What it does: adds an AOT definition for the existing runtime op et_vk.linear_q8csw and a new replace_quantized_linear_patterns branch. torchao int8 per-channel weight-only linears (dequantize_affine with [N, 1] scales) now fuse into that op instead of leaving torchao.dequantize_affine, which has no kernel, in the delegate. That case works: it delegates as a single et_vk.linear_q8csw for fp32, fp16, and dynamic shapes, and the pass test passes.
• Main problem: the branch does not check the dequant op or the bit width. It also matches quantized_decomposed.dequantize_per_channel ([N] scales) weight-only linears. Previously these fell through to FuseQuantizedOpsTransform, which runs after FusePatternsPass. As a result:
  • PT2E 8-bit weight-only now lowers to linear_q8csw, not aten._weight_int8pack_mm plus add. The runtime implementation is the same, so this is benign but unannounced.
  • PT2E 4-bit weight-only now lowers to linear_q8csw, not linear_qcs4w. This is a regression (A1).
• Merges into current main with no conflicts. Merged #24055 and #24116 are already in the PR base, and no semantic overlap was found with them.
• Semantic conflict with open #24971. That PR gates the detector on has_fused_op(), which returns False for torchao int8 per-channel, and it asserts that this case stays undelegated. Whichever lands second must reconcile both the code and the tests (A2). This PR is right for its case: after #24971 alone, a torchao int8 weight-only linear fails at export with "Missing out variants". #24971's gating is right for the cases that still have no route.
• Still broken (acknowledged in the PR summary, not a regression): int8 per-group, 8da8w, K % 4 != 0, and non-zero weight zero points all still delegate torchao.dequantize_affine without a kernel. I verified the int8 PerGroup(32) and K=66 cases.
• Open #22776 conflicts textually in the same region (method added after is_weight_perchannel_quantized). It is already stale against main, and there is no semantic interaction.

Findings not tied to a line

A6 (Nit): PR description understates scope

• Where: PR summary
• Issue: The summary describes only the torchao case. PT2E per-channel 8-bit weight-only linears also move from aten._weight_int8pack_mm plus a separate aten.add to linear_q8csw with fused bias. That leaves the 8-bit branch of FuseQuantizedOpsTransform reachable only for linears the pattern rejects (e.g. K % 4 != 0).
• Suggested fix: Mention the PT2E routing change in the summary and test plan.

[SS-JIA]

A1 (High): PT2E 4-bit per-channel weight-only linear rerouted from linear_qcs4w to linear_q8csw; breaks WebGPU

• Where: backends/vulkan/patterns/quantized_linear.py:725-731

• Issue: The new branch accepts any weight-only, per-output-channel linear with an int8-typed weight and zero zero points. PT2E 4-bit per-channel quantization (VulkanQuantizer with get_symmetric_quantization_config(is_dynamic=False, weight_bits=4)) stores 4-bit values in an int8 tensor and uses [N] scales. The q4gsw branch rejects it because is_weight_pergroup_quantized requires 2-D scales, so it reaches the new branch. FusePatternsPass runs before FuseQuantizedOpsTransform in vulkan_preprocess.py, so the linear is fused into linear_q8csw before the existing linear_qcs4w route can see it.

  • Vulkan: weights are no longer 4-bit packed, so weight memory doubles and the op uses a different kernel.
  • WebGPU: it reuses the Vulkan patterns and does not skip quantized_linear. The WebGPU runtime implements linear_qcs4w but not linear_q8csw, so these models fail at delegate init.

• Evidence: Executed. Ad-hoc lowering through VulkanPartitioner, nn.Linear(64, 128), delegate ops:

  • origin/main: ['et_vk.linear_qcs4w.default', 'et_vk.prepack.default', 'aten.add.Tensor']
  • PR head: ['et_vk.linear_q8csw.default']

  The existing WebGPU test backends/webgpu/test/ops/test_linear_qcs4w.py::LinearQcs4wTest::test_delegates fails on the PR head: AssertionError: b'linear_qcs4w' not found in .... The serialized buffer contains et_vk.linear_q8csw.default. This test is not in the PR's test plan.

• Suggested fix: Exclude 4-bit-range weights so they keep the FuseQuantizedOpsTransform / linear_qcs4w route. torchao [N, 1] 4-bit weights are already taken by the q4gsw branch above.

          and weight_tensor.dtype == torch.int8
          and not utils.is_in_4bit_range(weight_tensor)
          and bool(torch.all(weight_zeros_tensor == 0))

  Verified: with this line, the PT2E 4-bit case lowers to linear_qcs4w again, the torchao int8 cases still fuse to linear_q8csw, and test_linear_qcs4w.py plus the -k "qcs4w or delegates or linear" subset of test_vulkan_passes.py pass (9 passed). An alternative is to restrict the branch to match.dequantize_weight_node.target == exir_ops.edge.torchao.dequantize_affine.default, which also leaves PT2E 8-bit on _weight_int8pack_mm.

[SS-JIA]

A2 (Medium): Semantic conflict with open #24971 ("Do not delegate quantized linears that have no Vulkan kernel")

• Where: backends/vulkan/patterns/quantized_linear.py (detector find_quantized_linear_patterns vs. replacement routing); backends/vulkan/test/test_vulkan_passes.py:1338

• Issue: Do not delegate quantized linears that have no Vulkan kernel (#24971) #24971 adds has_fused_op() to the detector. For a torchao weight-only linear that is not 4-bit per-group, it returns False: is_weight_4bit_quantized() is False for the int8 range, and dequantize_affine is not a dequantize_per_channel node. Do not delegate quantized linears that have no Vulkan kernel (#24971) #24971 also asserts matches(w8, 8) == [False] and that dequantize_affine stays outside the delegate. If both land as written:

  • this PR's branch becomes dead code and its test fails (result.modified / linear_q8csw count), or
  • Do not delegate quantized linears that have no Vulkan kernel (#24971) #24971's tests fail, depending on how the conflict is resolved.

  Both PRs also edit the same lines, which caused textual conflicts in quantized_linear.py and test_vulkan_passes.py when I applied Do not delegate quantized linears that have no Vulkan kernel (#24971) #24971's Vulkan hunks onto this head with git apply -3.

• Evidence: Static reasoning from gh pr diff 24971. The textual conflicts were observed with git apply --check -3.

• Suggested fix: Keep both behaviors. Do not delegate quantized linears that have no Vulkan kernel (#24971) #24971's has_fused_op() should return True for torchao weight-only per-output-channel 8-bit, and its w8 assertions should flip to delegated linear_q8csw. Better still, have the detector and the replacement share one routing predicate so they cannot diverge again. That divergence is the root cause of the original bug: the partitioner accepts every node of a match unconditionally (vulkan_partitioner.py:299). Coordinate the landing order with Do not delegate quantized linears that have no Vulkan kernel (#24971) #24971's author.

[SS-JIA]

A3 (Medium): Tests miss the routing side effects and the end-to-end lowering

• Where: backends/vulkan/test/test_vulkan_passes.py:1338-1393
• Issue: The new test runs only FusePatternsPass on one torchao config. It does not cover:
  • PT2E per-channel 8-bit and 4-bit, whose routing this PR changes (A1 would have been caught);
  • partitioner-level lowering (no remaining dequantize_affine in the delegate, delegate op list);
  • the no-bias / mm anchor variant.
• Evidence: A1 passes the PR's test plan but fails an existing WebGPU test.
• Suggested fix: Add a to_edge_transform_and_lower(..., partitioner=[VulkanPartitioner()]) test asserting that the delegate contains linear_q8csw and no dequantize_affine. Add a PT2E weight_bits=4 per-channel case asserting linear_qcs4w (or rely on the WebGPU test in CI). Add a bias=False case.

[SS-JIA]

A4 (Low): Near-duplicate predicate and replacement body

• Where: backends/vulkan/patterns/quantized_linear.py:270-276, 368-405
• Issue: is_weight_per_output_channel_quantized duplicates is_weight_perchannel_quantized apart from accepting [N, 1]. make_linear_q8csw_op repeats the bias-padding and create/replace block from make_linear_q4gsw_op almost line for line. Extending is_weight_perchannel_quantized directly is unsafe, because the q8ta_q8csw route would then pad [N, 1] scales along the wrong dimension, so a separate check is justified.
• Suggested fix: Give the existing predicate an allow_keepdim: bool = False parameter instead of adding a second method. Optionally factor the shared "pad bias, insert node before output_node, copy meta['val'], replace uses" logic into a helper used by the weight-only makers.

[SS-JIA]

A5 (Low): WebGPU still has no kernel for the newly emitted op

• Where: backends/webgpu/partitioner/partitioner.py:83 (_WEBGPU_SKIP_FUSION_PATTERNS), backends/webgpu/runtime/ops/
• Issue: With this PR, WebGPU lowers torchao int8 per-channel weight-only linears to et_vk.linear_q8csw. The WebGPU runtime does not implement that op; it registers only et_vk.linear_q4gsw among the weight-only linears. This is not a regression: before the PR, the same models delegated dequantize_affine, which WebGPU also lacks. The failure mode stays "delegate init failure" instead of falling back to CPU.
• Suggested fix: Out of scope for this PR. Note it for the WebGPU owners, or fold it into the shared-predicate work in A2 with a backend-specific route filter.