Added TorchAO FP8 weight-only and static quantization + user guides

docsrc/conf.py

@@ -104,7 +104,7 @@
         "tutorials/_rendered_examples/distributed_inference",
     ],
     # Exclude pure utility modules that aren't standalone runnable examples.
-    "ignore_pattern": r"(utils\.py|rotary_embedding\.py|tensor_parallel_initialize_dist\.py)",
+    "ignore_pattern": r"(utils\.py|static_fp8_utils\.py|rotary_embedding\.py|tensor_parallel_initialize_dist\.py)",
 }
 
 # Setup the breathe extension

docsrc/debugging/troubleshooting.rst

@@ -215,6 +215,14 @@ Accuracy / Performance Issues
     See `NVIDIA ModelOpt documentation <https://nvidia.github.io/TensorRT-Model-Optimizer/>`_
     for the full list of built-in quantization configs and customization options.
 
+    For **FP8 weight-only** quantization with TorchAO (no calibration loop), see
+    :ref:`quantize_linear_fp8_woq` and :ref:`torch_export_flux_fp8_woq`. Those
+    examples keep ``dequantize_affine`` in the graph so TensorRT can emit an
+    ``IDequantizeLayer`` on an FP8 weight constant.
+
+    For **static FP8** (activations and weights, with a calibration loop), see
+    :ref:`quantize_linear_fp8_static`.
+
 ----
 
 Distributed / Tensor-Parallel Issues

docsrc/user_guide/shapes_precision/index.rst

@@ -2,7 +2,7 @@ Precision & Quantization
 =========================
 
 Control numerical precision with FP16, BF16, and mixed-precision autocast,
-and reduce model size with INT8/FP8/FP4 quantization via ModelOpt.
+and reduce model size with INT8/FP8/FP4 quantization via ModelOpt or TorchAO.
 
 .. toctree::
    :maxdepth: 1
@@ -12,3 +12,6 @@ and reduce model size with INT8/FP8/FP4 quantization via ModelOpt.
    quantization
    ../../tutorials/_rendered_examples/dynamo/vgg16_ptq
    Example: ViT FP8 Quantization <../../tutorials/_rendered_examples/dynamo/quantize_vit_fp8>
+   Example: TorchAO FP8 WOQ Linear <../../tutorials/_rendered_examples/dynamo/torchao/quantize_linear_fp8_woq>
+   Example: TorchAO Static FP8 Linear <../../tutorials/_rendered_examples/dynamo/torchao/quantize_linear_fp8_static>
+   Example: FLUX.1-dev FP8 WOQ <../../tutorials/_rendered_examples/dynamo/torchao/torch_export_flux_fp8_woq>

docsrc/user_guide/shapes_precision/quantization.rst

@@ -5,20 +5,23 @@ Quantization (INT8 / FP8 / FP4)
 
 Torch-TensorRT supports post-training quantization (PTQ) with **INT8**, **FP8**, and
 **FP4** precisions via NVIDIA's
-`ModelOpt <https://github.com/NVIDIA/TensorRT-Model-Optimizer>`_ library. ModelOpt
-inserts quantize/dequantize (QDQ) nodes into the model graph; Torch-TensorRT then
+`ModelOpt <https://github.com/NVIDIA/TensorRT-Model-Optimizer>`_ library, and
+**FP8 weight-only** and **static FP8** quantization via
+`TorchAO <https://github.com/pytorch/ao>`_. Quantizers insert
+quantize/dequantize (QDQ) nodes into the model graph; Torch-TensorRT then
 converts those nodes into TRT quantization layers and sets the appropriate builder flags.
 
 ----
 
 Prerequisites
 -------------
 
-Install ModelOpt (requires ``nvidia-modelopt``):
+Install ModelOpt (requires ``nvidia-modelopt``) and/or TorchAO:
 
 .. code-block:: bash
 
     pip install nvidia-modelopt
+    pip install torchao
 
 Hardware requirements:
 
@@ -152,6 +155,73 @@ Quantization also works with ``torch.compile``:
 
 ----
 
+TorchAO FP8 Weight-Only Quantization
+------------------------------------
+
+TorchAO ``Float8WeightOnlyConfig`` quantizes Linear weights to FP8 (e4m3) while
+leaving activations in BF16/FP16. Unlike ModelOpt PTQ, no calibration dataset is
+required.
+
+Default TorchAO ``Float8Tensor.dequantize`` decomposes into primitive ops. The
+examples promote weights to a ``Float8TensorNonDecomposed`` subclass so export
+emits ``torch.ops.torchao.dequantize_affine``, which Torch-TensorRT converts to
+``IDequantizeLayer``.
+
+.. code-block:: python
+
+    from torchao.quantization import Float8WeightOnlyConfig, quantize_
+
+    quantize_(model, Float8WeightOnlyConfig())
+    model = pre_process_model_for_export(model)  # emit dequantize_affine
+
+    with exclude_dq_from_constant_folding():
+        exp_program = torch.export.export(model, (example_input,), strict=True)
+
+    trt_model = torch_tensorrt.dynamo.compile(
+        exp_program,
+        inputs=[example_input],
+        min_block_size=1,
+        use_explicit_typing=True,
+        require_full_compilation=True,
+    )
+
+The intended engine keeps an FP8 weight constant plus a DQ prologue into GEMM.
+On **Blackwell**, Myelin can fuse that prologue into the matmul. On other GPUs
+DQ + GEMM may run as two kernels — that is still correct as long as the FP8
+weight is not constant-folded into a dense high-precision weight.
+
+See :ref:`quantize_linear_fp8_woq` for a toy Linear model and
+:ref:`torch_export_flux_fp8_woq` for FLUX.1-dev.
+
+----
+
+TorchAO Static FP8 Quantization
+--------------------------------
+
+Static FP8 quantizes **activations and weights**. Activation scales are chosen
+offline with min/max observers (per-tensor activations, per-channel weights),
+then Linear layers are rewritten so export emits
+``quantize_affine_float8_non_decomposed`` and
+``dequantize_affine_float8_non_decomposed``. Torch-TensorRT maps those ops to
+``IQuantizeLayer`` / ``IDequantizeLayer``. After fusion, GEMMs can run in FP8.
+
+.. code-block:: python
+
+    quantize_static_fp8(model, (example_input,), calibration_steps=10)
+    exp_program = torch.export.export(model, (example_input,), strict=True)
+    trt_model = torch_tensorrt.dynamo.compile(
+        exp_program,
+        inputs=[example_input],
+        enabled_precisions={torch.float8_e4m3fn},
+        min_block_size=1,
+        require_full_compilation=True,
+    )
+
+This needs a calibration loop (unlike weight-only). See
+:ref:`quantize_linear_fp8_static`.
+
+----
+
 How QDQ Nodes Are Converted
 -----------------------------
 
@@ -161,6 +231,15 @@ graph (inserted by ModelOpt), the
 ``IDequantizeLayer`` pairs. The TRT builder then fuses these with adjacent compute layers
 (e.g. Conv, Linear) to produce INT8 or FP8 kernel variants.
 
+For TorchAO weight-only graphs, ``torch.ops.torchao.dequantize_affine.default`` is
+mapped to ``IDequantizeLayer`` (weight constant stays FP8/INT, activations stay
+high precision).
+
+For TorchAO static FP8 graphs,
+``quantize_affine_float8_non_decomposed`` / ``dequantize_affine_float8_non_decomposed``
+map to ``IQuantizeLayer`` / ``IDequantizeLayer`` so both activations and weights
+participate in FP8 GEMM fusion.
+
 For FP4, ``torch.ops.tensorrt.dynamic_block_quantize_op.default`` nodes are converted
 via the dynamic block quantize converter, which uses TRT's
 ``add_dynamic_quantize`` API (TRT ≥ 10.8).
@@ -228,3 +307,9 @@ Troubleshooting
 **FP4 "requires TRT ≥ 10.8" error**
     Upgrade TensorRT. FP4 uses ``add_dynamic_quantize`` which is only available in
     TRT 10.8 and newer.
+
+**TorchAO FP8 weights folded to BF16/FP16**
+    Export or Torch-TensorRT constant folding removed ``dequantize_affine``.
+    Promote weights with ``pre_process_model_for_export`` and wrap export in
+    ``exclude_dq_from_constant_folding`` as in :ref:`quantize_linear_fp8_woq`.
+    Install ``torchao`` so the converter and constant-folding exclusion register.

examples/dynamo/README.rst

@@ -25,6 +25,9 @@ Model Zoo
 * :ref:`_torch_export_llama2`: Compiling a Llama2 model using AOT workflow (`ir=dynamo`)
 * :ref:`_torch_export_sam2`: Compiling SAM2 model using AOT workflow (`ir=dynamo`)
 * :ref:`_torch_export_flux_dev`: Compiling FLUX.1-dev model using AOT workflow (`ir=dynamo`)
+* :ref:`quantize_linear_fp8_woq`: TorchAO FP8 weight-only quantization of a Linear layer (``examples/dynamo/torchao``)
+* :ref:`quantize_linear_fp8_static`: TorchAO static FP8 (act + weight) quantization of a Linear layer (``examples/dynamo/torchao``)
+* :ref:`torch_export_flux_fp8_woq`: Compiling FLUX.1-dev with TorchAO FP8 weight-only quantization (``examples/dynamo/torchao``)
 * :ref:`debugger_example`: Debugging Torch-TensorRT Compilation
 * :ref:`torch_export_3d_rope`: Compiling a 3D RoPE video-transformer block with complex numerics support
 * :ref:`engine_converter_binding_names`: Naming input / output bindings when emitting a raw serialized TRT engine

examples/dynamo/requirements.txt

@@ -6,3 +6,4 @@ matplotlib
 pandas
 huggingface_hub
 opencv-python
+torchao

examples/dynamo/torchao/README.rst

@@ -0,0 +1,22 @@
+"""
+TorchAO quantization
+====================
+
+Compile models quantized with `TorchAO <https://github.com/pytorch/ao>`_ using
+the Torch-TensorRT Dynamo backend.
+
+Weight-only FP8 keeps activations in BF16/FP16. Export emits
+``dequantize_affine``, which Torch-TensorRT maps to TensorRT
+``IDequantizeLayer`` so the engine can keep an FP8 weight constant.
+
+Static FP8 also quantizes activations. Calibrate observers, then export
+``quantize_affine_float8_non_decomposed`` / ``dequantize_affine_float8_non_decomposed``
+so TensorRT can fuse Q/DQ into FP8 GEMMs.
+
+.. code-block:: bash
+
+    pip install -r ../requirements.txt
+    python quantize_linear_fp8_woq.py
+    python quantize_linear_fp8_static.py
+    python torch_export_flux_fp8_woq.py
+"""

examples/dynamo/torchao/quantize_linear_fp8_static.py

@@ -0,0 +1,119 @@
+"""
+.. _quantize_linear_fp8_static:
+
+TorchAO Static FP8 Quantization (Linear)
+========================================
+
+This example calibrates a two-layer Linear model with TorchAO observers, then
+rewrites those layers so **activations and weights** are quantized to FP8
+(e4m3). Export keeps explicit ``quantize_affine_float8_non_decomposed`` /
+``dequantize_affine_float8_non_decomposed`` nodes, which Torch-TensorRT maps
+to ``IQuantizeLayer`` / ``IDequantizeLayer``.
+
+Contrast with :ref:`quantize_linear_fp8_woq`, which quantizes **weights only**
+and needs no calibration. Static FP8 can run the GEMM itself in FP8 (Tensor
+Cores) after Q/DQ fusion.
+
+Graph after export (one Linear)::
+
+    BF16 act ──► Q ──► FP8 act ──► DQ ──┐
+                                        ▼
+                                     aten.linear
+                                        ▲
+    FP8 weight ──► DQ ──────────────────┘
+
+Requirements:
+
+* NVIDIA GPU with FP8 support (Hopper or newer)
+* ``torchao``
+* ``torch-tensorrt`` with the TorchAO float8_non_decomposed converters
+
+"""
+
+# %%
+# Imports
+# ^^^^^^^
+# This example lives in ``examples/dynamo/torchao/``. Move that directory off
+# the front of ``sys.path`` so ``import torchao`` resolves the PyPI package
+# instead of this folder.
+
+import sys
+from pathlib import Path
+
+_EXAMPLE_DIR = str(Path(__file__).resolve().parent)
+if sys.path and Path(sys.path[0]).resolve() == Path(_EXAMPLE_DIR):
+    sys.path.pop(0)
+
+import torch
+import torch_tensorrt as torchtrt
+
+sys.path.insert(0, _EXAMPLE_DIR)
+from static_fp8_utils import quantize_static_fp8
+
+
+def sqnr(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Signal-to-quantization-noise ratio in dB (higher is closer)."""
+    a = a.float().flatten()
+    b = b.float().flatten()
+    signal = torch.norm(a)
+    noise = torch.norm(a - b)
+    return 20 * torch.log10(signal / noise.clamp_min(1e-12))
+
+
+# %%
+# Define a small two-layer Linear model
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+class LinearModel(torch.nn.Module):
+    def __init__(self, in_features=256, hidden=512, out_features=128):
+        super().__init__()
+        self.linear1 = torch.nn.Linear(in_features, hidden, bias=False)
+        self.linear2 = torch.nn.Linear(hidden, out_features, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear2(self.linear1(x))
+
+
+model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda")
+example_input = torch.randn(32, 256, dtype=torch.bfloat16, device="cuda")
+
+# %%
+# Calibrate, then insert static FP8 Q/DQ
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Observers record per-tensor activation min/max and per-channel weight min/max
+# over a few forward passes. Those scales are baked into ``QuantizedLinearQDQ``.
+
+with torch.no_grad():
+    fp_out = model(example_input)
+
+quantize_static_fp8(model, (example_input,), calibration_steps=10)
+
+with torch.no_grad():
+    quant_out = model(example_input)
+print(f"eager SQNR after static FP8: {sqnr(fp_out, quant_out):.2f} dB")
+
+# %%
+# Export and compile
+# ^^^^^^^^^^^^^^^^^^
+# The exported graph should contain
+# ``quantize_affine_float8_non_decomposed`` (activations) and
+# ``dequantize_affine_float8_non_decomposed`` (activations and weights).
+
+exp_program = torch.export.export(model, (example_input,), strict=True)
+exp_program.graph_module.print_readable()
+
+trt_model = torchtrt.dynamo.compile(
+    exp_program,
+    inputs=[example_input],
+    enabled_precisions={torch.float8_e4m3fn},
+    min_block_size=1,
+    require_full_compilation=True,
+)
+
+with torch.no_grad():
+    trt_out = trt_model(example_input)
+    if isinstance(trt_out, (list, tuple)):
+        trt_out = trt_out[0]
+print(f"TRT SQNR vs quantized eager: {sqnr(quant_out, trt_out):.2f} dB")
+print(trt_out)