[ET-VK] Static activation quantization passes the int8 tensor to linear_q8ta_q8csw, which expects the float #22775

🐛 Describe the bug

Summary

The AOT fusion that builds et_vk.linear_q8ta_q8csw passes the quantised
activation when the input is statically quantised, while the runtime expects
the float activation and quantises it itself. Every statically quantised
linear therefore lowers and serialises cleanly and then fails at the first
inference, asking for a shader whose name contains an input dtype that cannot
exist:

get_shader_info at backends/vulkan/runtime/api/ShaderRegistry.cpp:51:
  (it != listings_.end()) is false!
Could not find ShaderInfo with name clone_buffer_to_image_int8_int32

The dynamically quantised path is unaffected, because it already steps back to
the float.

Where the two sides disagree

AOT — backends/vulkan/patterns/quantized_linear.py, QuantizedLinearMatch.__init__:

if utils.is_dequant_node(anchor_primary_input_node):
    # Assume that this is a static quantization pattern; the input to the
    # pattern is a statically quantized int8 tensor.
    self.dequantize_input_node = anchor_primary_input_node
    input_to_dq_node = self.dequantize_input_node.args[0]
    self.pattern_input_node = input_to_dq_node          # <- the int8 quantize output
    ...
    if utils.is_quant_node(input_to_dq_node) and utils.is_dynamic_qscale(
        self.input_scales_node
    ):
        self.quantize_input_node = input_to_dq_node
        self.pattern_input_node = self.quantize_input_node.args[0]   # <- the float

pattern_input_node becomes argument 0 of the op in
make_linear_q8ta_q8csw_custom_op. Only the dynamic branch steps back past the
quantize node; the static path stops at the int8 tensor, exactly as its comment
says it intends to.

Runtime — backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp
names that same argument fp_input and feeds it to the quantiser:

    const ValueRef fp_input,
    ...
    // Allocate temporary tensor to store quantized and packed input
    TmpTensor packed_int_input(
        &graph, graph.sizes_of(fp_input), vkapi::kInt8x4,
        utils::kBuffer, utils::kPackedInt8_4H4W);

    if (!input_quant_config.is_dynamic) {
      add_quantize_and_pack_4h4w_node(
          graph, input_quant_config, fp_input, ...);

So the runtime quantises the activation itself, and the AOT hands it something
already quantised. The delegate then tries to move that int8 buffer into a
texture and looks for clone_buffer_to_image_int8_int32.

The partitioner log states the mismatch plainly — argument 0 is the
quantize_per_tensor node:

[Vulkan Delegate] Inserting transition(s) for %et_vk_linear_q8ta_q8csw_default :
  args = (%quantized_decomposed_quantize_per_tensor_default, 0.05, 0, %b_w_q, ...)
  arg 0 (quantized_decomposed_quantize_per_tensor_default):
    (TensorRepr(BUFFER, PACKED_INT8_4W4C)) -> (TensorRepr(TEXTURE_3D, TENSOR_WIDTH_PACKED))

Reproduction

No model weights needed. This lowers, serialises, and then fails on forward.

import torch
from executorch.exir import to_edge
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.extension.pybindings.portable_lib import _load_for_executorch

qd = torch.ops.quantized_decomposed
K, N, M = 256, 256, 8
torch.manual_seed(0)
W = torch.randn(N, K)
S_W = W.abs().amax(dim=1).clamp(min=1e-8) / 127.0
Z_W = torch.zeros(N, dtype=torch.int64)
W_Q = qd.quantize_per_channel.default(W, S_W, Z_W, 0, -127, 127, torch.int8)
ACT_S = 0.05

class Mod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("w_q", W_Q)
        self.register_buffer("s_w", S_W)
        self.register_buffer("z_w", Z_W)
        self.bias = torch.nn.Parameter(torch.randn(N))
    def forward(self, x):
        xq  = qd.quantize_per_tensor.default(x, ACT_S, 0, -128, 127, torch.int8)
        xdq = qd.dequantize_per_tensor.default(xq, ACT_S, 0, -128, 127, torch.int8)
        wdq = qd.dequantize_per_channel.default(
            self.w_q, self.s_w, self.z_w, 0, -127, 127, torch.int8)
        return torch.nn.functional.linear(xdq, wdq, self.bias)

mod, x = Mod().eval(), torch.randn(M, K)
low = to_edge(torch.export.export(mod, (x,))).to_backend(VulkanPartitioner())
with open("/tmp/q8ta.pte", "wb") as f:
    low.to_executorch().write_to_file(f)
out = _load_for_executorch("/tmp/q8ta.pte").forward((x,))[0]   # <- raises here
print("max|d| =", (out - mod(x)).abs().max().item())

Observed:

RuntimeError: Exception raised from get_shader_info at
  .../backends/vulkan/runtime/api/ShaderRegistry.cpp:51:
  (it != listings_.end()) is false!
  Could not find ShaderInfo with name clone_buffer_to_image_int8_int32

Suggested fix, and confirmation that it works

Step pattern_input_node back past the quantize node in the static case too,
as the dynamic case already does — the scales and zero points still come from
the dequantize node, so nothing else changes.

Monkey-patching exactly that, on top of the unmodified build, makes the
reproduction above run and produce the right answer:

from executorch.backends.vulkan.patterns import quantized_linear as ql
from executorch.backends.vulkan import utils as vk_utils

_orig = ql.QuantizedLinearMatch.__init__

def _patched(self, *a, **kw):
    _orig(self, *a, **kw)
    dq = getattr(self, "dequantize_input_node", None)
    if dq is not None and self.quantize_input_node is None:
        producer = dq.args[0]
        if isinstance(producer, torch.fx.Node) and vk_utils.is_quant_node(producer):
            self.quantize_input_node = producer
            self.pattern_input_node = producer.args[0]

ql.QuantizedLinearMatch.__init__ = _patched

(The eager reference simulates the same quantisation, so that residual is the
shader's own accumulation error, not quantisation error.)

Secondary note: the reference implementation ignores the input scale

backends/vulkan/custom_ops_lib.py's CompositeExplicitAutograd implementation
takes input_scale and input_zero_point and never uses them:

def linear_q8ta_q8csw(x, input_scale, input_zero_point, weights,
                      weight_sums, weight_scales, bias=None):
    weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32)
    weights = torch.ops.quantized_decomposed.dequantize_per_channel(
        weights, weight_scales, weight_zeros, 0, -127, 127, torch.int8)
    out = torch.nn.functional.linear(x, weights)
    ...

Given x is meant to be the float activation, the dtypes are consistent — but
the op models no activation quantisation at all, so it cannot be used as a
numerical reference for the shader, and calling it with the int8 tensor the AOT
currently passes fails outright:

RuntimeError: expected m1 and m2 to have the same dtype, but got: signed char != float

Quantising x with the scale and zero point it already receives would make it a
usable reference and would have surfaced the argument mismatch above.

Impact

This blocks et_vk.linear_q8ta_q8csw for statically quantised activations,
which on RDNA 3.5 is the v_wmma_i32_16x16x16_iu8 path — the INT8 matrix
instruction. On the part below, that is worth roughly 2.1× on the dominant
operator: llama.cpp's Vulkan backend reaches 13.71 TFLOPS with int8 weights
on the same GPU and shape where this backend reaches 6.5 in fp32. The fp32
gap is small by comparison — 6.50 against llama.cpp's 7.40 — so the quantised
path is where the remaining Vulkan performance on this class of hardware lives.

Possibly related: #22431, on the accuracy of the dynamic activation
quantisation spec.

Versions

Collecting environment information...
PyTorch version: 2.12.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCm SDK used to build PyTorch: N/A
HIP used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.39

Python version: 3.12.3 (main, Jul 15 2026, 23:46:41) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-7.0.0-30-generic-x86_64-with-glibc2.39
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A
CPU: AMD RYZEN AI MAX+ 395 w/ Radeon 8060S, 16 cores / 32 threads (full lscpu elided)

Versions of relevant libraries:
[pip3] executorch==1.4.0a0+b20f16a
[pip3] numpy==2.4.6
[pip3] pytorch_tokenizers==1.4.1
[pip3] torch==2.12.1+cpu
[pip3] torchao==0.18.0.dev20260715+cpu
[conda] Could not collect