ao

Highlights

We are excited to announce the 0.17 release of torchao! This release adds support for cuteDSL MXFP8 MoE kernels, per-head FP8 quantized low precision attention, ABI stability, and more!

CuteDSL MXFP8 MoE Kernels

We added a new CuteDSL MXFP8 quantization kernel for 3d expert weights that writes scale factors directly to blocked layout for tensorcores: https://github.com/pytorch/ao/pull/4090

• Used for scaling along dim1 in the backward pass of MoE training with grouped GEMMs.
• ~12% speedup over previous 2 kernel “quantize then scale layout transformation” approach!

Per-Head FP8 Quantized Low Precision Attention

We added a new API for per-head fp8 quantized attention with FA3 as the backend (https://github.com/pytorch/ao/pull/3959 and https://github.com/pytorch/ao/pull/3857)

• Users can either choose to use the elementary blocks as direct replacements for `F.scaled_dot_product_attention` or use the high-level wrapper, which replaces all F.SDPA calls within a module with the low precision attention variant.
• Running torch.compile on a wrapped module will enable RoPE fusion where appropriate
• Results show a 1.18x speedup on Wan2.1-T2V-1.3B, 1.23x speedup on LLaMA 3 prefill with high sequence lengths (131k), 1.07x speedup on flux.1-schnell with 2048x2048 image size

Example Usage of Direct Replacement:

from torchao.prototype.attention.fp8_fa3 import fp8_fa3_sdpa, fp8_fa3_rope_sdpa
out = fp8_fa3_sdpa(q, k, v)

Example Usage of Wrapper:

from torchao.prototype.attention import (
    AttentionBackend,
    LowPrecisionAttentionConfig,
    apply_low_precision_attention,
)
# Instantiate any nn.Module()
model = MyModel()

# Simple SDPA replacement
config = LowPrecisionAttentionConfig(backend=AttentionBackend.FP8_FA3)
model = apply_low_precision_attention(model, config)

# Flash activation is handled internally by the wrapper
output = model(inputs)

# Torch.compile will enable rope fusion
model = torch.compile(model)

PyTorch ABI stability

As of https://github.com/pytorch/ao/issues/3516, torchao is now ABI stable for all cuda kernels! This means if the user is running torch 2.11+, they will be able to access torchao’s cuda kernels without having to upgrade torch version for each new torchao version. This applies to the current and all future torchao releases (0.17.0+). Note that python-only API compatibility is the same as before: we support the latest 3 torch minor versions.

Before:

# Compatible versions for cpp extensions:
torchao 0.16.0 + torch 2.10.0
torchao 0.15.0 + torch 2.9.1
torchao 0.14.1 + torch 2.9.0

# Compatible versions for python-only API
# 3 most recent torch versions:
torchao 0.16.0 + torch 2.10.0, 2.9.1, 2.8.0
torchao 0.15.0 + torch 2.9.1, 2.8.0, 2.7.1
torchao 0.14.0 + torch 2.9.0, 2.8.0, 2.7.0

After:

# Compatible versions for cpp extensions (all future releases):
torchao 0.17.0 onwards + any torch 2.11+

# Compatible versions for python-only API
# 3 most recent torch versions (same as before)
torchao 0.17.0 + torch 2.11.0, 2.10.0, 2.9.1

Deprecations

• Deprecate AQT and related classes (https://github.com/pytorch/ao/pull/4074)

We are planning to remove these classes in the future, so we are deprecating them in advance. Here are all the classes that now have new deprecation warnings:

  Base classes (torchao/dtypes/utils.py):
  1. Layout
  2. PlainLayout
  3. AQTTensorImpl

  torchao/dtypes/affine_quantized_tensor.py:
  4. AffineQuantizedTensor

  torchao/dtypes/floatx/:
  5. Float8Layout
  6. Float8AQTTensorImpl
  7. CutlassSemiSparseLayout
  8. CutlassSemiSparseTensorImpl

  torchao/dtypes/

Highlights

We are excited to announce the 0.16.0 release of torchao! This release adds support for MXFP8 MoE Building Blocks for Training with Expert Parallelism and deprecated older versions of some configs and less used quantization options to keep torchao leaner! We also revamped our doc page, README and made some progress in making torchao ABI stable.

MXFP8 MoE Building Blocks for Training with Expert Parallelism

This release includes the following differentiable building blocks for MXFP8 MoE Training with expert parallelism:

• a2a_dispatch_mxfp8_fwd_hp_bwd: All-to-all token dispatch (MXFP8 forward pass, BF16 backward pass)
• permute_mxfp8_fwd_hp_bwd: Permute and pad tokens for MXFP8 computation (MXFP8 forward pass, BF16 backward pass)
• _to_mxfp8_then_scaled_grouped_mm: MXFP8 grouped GEMM for routed expert computation (new: optionally accepts pre-quantized inputs). Produces bfloat16 output.
• unpermute_hp_fwd_mxfp8_bwd: Unpermute tokens back to original order (BF16 forward pass, MXFP8 backward pass)
• a2a_combine_hp_fwd_mxfp8_bwd: All-to-all token combine (BF16 forward pass, MXFP8 backward pass). Note the actual combine/aggregation op does not happen here, the naming is just to indicate it is intended to be used for the all2all immediatley preceding the aggregation.

These autograd functions can be chained together to implement efficient MoE training with expert parallel comms and grouped GEMMs in MXFP8.

This approach achieves 10% - 25% tokens/second speedup for DeepSeekV3 16b training:

• +10% tokens/second on single node 8xB200 with NVLink intra-node networking for inter-device communication.
• +25% tokens/second on multi-node B200 cluster with IB inter-node networking and NVLink intra-node networking.

Deprecations

• Deprecate v1 of Float8WeightOnlyConfig, Float8DynamicActivationFloat8WeightConfig, Int8DynamicActivationIntxWeightConfig, IntxWeightOnlyConfig, Int4WeightOnlyConfig (https://github.com/pytorch/ao/pull/3510, https://github.com/pytorch/ao/pull/3511, https://github.com/pytorch/ao/pull/3512, https://github.com/pytorch/ao/pull/3513)

# v0.15.0 - version 1 was available.
config = Float8WeightOnlyConfig(version=1, ...)
config = Float8DynamicActivationFloat8WeightConfig(version=1, ...)
config = Int8DynamicActivationIntxWeightConfig(version=1, ...)
config = IntxWeightOnlyConfig(version=1, ...)
config = Int4WeightOnlyConfig(version=1, ...)

# v0.16.0 - use version 2 (default). Using version 1 is no longer supported.
config = Float8WeightOnlyConfig(version=2, ...)
config = Float8DynamicActivationFloat8WeightConfig(version=2, ...)
config = Int8DynamicActivationIntxWeightConfig(version=2, ...)
config = IntxWeightOnlyConfig(version=2, ...)
config = Int4WeightOnlyConfig(version=2, ...)

• Move Int8DynamicActivationInt4WeightConfig, Int4DynamicActivationInt4WeightConfig, GemliteUIntXWeightOnlyConfig, Float8StaticActivationFloat8WeightConfig, `UI

Highlights

We are excited to announce the 0.15.0 release of torchao! This release adds:

• MXFP8 MoE training demonstrates 1.2x e2e training speedup with identical convergence versus bf16, training Llama4 Scout on a 64 node GB200 Crusoe cluster!
• MXFP8 MoE kernels shipped with torchao builds for CUDA 12.8+ (just pip install instead of building from source to use!)
• Safetensors enablement
• Quantization with parameter level targeting

MXFP8 MoE training demonstrates 1.2x e2e training speedup with identical convergence versus bf16, training Llama4 Scout on a 64 node GB200 Crusoe cluster

Training runs on 64 node GB200 cluster with TorchTitan Llama4 Scout demonstrated a 1.2x e2e training speedup with equivalent convergence to bfloat16 training baseline. In fact, after 3,000 steps it finishes with slightly lower loss than bfloat16! This is consistent with our scaling experiments with MXFP8 training for dense models.

See the TorchAO MXFP8 MoE training documentation for more details. You can also check out the TorchTitan MXFP8 documentation to run pretraining jobs with TorchAO MXFP8 by adding a single config.

Safetensors Enablement

You can now save and load TorchAO model checkpoints using safetensors! This feature is integrated with Hugging Face transformers starting from v5.0.0 and vLLM 0.13.0 for model inference/serving.

We currently support the following stable configs: Float8DynamicActivationFloat8WeightConfig Int4WeightOnlyConfig IntxWeightOnlyConfig Int8DynamicActivationIntxWeightConfig Int8WeightOnlyConfig Int8DynamicActivationInt8WeightConfig

and will continue to add support for configs as they become stable in the future.

Example:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from torchao.quantization import Float8WeightOnlyConfig


model_id = "facebook/opt-125m"
quant_config = Float8WeightOnlyConfig()
quantization_config = TorchAoConfig(quant_type=quant_config)
quantized_model = AutoModelForCausalLM.from_pretrained(
   model_id,
   device_map="auto",
   torch_dtype=torch.bfloat16,
   quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)


#### Push to hub
MODEL_NAME = model_id.split("/")[-1]
save_to = f"torchao-testing/{MODEL_NAME}-Float8WeightOnlyConfig-v2-0.15.0.dev-safetensors"
quantized_model.push_to_hub(save_to, safe_serialization=True)
tokenizer.push_to_hub(save_to)

To serve a safetensors model checkpoint on vLLM, use the CLI command vllm serve torchao-testing/Qwen3-8B-INT4-0.15.0dev-safetensors.

int4 weight only quantization with preshuffled packing format is supported in vllm (#3245, #26066)

Now we support running int4 weight only with preshuffled packing format in vllm. We can ship int4 weight only quantized checkpoint with plain packing format (e.g. https://huggingface.co/torchao-testing/opt-125m-Int4WeightOnlyConfig-v2-0.14.0.dev), and the Tensor with plain format will be converted to preshuffled format automatically in SM90+ and when fbgemm_gpu_genai is available.

Support for quantizing any parameter by its fully qualified name (FQN) via FqnToConfig (#[3083](https://github.com/pytorch/ao/pul

Highlights

We are excited to announce the 0.14.1 release of torchao! This release adds support for MoE training on Backwell GPUs and NVFP4 QAT!

(Prototype) MoE training on Blackwell GPUs

We’ve added a quantized building block for speeding up MoE training on Blackwell GPUs: torchao’s `_scaled_grouped_mm`! It is a differentiable drop-in replacement for `torch._grouped_mm` that dynamically quantizes inputs using the given recipe, performs a scaled grouped GEMM, then returns the results in original precision. This results in significant speedups (see benchmarks below)!

import torch
from torch.nn import functional as F
from torchao.prototype.moe_training import (
    _scaled_grouped_mm as torchao_scaled_grouped_mm
)
from torchao.prototype.moe_training.conversion_utils import MoEScalingType
from torchao.prototype.moe_training.utils import generate_jagged_offs

num_groups, total_M, N, K = 8, 131072, 8192, 5120

# A = input actvations, B = expert weights
A = torch.randn(total_M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True)
B = torch.randn(num_groups, N, K, dtype=torch.bfloat16, device="cuda", requires_grad=True)

# Token group offsets computed by router in actual MoE layer
offs = generate_jagged_offs(num_groups, total_M, device="cuda")

# Forward and backward example
out = torchao_scaled_grouped_mm(
        A,
        B.transpose(-2, -1),
        offs=offs,
        scaling_type=MoEScalingType.MXFP8,
)
labels = torch.ones_like(out)
loss = F.mse_loss(out, labels)
loss.backward()

Microbenchmarks (see README for commands to reproduce benchmarks):

• Forward + backward pass vs torch._grouped_mm:
  • ~1.4-1.8x faster for Llama4 17bx16e shapes
  • ~1.2-1.4x faster for DeepSeekV3 671b shapes
• Full MoE layer forward + backward pass:
  • ~1.4x faster (Llama4 17bx16e shapes, batch_size=8, seq_len=16384)
  • ~1.2x faster (DeepSeekV3 671b shapes, batch_size=8, seq_len=16384).

It’s also already integrated into TorchTitan for E2E training with DeepSeekV3 and Llama4! Just use the command line flag: `--model.converters=”quantize.grouped_mm.mx”, which will convert all `torch._grouped_mm` ops to torchao _scaled_grouped_mm ops under the hood:

Torchtitan e2e training benchmarks (see README for commands to reproduce benchmarks):

• ~1.4x e2e speedup of 2 layer Llama4 16e model with dp2ep (without TP) parallelism on 4 B200 GPUs connected via NVLink

(Prototype) NVFP4 QAT (https://github.com/pytorch/ao/pull/3050)

We added quantization-aware training (QAT) support for NVFP4 as a prototype feature! This feature is currently only available on blackwell GPUs:

from torchao.quantization import quantize_
from torchao.prototype.mx_formats import NVFP4InferenceConfig

# NVFP4 activation and weight quantization
base_config = NVFP4InferenceConfig()
quantize_(model, QATConfig(base_config, step="prepare"))
train(model)
quantize_(model, QATConfig(base_config, step="convert"))

Our NVFP4 QAT support is also integrated into axolotl:

axolotl train examples/llama-3/3b-qat-fsdp2-nvfp4.yaml

Initial experimentation demonstrated that QAT recovered up to 41% mmlu_pro accuracy degradation and 33% bbh accuracy degradation from NVFP4 quantization, for Qwen3-8B and Gemma3-12B-it respectively:

# Qwen3-8B
bbh: 0.779 (baseline) -> 0.7254 (quant) -> 0.7368 (qat), recovered 21.269%
mmlu_pro: 0.4969 (baseline) -> 0.4521 (quant) -> 0.4707 (qat), recovered 41.518%

# Gemma3-12B-it accuracy
bbh: 0.7527 (baseline) -> 0.7068 (quant) -> 0.7222 (qat), recovered 33.551%
mmlu_pro: 0.4074 (baseline) -> 0.3621 (quant) -> 0.3702 (qat), recovered 17.881%

BC Breaking

Rename In