Int8Tensor: A PyTorch-Native INT8 Quantization Subclass for torchao

01. Intro

Int8Tensor is a quantized tensor subclass I built for PyTorch’s torchao library (PR #3391). Its primary contribution is dynamic activation quantization (W8A8-INT): both weights and activations are represented in INT8, with activation scales computed at runtime, enabling INT8×INT8 matrix multiplication on hardware with native integer arithmetic support.

W8A8-INT is harder than weight-only quantization because activation distributions vary per input — scales can’t be precomputed offline, and the kernel that consumes them is hardware-specific. Int8Tensor integrates via __torch_dispatch__, exposes a Granularity-based API for per-tensor and per-row quantization, and propagates scales through slice/select so it survives tensor parallelism.

📄 Full write-up: Int8Tensor report (PDF)

02. Quantization Modes

Int8Tensor supports two complementary strategies:

• Dynamic activation quantization (W8A8-INT). Both weights and activations are quantized to INT8; scales are computed at runtime since activation distributions vary per batch. This is the primary focus of this post.
• Weight-only quantization (FP16/BF16×INT8). Only weights are pre-quantized to INT8; activations stay in floating point. Included for broader hardware compatibility.

03. Technical Foundations

Asymmetric integer quantization

\[Q_x = \left\lfloor \frac{X}{s} + z \right\rceil\]

where $X$ is the floating-point tensor, $Q_x$ its $n$-bit quantized counterpart, and $s$ the scaling factor

\[s = \frac{X_{\max} - X_{\min}}{q_{\max} - q_{\min}}\]

with zero point $z = \lfloor q_{\min} - X_{\min}/s \rceil$. Dequantization recovers the approximation:

\[\hat{X} = (Q_x - z) \cdot s\]

Supported tensor operations

Operations are dispatched via __torch_dispatch__:

Operation	Description	Limitations
aten.linear.default	INT8×INT8 or FP×INT8 matmul	None
aten.slice.Tensor	Slicing with scale adjustment	dim ∈ {0, 1, 2}, step=1
aten.select.int	Single-index selection	dim=0 only
aten.index.Tensor	Advanced indexing	None

04. Technical Design

Quantization granularity

Int8Tensor supports per-tensor (single scale) and per-row (one scale per output channel) granularity; per-row is preferred for W8A8-INT. Rather than exposing block_size directly — per-row on an [N, K] matrix maps to block_size=[1, K], which is non-obvious — the API takes Granularity objects that express intent:

from ao.quantization.granularity import PerRow, PerTensor

Int8Tensor.from_hp(weight, granularity=PerRow())
Int8Tensor.from_hp(weight, granularity=PerTensor())

Dequantization

The scale shape depends on granularity: scalar for per-tensor, [N] or [N, 1] for per-row. Rather than implementing scale broadcasting manually, Int8Tensor delegates to dequantize_affine:

def dequantize(self, output_dtype=None):
    return dequantize_affine(
        input=self.qdata,
        block_size=self.block_size,
        scale=self.scale,
        output_dtype=output_dtype,
    )

W8A8-INT linear kernel

The core challenge in W8A8-INT is computing $Y = XW^\top$ when both $X$ and $W$ are INT8 with separate scale tensors. Converting to float before the matmul eliminates the benefit of INT8 arithmetic. Integer accumulation into INT64 works on CPU but is unsupported by addmm_cuda on GPU. The implementation instead uses int_scaled_matmul, an ao kernel with CPU and CUDA coverage that performs fused INT8 matmul with scale application:

from ao.kernel import int_scaled_matmul

y_dot_scaled = int_scaled_matmul(
    tmp, w_vals_t, x_scales.reshape(-1, 1))
result = y_dot_scaled * w_scales

Weight-only path. For FP×INT8, weights are cast to the activation dtype before the matmul, deferring scale multiplication to avoid materializing a full floating-point weight copy:

w_vals_t = weight.qdata.t().to(activation.dtype)
m = torch.mm(activation.reshape(-1, activation.shape[-1]), w_vals_t)
result = m * weight.scale

Scale slicing

Slicing requires consistent scale adjustment across all granularity cases; _slice_scale handles this self-contained, without external dependencies:

def _slice_scale(scale, data_shape, dim, start, end, step):
    if scale.numel() <= 1:      # Per-tensor: scalar unchanged
        return scale
    if scale.ndim == 1:         # Per-row: slice along dim 0 only
        return (aten.slice.Tensor(scale, 0, start, end, step)
                if dim == 0 else scale)
    # Per-block: map data indices to scale indices
    block_size_for_dim = data_shape[dim] // scale.shape[dim]
    scale_start = start // block_size_for_dim
    scale_end = (end + block_size_for_dim - 1) // block_size_for_dim
    return aten.slice.Tensor(scale, dim, scale_start, scale_end, 1)

05. Design Principles

Three principles guided the design:

• Express intent, not implementation. Granularity objects let callers specify what quantization is intended without knowing the internal block_size mapping.
• Reuse validated primitives. Delegating dequantization to dequantize_affine avoids error-prone per-rank scale broadcasting.
• Verify hardware constraints early. int_scaled_matmul was chosen specifically for its CPU+CUDA coverage, since generic integer accumulation does not transfer across devices; _slice_scale is kept self-contained to avoid pulling in unrelated subsystems.

06. Benchmarks

Benchmarks run on Qwen3-8B (RTX 5090). Configurations: BF16 baseline, W8A8-INT, and W8A8-INT with torch.compile (INT8+compile).

Peak memory — C++ kernel

Metric	BF16	W8A8-INT	INT8+compile
Allocated (MiB)	15,642.3	11,397.5	8,431.5
Reserved (MiB)	15,658.0	20,678.0	16,086.0
Fragmentation (%)	0.1	44.9	47.6

Peak memory — Triton kernel

Metric	BF16	W8A8-INT	INT8+compile
Allocated (MiB)	15,642.3	8,429.3	8,431.5
Reserved (MiB)	15,658.0	18,302.0	17,274.0
Fragmentation (%)	0.1	53.9	51.2

The Triton backend reduces allocated memory by ≈46% (matching INT8+compile), while the C++ backend achieves only ≈27%. All quantized variants show high fragmentation (45–54%) versus near-zero for BF16, attributable to short-lived per-batch activation scale allocations.

07. Limitations

Int8Tensor’s dispatch surface (Table 1) is deliberately narrow: aten.linear.default plus a handful of indexing ops. That mirrors a broader constraint in torchao itself — INT8 support is scoped to GEMM, not the whole model graph.

That scoping matters because GEMM is not where a transformer forward pass ends. Two op classes are explicitly out of scope:

• Attention. scaled_dot_product_attention is not a linear, so it isn’t quantized by this subclass. QK^T and the attention-weight×V matmul stay in the original activation dtype (BF16/FP16), and even fused SDPA kernels expect FP inputs, not Int8Tensor operands. Serious INT8 attention support needs its own quantized kernel — quantizing Q/K/V and the softmax output with numerically stable scales — which is a separate line of work from the int_scaled_matmul path used for linear.
• Normalization. LayerNorm/RMSNorm involve reductions (mean, variance) over the un-quantized distribution and are numerically sensitive to low precision; they’re conventionally left in FP16/BF16 regardless of the quantization scheme used elsewhere. Int8Tensor doesn’t attempt to dispatch through them, so a quantized model still round-trips to floating point at every norm.

The practical effect: in a transformer block, only the linear projections (QKVO, MLP up/down) run in INT8×INT8. Attention math and normalization remain FP, and each transition between an Int8Tensor output and an FP-only op costs a dequantize. On memory-bound decode workloads this is usually fine — GEMM dominates FLOPs and bandwidth — but it caps the achievable speedup versus a hypothetical scheme that also quantizes attention, and it means the benchmarks in §6 measure GEMM-only savings, not end-to-end INT8 inference.

Acknowledgements

This work was contributed to ao (PR #3391). I’m especially grateful to Zerry Zhang of the PyTorch core team for invaluable mentorship and technical guidance, and thank the ao maintainers for their constructive feedback.

References

1. ao Contributors, “ao: PyTorch-Native Training-to-Serving Model Optimization,” GitHub, 2024. https://github.com/pytorch/ao
2. “QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving,” MLSys’25. https://arxiv.org/abs/2405.04532