X86InductorQuantizer + torch.export/aoti_compile_and_package never fuses quantized Linear into onednn::qlinear_pointwise (silent fp32 fallback)

[AlexPoone]

Description

Summary

Quantizing a model with X86InductorQuantizer (static PT2E quantization) and then
compiling it via torch.export.export() + torch._inductor.aoti_compile_and_package()
(the AOTInductor deployment path, as opposed to torch.compile()) produces a package
where every Linear/GEMM op still executes in fp32. Weights are stored compactly as
int8, and quantize_per_tensor/dequantize_per_channel/dequantize_per_tensor nodes
are correctly inserted by convert_pt2e(), but the weight-prepack fusion that should
convert dequantize -> linear -> quantize into a fused onednn::qlinear_pointwise
(and lower to a VNNI/AVX-512 int8 GEMM) never fires. Every GEMM in the compiled
.pt2 package instead dequantizes to fp32, computes with the plain
CppMicroGemmFP32Vec template, and requantizes the output afterward -- i.e. pure
quantize/dequantize overhead with zero compute benefit, and plausibly a net
slowdown versus skipping quantization entirely.

What we verified on the real model (53MB conformer-style model, same pipeline)

• convert_pt2e() output graph: 59 aten.linear.default, 34 quantized_decomposed.dequantize_per_channel.default,
  34 quantized_decomposed.quantize_per_tensor.default, 34 quantized_decomposed.dequantize_per_tensor.default
  -- i.e. 34 of the 59 Linear ops are correctly annotated/quantized at this stage.
• Confirmed torch._inductor.config.pre_grad_custom_pass is xiq.quant_lift_up (QuantLiftUp
  instance) -- correctly registered via torchao/quantization/pt2e/quantizer/x86_inductor_quantizer.py's
  module-level side effect.
• Manually invoked QuantLiftUp.__call__ directly on the graph -- zero nodes added/removed/changed
  (expected: it only handles attention-style Linear -> view/permute -> quantize chains; not
  the bug, just ruled out as the cause).
• Ran the real aoti_compile_and_package() three separate times (once with EXHAUSTIVE
  autotune + LTO + freezing=True, once with Inductor defaults, once repeated) --
  torch._dynamo.utils.counters["inductor"] never contains any qlinear/qconv-related
  key after any of them.
• Directly inspected the generated *.kernel.cpp in all three compiled packages: every
  _micro_gemm_kernel function takes const float* A, const float* B, float* C
  (0 out of 14 signatures use int8_t*/uint8_t*), and grep -c onednn on the generated
  .cpp files returns 0.
• .wrapper.so is ~44MB vs. ~53MB of fp32 params -- weights are somewhat more compact
  (consistent with int8 storage) but nowhere near the ~4x reduction genuine int8 packing
  would give, and irrelevant anyway since compute happens in fp32 regardless.

Expected behavior

dequantize_per_channel(weight) + dequantize_per_tensor(activation) + aten.linear + quantize_per_tensor(output) should be fused into torch.ops.onednn.qlinear_pointwise
and lowered to an accelerated AVX-512 VNNI int8 GEMM (CppMicroGemmAVX512VNNI in
torch/_inductor/codegen/cpp_micro_gemm.py), the same way it apparently does through
the documented torch.compile() path.

Actual behavior

Silent fallback: the compiled package is numerically correct but computes every
Linear/GEMM in fp32, with quantize/dequantize conversions wrapped around each one for
no benefit. No warning or error is raised anywhere in the pipeline.

Questions for maintainers

1. Is torch.export() + aoti_compile_and_package() an officially supported path for
   PT2E-quantized CPU models at all, or is torch.compile() currently the only
   supported route to get real int8 acceleration?
2. If it's meant to work: what's missing from the pipeline above to get the
   weight-prepack fusion (_register_qlinear_weight_prepack) to actually match?
3. If it's not meant to work yet: could convert_pt2e/aoti_compile_and_package
   at least emit a warning when quantize/dequantize nodes survive all the way into
   the compiled artifact unfused, rather than silently shipping a package that's
   doing pure quantization overhead with no compute benefit? This seems like an easy
   footgun for anyone following the AOTInductor deployment tutorials with a quantized
   model.

Steps to reproduce

Expected (if fusion works): "onednn references" > 0, and at least one
_micro_gemm_kernel with int8_t*/uint8_t* args.

Actual (observed bug): onednn references == 0, all _micro_gemm_kernel
signatures use float* -- i.e. weights are quantized to int8 in storage,
but every GEMM still dequantizes to fp32, computes in fp32, then requantizes
the output. Pure quantize/dequantize overhead, zero compute benefit.


import os
import zipfile

import torch
import torch.nn as nn
from torch.export import Dim, export
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
import torchao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq
from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import X86InductorQuantizer


class MiniConformerLike(nn.Module):
    """Mimics the structural shape of the real model that triggers this:
    a self-attention block (nn.MultiheadAttention) feeding a plain Linear FFN,
    both under the same 2D/3D reshape-heavy pattern PT2E quantization targets.
    """

    def __init__(self, d_model=256, nhead=4, dim_ff=768):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.ff1 = nn.Linear(d_model, dim_ff)
        self.ff2 = nn.Linear(dim_ff, d_model)
        self.act = nn.ReLU()

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x, need_weights=False)
        x = x + attn_out
        x = x + self.ff2(self.act(self.ff1(x)))
        return x


def main():
    torch.manual_seed(0)
    model = MiniConformerLike().eval()

    example_inputs = (torch.randn(1, 128, 256),)
    seq_len = Dim("seq_len", min=16, max=512)
    dynamic_shapes = {"x": {1: seq_len}}

    with torch.no_grad():
        ep = export(model, example_inputs, dynamic_shapes=dynamic_shapes)

    quantizer = X86InductorQuantizer()
    quantizer.set_global(xiq.get_default_x86_inductor_quantization_config())
    prepared = prepare_pt2e(ep.module(), quantizer)

    with torch.no_grad():
        prepared(*example_inputs)  # calibrate

    quantized = convert_pt2e(prepared)

    with torch.no_grad():
        quantized_ep = export(quantized, example_inputs, dynamic_shapes=dynamic_shapes)

    import torch._inductor.config as inductor_config
    inductor_config.cpp_wrapper = True
    inductor_config.freezing = True

    package_path = "/tmp/repro_qlinear_fusion_bug.pt2"
    torch._inductor.aoti_compile_and_package(
        quantized_ep,
        package_path=package_path,
    )

    unzip_dir = "/tmp/repro_qlinear_fusion_bug_unzip"
    with zipfile.ZipFile(package_path) as zf:
        zf.extractall(unzip_dir)

    onednn_refs = 0
    float_gemm_kernels = 0
    int8_gemm_kernels = 0
    for root, _dirs, files in os.walk(unzip_dir):
        for f in files:
            if not f.endswith(".cpp"):
                continue
            text = open(os.path.join(root, f), errors="replace").read()
            onednn_refs += text.count("onednn")
            for line in text.splitlines():
                if "_micro_gemm_kernel(" in line and "inline void" in line:
                    pass  # signature spans multiple lines; checked via arg scan below
            # crude scan: look at declared arg types following each micro_gemm_kernel def
            idx = 0
            while True:
                idx = text.find("_micro_gemm_kernel(", idx)
                if idx == -1:
                    break
                snippet = text[idx: idx + 300]
                if "const float*" in snippet:
                    float_gemm_kernels += 1
                if "const int8_t*" in snippet or "const uint8_t*" in snippet:
                    int8_gemm_kernels += 1
                idx += 1

    print(f"onednn references:        {onednn_refs}")
    print(f"float32 micro_gemm kernels: {float_gemm_kernels}")
    print(f"int8/uint8 micro_gemm kernels: {int8_gemm_kernels}")

    if onednn_refs == 0 and float_gemm_kernels > 0 and int8_gemm_kernels == 0:
        print("\nBUG REPRODUCED: quantized Linear layers compiled to fp32 GEMM, "
              "no oneDNN/VNNI fusion.")
    else:
        print("\nDid NOT reproduce with this minimal model -- fusion may be "
              "specific to a different structural pattern. Do not use this "
              "script as-is in the bug report if this happens; investigate "
              "further before filing.")


if __name__ == "__main__":
    main()

Environment

- `torch==2.14.0.dev20260710+cpu`
- `torchao==0.18.0.dev20260712+cpu`
- Linux, Intel Xeon Gold 6338 (Ice Lake-SP, AVX-512 + AVX512_VNNI)
- `torch.backends.mkldnn.enabled == True`
- `torch.backends.mkldnn.is_available() == True`
- `torch.ops.mkldnn._is_mkldnn_acl_supported() == False` (x86, not ARM/ACL)