Int8Tensor/Float8Tensor missing aten.view/abs/arithmetic op registrations in 0.18 (quantize fails silently, unquantized)

[JWLHS]

Summary

torchao 0.18.0+xpu rewrote the quantized tensor subclasses (Int8Tensor / Float8Tensor) to use __torch_dispatch__ tables, but several operators are missing from the dispatch registrations of exactly these two classes. This breaks quantize_() whenever the flow calls aten.view on an Int8Tensor or aten.abs/arithmetic ops on a Float8Tensor — e.g. re-quantizing an already-quantized module goes through _choose_qparams_affine which calls input.view(...).

Worse, the failure is swallowed by per-module try/except in callers: the module is left unquantized and only a Failed to quantize ... message is printed, so users believe quantization is on when it isn't.

Environment

• OS: Windows 11
• Python: 3.12.9
• torch: 2.13.0+xpu (Intel Arc A770)
• torchao: 0.18.0+xpu (from https://download.pytorch.org/whl/xpu)

Reproduction

import torch
from torchao.quantization.quant_api import quantize_ as torchao_quantize_
from torchao.quantization.quant_api import Int8WeightOnlyConfig, Float8WeightOnlyConfig

model = torch.nn.Sequential(torch.nn.Linear(64, 128), torch.nn.Linear(128, 16)).to("xpu")

# First pass: OK (weights become Int8Tensor)
torchao_quantize_(model, Int8WeightOnlyConfig())

# Second pass on an already-quantized module (what ai-toolkit's quantize() does
# when it quantizes the container, then re-visits each child Linear):
torchao_quantize_(model[0], Int8WeightOnlyConfig())
# NotImplementedError: Int8Tensor dispatch: attempting to run unimplemented
# operator/function: func=<OpOverload(op='aten.view', overload='default')>, ...

# Direct repro:
model[0].weight.view(64, 128)  # same NotImplementedError

Float8Tensor equivalent:

model = torch.nn.Sequential(torch.nn.Linear(64, 128), torch.nn.Linear(128, 16)).to("xpu")
torchao_quantize_(model, Float8WeightOnlyConfig())
model[0].weight.abs()  # NotImplementedError: Float8Tensor dispatch: aten.abs unimplemented
torchao_quantize_(model[0], Float8WeightOnlyConfig())
# after registering abs, next missing: aten.div (see below)

Full stack for int8:

torchao_quantize_ -> ... -> _choose_qparams_affine -> input.view(shape_for_reduction)
-> torchao.utils._dispatch__torch_dispatch__ -> NotImplementedError

Root cause analysis (dispatch tables in 0.18.0)

torchao/quantization/quantize_/workflows/int8/int8_tensor.py registers only:
aten.linear, aten.slice, aten.index, aten.embedding, aten.is_pinned,
aten._pin_memory, aten.select — no aten.view.

torchao/quantization/quantize_/workflows/float8/float8_tensor.py registers
aten.view (line ~863) but is missing aten.abs, and the requantize path
then also hits missing aten.div (and likely mul/add/sub/round/clamp/neg).

Sibling classes (Int4Tensor, NF4Tensor, NVFP4Tensor, MXTensor) all register
aten.view, which suggests this is an omission during the 0.18 rewrite rather than
an intentional design.

Because the dispatch tables are pure-Python per-class tables, this is not XPU-specific:
the same missing registrations should fail on CUDA as well (we only verified on XPU here).

Expected behavior

aten.view on Int8Tensor and aten.abs/basic arithmetic on Float8Tensor should
either be implemented (shape/view on the quantized data; abs/arith via dequantize
then re-quantize), or quantize_() should explicitly no-op on already-quantized
modules instead of failing per-module.

Workaround (verified working)

Runtime registration into the class dispatch tables works:

from torchao.quantization import Int8Tensor, Float8Tensor

@Int8Tensor.implements(torch.ops.aten.view.default)
def _view(func, types, args, kwargs):
    t = args[0]
    size = args[1] if len(args) > 1 else (kwargs or {}).get("size")
    return t.dequantize().view(size)

@Float8Tensor.implements(torch.ops.aten.abs.default)
def _abs(func, types, args, kwargs):
    return torch.abs(args[0].dequantize())
# plus div/mul/add/sub/round/clamp/neg with dequantize-based impls

With these registrations, both int8 and float8 quantize_() + forward pass on XPU.

[JWLHS]

Follow-up correction from deeper testing:

• Float8Tensor missing aten.abs is NOT new in 0.18 — it already exists in torchao 0.17.0+xpu. Reproduced there: second quantize pass on an already-quantized Float8Tensor module fails identically with \Float8Tensor dispatch: aten.abs unimplemented\ (silently skipped by per-module try/except).
• 0.17 int8 is unaffected because it still uses the legacy \AffineQuantizedTensor\ (aten.view registered, and re-quantize is idempotent — it skips already-quantized weights).
• 0.18 regression for int8: int8 weight-only switched to the new \Int8Tensor\ (missing aten.view) and lost the idempotent skip, so the double-quantize path now fails for both int8 and float8.

So the concrete asks stand: implement aten.view for Int8Tensor and aten.abs (+ the arithmetic family hit by re-quantization) for Float8Tensor, or make quantize_() idempotent for already-quantized modules (as 0.17 int8 effectively was).

[yashb98]

Hi @jerryzh168 @andrewor14. I’d like to investigate the repeated-quantization behavior here, starting with CPU and DGX Spark reproduction. I would keep this separate from #4746’s operator/FSDP work and propose explicit behavior for repeated calls with the same versus different recipes. Is that scope available, and is it the direction you’d prefer?

[JWLHS]

Hi @yashb98, thanks for picking this up — the plan sounds good.

One scope clarification on our side: we only observed and tested this on Intel XPU (torch 2.13.0+xpu, torchao 0.18.0+xpu). The "likely affects CUDA too" note in our report is an inference from the pure-Python dispatch tables, not something we verified — please treat it as unconfirmed until your CPU/DGX Spark reproduction.

For reference, what works for us on XPU today is runtime registration into the class dispatch tables (dequantize-based implementations):

import torch
from torchao.quantization import Int8Tensor, Float8Tensor

@Int8Tensor.implements(torch.ops.aten.view.default)
def _view(func, types, args, kwargs):
    t = args[0]
    size = args[1] if len(args) > 1 else (kwargs or {}).get("size")
    return t.dequantize().view(size)

@Float8Tensor.implements(torch.ops.aten.abs.default)
def _abs(func, types, args, kwargs):
    return torch.abs(args[0].dequantize())
# plus div/mul/add/sub/round/clamp/neg with dequantize-based impls

This makes both int8 and float8 quantize_() plus forward pass work on XPU. The full script lives in our fork: https://github.com/JWLHS/ai-toolkit-xpu/blob/main/fix_torchao_018_xpu.py

Happy to run additional tests once you have a proposed fix. Good luck with the CPU/DGX Spark reproduction!

[yashb98]

I opened a focused PR for the Int8 and Float8 weight-only
reapplication path reported here.

The problem was that a second quantize_() call selected an nn.Linear whose
weight was already a TorchAO tensor, then sent that tensor through from_hp()
again. Int8 reached unsupported aten.view, and Float8 reached unsupported
aten.abs.

The patch now checks the stored tensor representation before either
weight-only handler calls from_hp(). An equivalent config leaves an
already-quantized ordinary nn.Linear, its parameter, stored tensors,
metadata, and eager output unchanged. A TorchAOBaseTensor with an
incompatible family or readable stored metadata raises ValueError before
that weight is modified. A compatible Float8Linear child is unwrapped to
nn.Linear while preserving its weight and bias identities and evaluation
state, with output matching torch.nn.functional.linear. Meta and FakeTensor
Int8 zero-point values are unavailable, so that path checks the remaining
recipe-identifying stored metadata without reading tensor data.

I tested the patch on Linux aarch64 with Python 3.11.15 and torch
2.13.0+cu130, on CPU and on a GIGABYTE AI TOP ATOM with NVIDIA GB10 at
compute capability 12.1:

• focused CPU, CUDA, meta, and FakeTensor selection: 94 passed, 47 deselected
• full non-PT2E CPU suite: 1051 passed, 7067 skipped, 102 warnings
• Int8 tensor test file: 261 passed, 1 skipped, 87 warnings
• complete Float8 workflow: 6 failed, 542 passed, 92 skipped, 16 warnings

The six Float8 failures have the same test IDs and terminal exceptions on
clean main at b7ac3aacf9b6f1cb0c6bb29143329f601dfcfcb1. Four require an
unavailable sparse CUTLASS operator, and two dynamic-activation grouped-matmul
tests reject compute capability 12.1. None of the six exercises the
weight-only handlers changed in the PR.

The PR is limited to Int8WeightOnlyConfig and Float8WeightOnlyConfig
reapplication. It does not add direct view or abs support. Repeated
activation-quantization reapplication is outside this PR.

I could not run XPU locally because this torch build has no XPU support.
Validation on the reported Windows 11 and Arc A770 setup would cover that
remaining hardware gap.

[JWLHS]

XPU validation data (Windows 11, Intel Arc A770, torch 2.14.0+xpu, torchao 0.18.0+xpu):

1. Int8WeightOnlyConfig: first quantize_() passes. A direct weight.view(...) and a second quantize_() on an already-quantized nn.Linear both fail with NotImplementedError: Int8Tensor dispatch: ... aten.view.default. This is a clear dispatch-table omission (int8 itself is fully supported on XPU).
2. Float8WeightOnlyConfig: weight.abs() and a second quantize_() fail with Float8Tensor dispatch: ... aten.abs.default. Caveat for this half of the report: XPU allocates float8_e4m3fn/float8_e5m2 storage, but FP8 compute support on XPU is incomplete, so this failure may be caused by missing dispatch registration, by absent underlying FP8 ops, or both. Treating the Float8 result as less conclusive than the Int8 result is appropriate until XPU FP8 kernels are confirmed available.
3. The dequantize-based runtime registrations from this issue make both Int8 view/re-quantize and Float8 abs/re-quantize succeed on XPU in the small eager repro. For Float8 this works because the workaround routes through dequantized (non-FP8) arithmetic.

Summary: Int8 aten.view is a confirmed registration bug on XPU. The Float8 failures are reproduced but should not be attributed purely to torchao until XPU FP8 op support is separately verified.