If you have ever run a quantized LLM through plain PyTorch and wondered where the speed-up actually comes from, TorchAO is the answer hiding in plain sight. TorchAO (the torchao pip package, developed in the pytorch/ao repository) is PyTorch’s native quantization and sparsity library: it takes your ordinary nn.Linear layers, swaps their weights for compressed tensor representations, and dispatches the math to hand-written low-bit kernels. Unlike standalone quantization tools, everything it does is expressed in the language of the framework itself — tensor subclasses, dispatch, and torch.compile — which is exactly why its source is worth reading.

The second reason to study this codebase is architectural. TorchAO is not one big monolith; it is a thin public API on top of three cooperating layers: a recipe layer (dataclass configs that users touch), a transform layer (a registry that maps recipes to module surgery), and a tensor layer (subclasses like Int8Tensor and Float8Tensor that own the quantized data layout and their own kernels). Once you see how a config flows from quantize_() down to a .cpp file, you can predict where any new feature will land — and you can read llm-compressor or vLLM integration notes with far less guesswork.

This post is a source tour of the repository at main, walking the tree from the public API down to the metal. The two diagrams below were generated directly from the real file tree, and every symbol and line anchor in the text was verified against the current source.

The TorchAO overview: user recipes enter through quantize_(), get dispatched by a small handler registry into tensor subclasses, and terminate in native int8, int4, and float8 kernels.

Reading the overview from left to right:

• quantize_() — the eager entry point in quant_api.py, L292. It walks your model, filters for linear modules, and applies one recipe per module in place. Everything else in the library is downstream of this function.
• AOBaseConfig — the base class for all recipes at config.py, L27. Every public config dataclasses from it, which is what lets TorchAO serialize recipes to JSON and round-trip them through checkpoints.
• register_quantize_module_handler — the registry decorator in transform_module.py, L19. Each recipe type maps to exactly one transform function; quantize_() looks the handler up by config type and calls it.
• Int8Tensor — the dynamic int8 activation + int8 weight subclass in [int8tensor.py, L59](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/int8/int8_tensor.py#L59). Its linear override at L266 intercepts F.linear and routes to integer matmul kernels.
• Int4Tensor — the weight-only int4 subclass in [int4tensor.py, L57](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/int4/int4_tensor.py#L57). Its linear path at L190 now dispatches to fused rowwise int4 kernels instead of the older tinygemm op.
• NF4Tensor — the 4-bit NormalFloat representation in [nf4tensor.py, L580](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/nf4/nf4_tensor.py#L580), the same quantization family that made QLoRA popular, implemented as a first-class tensor.
• Float8Tensor — the fp8 inference tensor in [float8tensor.py, L84](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/float8/float8_tensor.py#L84), for Hopper-class GPUs where 8-bit floating point beats integer formats.
• quantize_affine — the numeric primitive in quant_primitives.py, L364. Scale/zero-point computation, requantization, and dequantization all reduce to this layer, with dequantize_affine as its mirror at L749.
• register_custom_op — the bridge in ops.py, L45. Ops like torchao::da8w4_linear_cpu declared here connect the Python tensor layer to the C++ and CUDA sources under torchao/csrc/.

From Install to Your First Model

Installation is deliberately boring: pip install torchao. The package builds shipped extension modules against a compatible PyTorch, and the top-level torchao/__init__.py checks versions at import time — if your torch is older than 2.11, it warns and skips loading the .so files rather than crashing, so the pure-Python APIs still work. On Apple Silicon there is an optional extra path: the torchao/experimental/ tree exposes Metal-backed low-bit kernels for MPS devices.

Your first quantized model is three lines:

import torch.nn as nn
from torchao import quantize_
from torchao.quantization.quant_api import Int4WeightOnlyConfig

model = MyModel().cuda()
quantize_(model, Int4WeightOnlyConfig(group_size=32))

What actually happens next is the interesting part. quantize_() uses the default filter to visit every nn.Linear, resolves the handler registered for Int4WeightOnlyConfig, and lets that handler replace each module’s weight with an Int4Tensor. The quantization itself respects a granularity object — PerRow at granularity.py, L79 or PerGroup at L58 — which maps down to the block_size argument of quantize_affine. After the swap, model.linear.weight is no longer a plain tensor; it is a dispatch table. Every subsequent forward pass hits the subclass’s __torch_dispatch__ implementations, so no model code changes at all.

Pick recipes by hardware. Int8DynamicActivationInt8WeightConfig (defined at quant_api.py, L852) is the CPU and torch.compile workhorse; Int8WeightOnlyConfig at L746 halves memory with near-zero accuracy loss; Int8StaticActivationInt8WeightConfig at L962 uses calibration-time scales for fastest integer math; and Float8DynamicActivationFloat8WeightConfig at L1156 targets fp8-capable GPUs. On the export side, the PT2E flow — prepare_pt2e at quantize_pt2e.py, L63 followed by convert_pt2e at L439 — produces a graph of quantized ops that ExecuTorch can deploy to phones and embedded targets.

CPU deployments deserve a special mention because TorchAO treats them as first-class. The Int8DynamicActivationIntxWeightConfig recipe at quant_api.py, L368 pairs dynamically quantized 8-bit activations with flexible 1-to-8-bit weights, and its CPU execution path bottoms out in the C++ da8w4_linear.cpp kernel rather than a Python fallback. If your architecture hides linear layers inside 1x1 convolutions — common in diffusion and vision models — the swap_conv2d_1x1_to_linear helper at quant_api.py, L236 rewrites those modules first so the recipes can see them. For maximum compression on memory-bound CPU inference, IntxWeightOnlyConfig at L1408 packs weights down to sub-byte widths and relies on lookup-table-style kernels to decode them on the fly.

The detailed view: recipe dataclasses carry packing formats and granularity into the transform registry, tensor subclasses own dispatch, and a thin custom-op bridge hands work to C++ and CUDA kernels.

Under the Hood: Recipes to Kernels

The transform chain is worth appreciating because it is so small. transform_module.py is a single dictionary plus a decorator; the intelligence lives in the handlers defined next to their recipes in quant_api.py. A handler receives the module and the config, calls the appropriate quantize-tensor function, and returns a new module whose weight is a subclass instance. There is no code generation and no global state — which is why recipes compose cleanly and why adding a new format rarely touches more than two files.

The same discipline shows up in the numeric plumbing. Every recipe ultimately calls choose_qparams_affine at quant_primitives.py, L1254 to derive scales and zero points for a block of weights, then hands the results to the tensor constructor. Storage layout is declared, not guessed: the PackingFormat enum at packing_format.py, L12 gives every subclass a stable identifier for how its bytes are arranged, which is what makes checkpoint round-tripping and cross-device transfers inspectable instead of magical.

Each tensor subclass is where layout meets execution. Int8Tensor stores quantized data plus scales and owns every aten op it cares about: linear, slice, embedding, and friends. Its kernel module defines safe_int_mm at kernels.py, L18 — a guarded integer matmul that pads dimensions to satisfy cuBLAS requirements — and _int_scaled_matmul_cpu at L79, which picks a u8s8 path when AVX512-VNNI is available and an s8s8 path on AMX hardware. The int4 side exposes Int4PackingFormat at [int4packing_format.py, L12](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/int4/int4_packing_format.py#L12) with plain, preshuffled, and tile-packed layouts, and Int4ChooseQParamsAlgorithm at [int4choose_qparams_algorithm.py, L12](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/int4/int4_choose_qparams_algorithm.py#L12) lets you trade a slower HQQ calibration for better rounding. The float8 recipes grow a KernelPreference field — the enum at kernel_preference.py, L17 — so a config can request Torch-native or specialized kernels and let the tensor resolve it at dispatch time.

Beneath the Python sits torchao/csrc/, organized by target: CPU shared kernels such as the 8-bit-activation x-bit-weight linear used by ExecuTorch deployments, an aten_kernels directory holding da8w4_linear.cpp for dynamic-activation int4-weight CPU inference, AArch64 NEON-dot headers for phone-class hardware, and CUDA sources for microscaling fp8 plus 2:4 semi-structured sparsity via CUTLASS on Hopper. The fp8 sparse tensor — Float8Sparse2x4_1DData1DMetadataTensor at [float8sparse_2x4_1d_data_1d_metadata_tensor.py, L42](https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize/workflows/float8/float8_sparse_2x4_1d_data_1d_metadata_tensor.py#L42) — shows the pattern end to end: the Python class declares a custom fp8_sparse_mm op, ops.py gives it a schema, and the CUDA file implements it.

Training is a sibling rather than an afterthought. The torchao/float8/ package converts a model for fp8 training with convert_to_float8_training at float8_linear_utils.py, L86, swapping eligible layers for Float8Linear at float8_linear.py, L206 with per-tensor or rowwise scaling strategies. QAT users get QATConfig at qat/api.py, L43, which inserts FakeQuantizedLinear at qat/linear.py, L42 modules so weights learn under the same numerics they will be quantized to later. Even optimizer memory is covered: OptimState8bit at subclass_8bit.py, L40 stores Adam moments as quantized tensors, extending the same TorchAOBaseTensor base at utils.py, L721 that the inference subclasses use.

Honest Limits

TorchAO is PyTorch-coupled by design, and that cuts both ways. The shipped binaries track specific torch releases; run an older torch and the extension load silently degrades to a warning, meaning your “quantized” model may fall back to dequantize-then-float math that is slower than the original. Kernel coverage is also hardware-shaped: the fast int4 rowwise paths assume modern NVIDIA GPUs or specific ARM CPUs, the fp8 and CUTLASS sparse paths need Hopper-class cards, and the fancy MPS Metal kernels are opt-in through the experimental API. On unsupported targets you get correct-but-unremarkable performance, not the benchmark numbers.

Accuracy is a per-model negotiation, not a guarantee. Weight-only int8 and group-wise int4 are broadly safe for LLMs, but dynamic-activation int8, fp8, and anything with coarse group sizes need numerical checking against your own eval — the repo’s testing utilities, including benchmark_model at utils.py, L75, help, but they do not replace a proper eval harness. The PT2E flow additionally requires an export-compatible model, and its prepare/convert contract has churned across PyTorch versions, so pin your versions. Finally, the repository moves fast: sparsity has migrated into prototype and float8 subtrees, the old linear_activation_quantized_tensor.py files are gone in favor of the workflows layout, and int4 kernels now route through the mslk ops. Any source tour — including this one — has a shelf life, so anchor to a commit, not to main.

Where TorchAO Sits in the Stack

Compared with llm-compressor, which produces compressed checkpoints for serving engines, or vLLM, which consumes them, TorchAO keeps everything inside eager PyTorch and torch.compile — you quantize a live module, not a file. That makes it the natural choice when you want quantization to survive fine-tuning, export cleanly through ExecuTorch, or compose with PyTorch training utilities without leaving the ecosystem. Read the three layers in order — recipe, transform, tensor — and the whole codebase opens up.

Enjoyed this post? Never miss out on future posts by following us