diff --git a/docsrc/conf.py b/docsrc/conf.py
index bd12a18854..cca631a7f3 100644
--- a/docsrc/conf.py
+++ b/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
diff --git a/docsrc/debugging/troubleshooting.rst b/docsrc/debugging/troubleshooting.rst
index 68ba01fa2f..e9abcb6aa8 100644
--- a/docsrc/debugging/troubleshooting.rst
+++ b/docsrc/debugging/troubleshooting.rst
@@ -215,6 +215,14 @@ Accuracy / Performance Issues
See `NVIDIA ModelOpt documentation `_
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
diff --git a/docsrc/user_guide/shapes_precision/index.rst b/docsrc/user_guide/shapes_precision/index.rst
index de0883e6e8..195e40bb1d 100644
--- a/docsrc/user_guide/shapes_precision/index.rst
+++ b/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>
diff --git a/docsrc/user_guide/shapes_precision/quantization.rst b/docsrc/user_guide/shapes_precision/quantization.rst
index e703e6f372..ade862fb7c 100644
--- a/docsrc/user_guide/shapes_precision/quantization.rst
+++ b/docsrc/user_guide/shapes_precision/quantization.rst
@@ -5,8 +5,10 @@ Quantization (INT8 / FP8 / FP4)
Torch-TensorRT supports post-training quantization (PTQ) with **INT8**, **FP8**, and
**FP4** precisions via NVIDIA's
-`ModelOpt `_ library. ModelOpt
-inserts quantize/dequantize (QDQ) nodes into the model graph; Torch-TensorRT then
+`ModelOpt `_ library, and
+**FP8 weight-only** and **static FP8** quantization via
+`TorchAO `_. 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.
----
@@ -14,11 +16,12 @@ converts those nodes into TRT quantization layers and sets the appropriate build
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.
diff --git a/examples/dynamo/README.rst b/examples/dynamo/README.rst
index 4b27ad92e6..847781cb39 100644
--- a/examples/dynamo/README.rst
+++ b/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
\ No newline at end of file
diff --git a/examples/dynamo/requirements.txt b/examples/dynamo/requirements.txt
index 09f1438780..cd95e6234d 100644
--- a/examples/dynamo/requirements.txt
+++ b/examples/dynamo/requirements.txt
@@ -6,3 +6,4 @@ matplotlib
pandas
huggingface_hub
opencv-python
+torchao
diff --git a/examples/dynamo/torchao/README.rst b/examples/dynamo/torchao/README.rst
new file mode 100644
index 0000000000..cb76b403d4
--- /dev/null
+++ b/examples/dynamo/torchao/README.rst
@@ -0,0 +1,22 @@
+"""
+TorchAO quantization
+====================
+
+Compile models quantized with `TorchAO `_ 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
+"""
diff --git a/examples/dynamo/torchao/quantize_linear_fp8_static.py b/examples/dynamo/torchao/quantize_linear_fp8_static.py
new file mode 100644
index 0000000000..6f1cb0d04c
--- /dev/null
+++ b/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)
diff --git a/examples/dynamo/torchao/quantize_linear_fp8_woq.py b/examples/dynamo/torchao/quantize_linear_fp8_woq.py
new file mode 100644
index 0000000000..f7f74982a6
--- /dev/null
+++ b/examples/dynamo/torchao/quantize_linear_fp8_woq.py
@@ -0,0 +1,84 @@
+"""
+.. _quantize_linear_fp8_woq:
+
+TorchAO FP8 Weight-Only Quantization (Linear)
+=============================================
+
+This example quantizes a toy ``nn.Linear`` with TorchAO
+``Float8WeightOnlyConfig``, keeps an explicit ``dequantize_affine`` in the
+exported graph, and compiles with the Torch-TensorRT Dynamo backend.
+
+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 the
+DQ + GEMM may run as two kernels — that is still correct as long as the FP8
+weight is **not** constant-folded into a dense FP16/BF16 weight.
+
+Requirements:
+
+* NVIDIA GPU with FP8 support (Hopper or newer). Blackwell for DQ fusion.
+* ``torchao``
+* ``torch-tensorrt`` with the ``torchao.dequantize_affine`` converter
+
+"""
+
+# %%
+# 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
+from torchao.quantization import Float8WeightOnlyConfig, quantize_
+
+sys.path.insert(0, _EXAMPLE_DIR)
+from utils import exclude_dq_from_constant_folding, pre_process_model_for_export
+
+# %%
+# Define a linear model and quantize weights to FP8
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+
+class LinearModel(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.linear = torch.nn.Linear(3072, 4096)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.linear(x)
+
+
+model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda")
+example_input = torch.randn(32, 3072, dtype=torch.bfloat16, device="cuda")
+
+quantize_(model, Float8WeightOnlyConfig())
+processed_model = pre_process_model_for_export(model)
+
+# %%
+# Export and compile
+# ^^^^^^^^^^^^^^^^^^
+# Wrap export in ``exclude_dq_from_constant_folding`` so inductor does not fold
+# ``dequantize_affine`` before Torch-TensorRT lowering. Torch-TensorRT also
+# marks that op impure in its own constant-folding pass.
+
+with exclude_dq_from_constant_folding():
+ exp_program = torch.export.export(processed_model, (example_input,), strict=True)
+
+trt_model = torchtrt.dynamo.compile(
+ exp_program,
+ inputs=[example_input],
+ min_block_size=1,
+ use_explicit_typing=True,
+ require_full_compilation=True,
+)
+
+output = trt_model(example_input)
+print(output)
diff --git a/examples/dynamo/torchao/static_fp8_utils.py b/examples/dynamo/torchao/static_fp8_utils.py
new file mode 100644
index 0000000000..1be42cd407
--- /dev/null
+++ b/examples/dynamo/torchao/static_fp8_utils.py
@@ -0,0 +1,174 @@
+"""Helpers for TorchAO static FP8 (activation + weight) quantization examples.
+
+Static FP8 calibrates activation and weight ranges with observers, then rewrites
+Linear layers to explicit ``quantize_affine_float8_non_decomposed`` /
+``dequantize_affine_float8_non_decomposed`` so Torch-TensorRT can emit
+``IQuantizeLayer`` / ``IDequantizeLayer`` pairs.
+"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torchao.core.config import AOBaseConfig
+from torchao.quantization import quantize_
+from torchao.quantization.granularity import PerAxis, PerTensor
+from torchao.quantization.observer import AffineQuantizedMinMaxObserver
+from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
+from torchao.quantization.quant_primitives import (
+ MappingType,
+ _dequantize_affine_float8_non_decomposed,
+ _quantize_affine_float8_non_decomposed,
+)
+from torchao.quantization.transform_module import register_quantize_module_handler
+
+
+class ObservedLinear(torch.nn.Linear):
+ """Linear that records activation and weight ranges during calibration."""
+
+ def __init__(
+ self,
+ in_features,
+ out_features,
+ act_obs,
+ weight_obs,
+ bias=True,
+ device=None,
+ dtype=None,
+ ):
+ super().__init__(in_features, out_features, bias, device, dtype)
+ self.act_obs = act_obs
+ self.weight_obs = weight_obs
+
+ def forward(self, input: Tensor):
+ observed_input = self.act_obs(input)
+ observed_weight = self.weight_obs(self.weight)
+ return F.linear(observed_input, observed_weight, self.bias)
+
+ @classmethod
+ def from_float(cls, float_linear, act_obs, weight_obs):
+ observed = cls(
+ float_linear.in_features,
+ float_linear.out_features,
+ act_obs,
+ weight_obs,
+ bias=float_linear.bias is not None,
+ device=float_linear.weight.device,
+ dtype=float_linear.weight.dtype,
+ )
+ observed.weight = float_linear.weight
+ observed.bias = float_linear.bias
+ return observed
+
+
+class QuantizedLinearQDQ(torch.nn.Module):
+ """Linear with explicit FP8 Q/DQ on activations and a pre-quantized FP8 weight."""
+
+ def __init__(
+ self,
+ act_obs,
+ weight_obs,
+ weight: torch.Tensor,
+ bias: torch.Tensor | None,
+ target_dtype: torch.dtype,
+ ):
+ super().__init__()
+ assert target_dtype == torch.float8_e4m3fn
+ self.act_scale, _ = act_obs.calculate_qparams()
+ weight_scale, _ = weight_obs.calculate_qparams()
+ self.target_dtype = target_dtype
+ self.bias = bias
+ self.output_dtype = weight.dtype
+ weight_scale_2d = (
+ weight_scale.view(-1, 1) if weight_scale.dim() == 1 else weight_scale
+ )
+ self.register_buffer(
+ "weight_fp8",
+ _quantize_affine_float8_non_decomposed(
+ weight, weight_scale_2d, target_dtype
+ ),
+ )
+ self.register_buffer("weight_scale", weight_scale_2d)
+
+ def forward(self, input: Tensor):
+ input_fp8 = _quantize_affine_float8_non_decomposed(
+ input, self.act_scale, self.target_dtype
+ )
+ input_hp = _dequantize_affine_float8_non_decomposed(
+ input_fp8, self.act_scale, self.output_dtype
+ )
+ weight_hp = _dequantize_affine_float8_non_decomposed(
+ self.weight_fp8, self.weight_scale, self.output_dtype
+ )
+ return F.linear(input_hp, weight_hp, self.bias)
+
+ @classmethod
+ def from_observed(cls, observed_linear, target_dtype):
+ return cls(
+ observed_linear.act_obs,
+ observed_linear.weight_obs,
+ observed_linear.weight,
+ observed_linear.bias,
+ target_dtype,
+ )
+
+
+@dataclass
+class StaticQuantConfigQDQ(AOBaseConfig):
+ target_dtype: torch.dtype
+
+
+@register_quantize_module_handler(StaticQuantConfigQDQ)
+def _apply_static_quant_qdq_transform(module, config):
+ return QuantizedLinearQDQ.from_observed(module, config.target_dtype)
+
+
+def insert_observers_(model, act_obs, weight_obs):
+ def replacement_fn(m):
+ return ObservedLinear.from_float(
+ m, copy.deepcopy(act_obs), copy.deepcopy(weight_obs)
+ )
+
+ _replace_with_custom_fn_if_matches_filter(
+ model,
+ replacement_fn,
+ lambda m, fqn: isinstance(m, torch.nn.Linear),
+ )
+
+
+def create_fp8_observers():
+ common_kwargs = dict(
+ mapping_type=MappingType.SYMMETRIC,
+ target_dtype=torch.float8_e4m3fn,
+ eps=torch.finfo(torch.float32).eps,
+ scale_dtype=torch.float32,
+ zero_point_dtype=torch.float32,
+ )
+ act_obs = AffineQuantizedMinMaxObserver(granularity=PerTensor(), **common_kwargs)
+ weight_obs = AffineQuantizedMinMaxObserver(
+ granularity=PerAxis(axis=0), **common_kwargs
+ )
+ return act_obs, weight_obs
+
+
+def quantize_static_fp8(
+ model: torch.nn.Module,
+ example_inputs,
+ calibration_steps: int = 10,
+) -> torch.nn.Module:
+ """Calibrate Linear layers and replace them with explicit FP8 Q/DQ modules."""
+ act_obs, weight_obs = create_fp8_observers()
+ insert_observers_(model, act_obs, weight_obs)
+ with torch.no_grad():
+ for _ in range(calibration_steps):
+ model(*example_inputs)
+ quantize_(
+ model,
+ StaticQuantConfigQDQ(torch.float8_e4m3fn),
+ lambda m, fqn: isinstance(m, ObservedLinear),
+ )
+ return model
diff --git a/examples/dynamo/torchao/torch_export_flux_fp8_woq.py b/examples/dynamo/torchao/torch_export_flux_fp8_woq.py
new file mode 100644
index 0000000000..092a1d56d3
--- /dev/null
+++ b/examples/dynamo/torchao/torch_export_flux_fp8_woq.py
@@ -0,0 +1,161 @@
+"""
+.. _torch_export_flux_fp8_woq:
+
+Compiling FLUX.1-dev with TorchAO FP8 weight-only quantization
+==============================================================
+
+This example quantizes the ``transformer`` of
+`FLUX.1-dev `_ with TorchAO
+``Float8WeightOnlyConfig``, then compiles it with the Torch-TensorRT Dynamo
+backend.
+
+Weight-only FP8 keeps activations in BF16 and stores Linear weights as FP8 plus
+per-channel scales. After export, Torch-TensorRT maps ``dequantize_affine`` to
+TensorRT ``IDequantizeLayer`` so the engine can keep an FP8 weight constant
+instead of folding DQ into a dense BF16 GEMM.
+
+You need access to FLUX.1-dev on Hugging Face and a GPU with enough memory to
+load and compile the 12B transformer (the unquantized Flux compile path wants
+>80GB; FP8 WOQ reduces the weight footprint).
+
+.. code-block:: bash
+
+ pip install torchao diffusers transformers accelerate sentencepiece protobuf
+
+On Blackwell, set Myelin prologue-fusion flags to encourage DQ+GEMM fusion::
+
+ export __LUNOWUD='-log:level=1 -log:dump=on -trace:use_id=on -mlir:prologue_fusion=1 -mlir:fusion_profit_threshold=0.01'
+
+"""
+
+# %%
+# 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 gc
+import os
+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
+from diffusers import FluxPipeline
+from torchao.quantization import Float8WeightOnlyConfig, quantize_
+
+sys.path.insert(0, _EXAMPLE_DIR)
+from utils import exclude_dq_from_constant_folding, pre_process_model_for_export
+
+DEVICE = "cuda:0"
+MODEL_ID = os.environ.get("FLUX_MODEL_ID", "black-forest-labs/FLUX.1-dev")
+
+# %%
+# Load FLUX.1-dev and quantize the transformer
+# --------------------------------------------
+# Only the transformer is quantized and compiled. Text encoders and the VAE stay
+# BF16. ``Float8WeightOnlyConfig`` is weight-only, so no calibration dataset is
+# required.
+
+pipe = FluxPipeline.from_pretrained(
+ MODEL_ID,
+ torch_dtype=torch.bfloat16,
+)
+config = pipe.transformer.config
+backbone = pipe.transformer.to(DEVICE)
+
+quantize_(backbone, Float8WeightOnlyConfig())
+backbone = pre_process_model_for_export(backbone)
+
+# %%
+# Export the quantized transformer
+# --------------------------------
+# Dummy inputs match the Flux transformer signature used by the BF16 Flux
+# example. ``exclude_dq_from_constant_folding`` keeps ``dequantize_affine`` in
+# the graph during ``torch.export``.
+
+dummy_inputs = {
+ "hidden_states": torch.randn(1, 4096, 64, dtype=torch.bfloat16, device=DEVICE),
+ "encoder_hidden_states": torch.randn(
+ 1, 512, 4096, dtype=torch.bfloat16, device=DEVICE
+ ),
+ "pooled_projections": torch.randn(1, 768, dtype=torch.bfloat16, device=DEVICE),
+ "timestep": torch.randn(1, dtype=torch.bfloat16, device=DEVICE),
+ "guidance": torch.randn(1, dtype=torch.float32, device=DEVICE),
+ "img_ids": torch.randn(4096, 3, dtype=torch.bfloat16, device=DEVICE),
+ "txt_ids": torch.randn(512, 3, dtype=torch.bfloat16, device=DEVICE),
+ "joint_attention_kwargs": {},
+ "return_dict": False,
+}
+
+with exclude_dq_from_constant_folding():
+ exp_program = torch.export.export(
+ backbone,
+ args=(),
+ kwargs=dummy_inputs,
+ strict=True,
+ )
+
+# %%
+# Compile with Torch-TensorRT
+# ---------------------------
+# .. note::
+# Compilation of the 12B transformer takes on the order of 20–30 minutes on
+# an H100. ``offload_module_to_cpu`` frees PyTorch weights after they are
+# ingested by TensorRT.
+
+trt_gm = torch_tensorrt.dynamo.compile(
+ exp_program,
+ inputs=dummy_inputs,
+ truncate_double=True,
+ min_block_size=1,
+ use_explicit_typing=True,
+ require_full_compilation=True,
+ immutable_weights=False,
+ offload_module_to_cpu=True,
+)
+
+# %%
+# Swap the compiled transformer into the pipeline
+# -----------------------------------------------
+
+pipe.transformer = None
+pipe.to(DEVICE)
+pipe.transformer = trt_gm
+pipe.transformer.config = config
+trt_gm.device = torch.device("cuda")
+del exp_program, backbone
+gc.collect()
+torch.cuda.empty_cache()
+
+# %%
+# Generate an image
+# -----------------
+
+
+def generate_image(pipe, prompt, image_name):
+ seed = 42
+ with torch.no_grad():
+ image = pipe(
+ prompt,
+ output_type="pil",
+ num_inference_steps=20,
+ generator=torch.Generator("cuda").manual_seed(seed),
+ ).images[0]
+ image.save(f"{image_name}.png")
+ print(f"Image generated using {image_name} model saved as {image_name}.png")
+
+
+generate_image(
+ pipe,
+ [
+ "Baroque style, a lavish palace interior with ornate gilded ceilings, "
+ "intricate tapestries, and dramatic lighting over a grand staircase."
+ ],
+ "flux_fp8_woq",
+)
diff --git a/examples/dynamo/torchao/utils.py b/examples/dynamo/torchao/utils.py
new file mode 100644
index 0000000000..6d99e1b25e
--- /dev/null
+++ b/examples/dynamo/torchao/utils.py
@@ -0,0 +1,59 @@
+"""Helpers for TorchAO FP8 weight-only quantization examples.
+
+TorchAO's default ``Float8Tensor.dequantize`` decomposes into primitive ops.
+``Float8TensorNonDecomposed`` keeps an explicit ``dequantize_affine`` in the
+exported graph so Torch-TensorRT can map it to ``IDequantizeLayer``.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator
+
+import torch
+from torchao.quantization import dequantize_affine
+from torchao.quantization.quantize_.workflows import Float8Tensor
+
+
+class Float8TensorNonDecomposed(Float8Tensor):
+ """``Float8Tensor`` that dequantizes via explicit ``dequantize_affine``."""
+
+ def dequantize(self, output_dtype=None):
+ if output_dtype is None:
+ output_dtype = torch.bfloat16
+ return dequantize_affine(
+ self.qdata,
+ self.block_size,
+ self.scale,
+ None,
+ self.qdata.dtype,
+ output_dtype=output_dtype,
+ )
+
+
+def pre_process_model_for_export(model: torch.nn.Module) -> torch.nn.Module:
+ """Promote ``Float8Tensor`` parameters so export emits ``dequantize_affine``."""
+ for param in model.parameters():
+ if isinstance(param, Float8Tensor) and not isinstance(
+ param, Float8TensorNonDecomposed
+ ):
+ param.__class__ = Float8TensorNonDecomposed
+ param.requires_grad_(False)
+ return model
+
+
+@contextmanager
+def exclude_dq_from_constant_folding() -> Iterator[None]:
+ """Keep ``dequantize_affine`` out of inductor constant folding during export."""
+ from torch._inductor.constant_folding import (
+ _dont_constant_fold,
+ add_dont_constant_fold,
+ )
+
+ op = torch.ops.torchao.dequantize_affine.default
+ add_dont_constant_fold(op)
+ try:
+ yield
+ finally:
+ if op in _dont_constant_fold:
+ _dont_constant_fold.remove(op)
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index b23205398a..598780d2d2 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -1030,6 +1030,108 @@ def aten_ops_dynamic_block_quantize_op(
)
+try:
+ import torchao # noqa: F401
+
+ assert torch.ops.torchao.dequantize_affine.default
+except Exception:
+ _LOGGER.debug(
+ "torchao not available; skipping torchao.dequantize_affine converter. "
+ "Install torchao to compile TorchAO weight-only quantized models."
+ )
+else:
+
+ @dynamo_tensorrt_converter(
+ torch.ops.torchao.dequantize_affine.default,
+ supports_dynamic_shapes=False,
+ )
+ def aten_ops_torchao_dequantize_affine(
+ ctx: ConversionContext,
+ target: Target,
+ args: Tuple[Argument, ...],
+ kwargs: Dict[str, Argument],
+ name: str,
+ ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+ # dequantize_affine(input, block_size, scale, zero_point, input_dtype,
+ # quant_min=None, quant_max=None, output_dtype=...)
+ output_dtype = kwargs.get("output_dtype")
+ if output_dtype is None:
+ output_dtype = args[7] if len(args) > 7 else torch.float16
+ input_dtype = args[4] if len(args) > 4 else None
+ return impl.quantize.dequantize_affine(
+ ctx,
+ target,
+ SourceIR.ATEN,
+ name,
+ args[0],
+ args[1],
+ args[2],
+ output_dtype,
+ input_dtype=input_dtype,
+ )
+
+
+try:
+ from torchao.quantization.quant_primitives import ( # noqa: F401
+ _dequantize_affine_float8_non_decomposed,
+ _quantize_affine_float8_non_decomposed,
+ )
+
+ assert torch.ops.torchao.quantize_affine_float8_non_decomposed.default
+ assert torch.ops.torchao.dequantize_affine_float8_non_decomposed.default
+except Exception:
+ _LOGGER.debug(
+ "torchao float8_non_decomposed ops not available; skipping static FP8 "
+ "Q/DQ converters. Import torchao quantization primitives to compile "
+ "TorchAO static FP8 graphs."
+ )
+else:
+
+ @dynamo_tensorrt_converter(
+ torch.ops.torchao.quantize_affine_float8_non_decomposed.default,
+ supports_dynamic_shapes=True,
+ )
+ def aten_ops_torchao_quantize_affine_float8(
+ ctx: ConversionContext,
+ target: Target,
+ args: Tuple[Argument, ...],
+ kwargs: Dict[str, Argument],
+ name: str,
+ ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+ return impl.quantize.quantize_affine_float8(
+ ctx,
+ target,
+ SourceIR.ATEN,
+ name,
+ args[0],
+ args[1],
+ )
+
+ @dynamo_tensorrt_converter(
+ torch.ops.torchao.dequantize_affine_float8_non_decomposed.default,
+ supports_dynamic_shapes=True,
+ )
+ def aten_ops_torchao_dequantize_affine_float8(
+ ctx: ConversionContext,
+ target: Target,
+ args: Tuple[Argument, ...],
+ kwargs: Dict[str, Argument],
+ name: str,
+ ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+ output_dtype = (
+ args[2] if len(args) > 2 else kwargs.get("output_dtype", torch.bfloat16)
+ )
+ return impl.quantize.dequantize_affine_float8(
+ ctx,
+ target,
+ SourceIR.ATEN,
+ name,
+ args[0],
+ args[1],
+ output_dtype,
+ )
+
+
@dynamo_tensorrt_converter(torch.ops.aten.squeeze.dim, supports_dynamic_shapes=True)
@dynamo_tensorrt_converter(torch.ops.aten.squeeze.dims, supports_dynamic_shapes=True)
def aten_ops_squeeze(
@@ -1227,7 +1329,7 @@ def _index_copy_kv_eligible(
if len(node.args) < 4:
return False
if input_node is None:
- input_node = node.args[0] # type: ignore[assignment]
+ input_node = node.args[0]
dim, _index_node, src_node = node.args[1:4]
if not isinstance(input_node, Node) or input_node.op != "placeholder":
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/quantize.py b/py/torch_tensorrt/dynamo/conversion/impl/quantize.py
index 8ec0895832..0258f0a099 100644
--- a/py/torch_tensorrt/dynamo/conversion/impl/quantize.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/quantize.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Sequence, Union
import numpy as np
import tensorrt as trt
@@ -6,6 +6,7 @@
from tensorrt import ITensor as TRTTensor
from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
from torch.fx.node import Target
+from torch_tensorrt import _enums
from torch_tensorrt.dynamo._SourceIR import SourceIR
from torch_tensorrt.dynamo.conversion import impl
from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
@@ -138,3 +139,125 @@ def quantize(
dq_output = dequantize_layer.get_output(0)
return dq_output
+
+
+def _block_size_as_ints(block_size: Sequence[object]) -> list[int]:
+ dims: list[int] = []
+ for bs in block_size:
+ if isinstance(bs, torch.Tensor):
+ dims.append(int(bs.item()))
+ elif isinstance(bs, (int, float)):
+ dims.append(int(bs))
+ else:
+ raise TypeError(f"Unsupported block_size dim type {type(bs)}: {bs!r}")
+ return dims
+
+
+def dequantize_affine(
+ ctx: ConversionContext,
+ target: Target,
+ source_ir: Optional[SourceIR],
+ name: str,
+ qdata: Union[torch.Tensor, TRTTensor],
+ block_size: Sequence[object],
+ scale: Union[np.ndarray, torch.Tensor, TRTTensor],
+ output_dtype: torch.dtype,
+ input_dtype: Optional[torch.dtype] = None,
+) -> TRTTensor:
+ """Map ``torchao.dequantize_affine`` to TensorRT ``IDequantizeLayer``.
+
+ Used by TorchAO weight-only quantization (FP8/INT8/INT4). The quantized
+ weight stays a low-precision constant and is dequantized at the GEMM
+ boundary so Myelin can fuse DQ into the matmul prologue instead of
+ folding it into a dense high-precision weight.
+ """
+ qdata_trt = get_trt_tensor(ctx, qdata, f"{name}_qdata", dtype=input_dtype)
+
+ axis = None
+ if isinstance(scale, torch.Tensor):
+ scale_for_trt = scale.squeeze()
+ if scale_for_trt.numel() != 1:
+ # Per-channel axis is the dimension whose block size is 1
+ # (quantized independently per slice). Example: weight (3072, 64)
+ # with block_size [3072, 1] → axis 1.
+ bs = _block_size_as_ints(block_size)
+ try:
+ axis = next(i for i, dim in enumerate(bs) if dim == 1)
+ except StopIteration as exc:
+ raise ValueError(
+ f"Unable to derive IDequantizeLayer axis from block_size={bs} "
+ f"and scale shape {tuple(scale.shape)}"
+ ) from exc
+ else:
+ scale_for_trt = scale
+
+ scale_trt = get_trt_tensor(ctx, scale_for_trt, f"{name}_scale", dtype=torch.float32)
+ trt_output_dtype = _enums.dtype._from(output_dtype).to(trt.DataType)
+
+ dequantize_layer = ctx.net.add_dequantize(
+ qdata_trt,
+ scale_trt,
+ output_type=trt_output_dtype,
+ )
+ if axis is not None:
+ dequantize_layer.axis = axis
+ set_layer_name(dequantize_layer, target, f"{name}_dequantize", source_ir)
+ return dequantize_layer.get_output(0)
+
+
+def _fp8_scale_and_axis(
+ ctx: ConversionContext,
+ scale: Union[np.ndarray, torch.Tensor, TRTTensor],
+ name: str,
+) -> tuple[TRTTensor, Optional[int]]:
+ axis = None
+ if isinstance(scale, torch.Tensor):
+ scale_for_trt = scale.squeeze()
+ if scale_for_trt.numel() != 1:
+ axis = 0
+ else:
+ scale_for_trt = scale
+ scale_trt = get_trt_tensor(ctx, scale_for_trt, f"{name}_scale", dtype=torch.float32)
+ return scale_trt, axis
+
+
+def quantize_affine_float8(
+ ctx: ConversionContext,
+ target: Target,
+ source_ir: Optional[SourceIR],
+ name: str,
+ input_tensor: Union[torch.Tensor, TRTTensor],
+ scale: Union[np.ndarray, torch.Tensor, TRTTensor],
+) -> TRTTensor:
+ """Map TorchAO ``quantize_affine_float8_non_decomposed`` to ``IQuantizeLayer``."""
+ input_trt = get_trt_tensor(ctx, input_tensor, f"{name}_input")
+ scale_trt, axis = _fp8_scale_and_axis(ctx, scale, name)
+ quantize_layer = ctx.net.add_quantize(input_trt, scale_trt, trt.DataType.FP8)
+ if axis is not None:
+ quantize_layer.axis = axis
+ set_layer_name(quantize_layer, target, f"{name}_quantize", source_ir)
+ return quantize_layer.get_output(0)
+
+
+def dequantize_affine_float8(
+ ctx: ConversionContext,
+ target: Target,
+ source_ir: Optional[SourceIR],
+ name: str,
+ input_tensor: Union[torch.Tensor, TRTTensor],
+ scale: Union[np.ndarray, torch.Tensor, TRTTensor],
+ output_dtype: torch.dtype = torch.bfloat16,
+) -> TRTTensor:
+ """Map TorchAO ``dequantize_affine_float8_non_decomposed`` to ``IDequantizeLayer``."""
+ input_trt = get_trt_tensor(ctx, input_tensor, f"{name}_input")
+ scale_trt, axis = _fp8_scale_and_axis(ctx, scale, name)
+ trt_output_dtype = _enums.dtype._from(output_dtype).to(trt.DataType)
+ dequantize_layer = ctx.net.add_dequantize(
+ input_trt,
+ scale_trt,
+ output_type=trt_output_dtype,
+ )
+ if axis is not None:
+ dequantize_layer.axis = axis
+ set_layer_name(dequantize_layer, target, f"{name}_dequantize", source_ir)
+ return dequantize_layer.get_output(0)
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
index a88aed4a9c..f7d999144c 100644
--- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
+++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py
@@ -107,8 +107,8 @@ def replace_node_with_constant(
class _TorchTensorRTConstantFolder(ConstantFolder): # type: ignore[misc]
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
- # Set of known quantization ops to be excluded from constant folding.
- # Currently, we exclude all quantization ops coming from modelopt library.
+ # Quantization ops excluded from constant folding so TRT sees QDQ
+ # (ModelOpt tensorrt.quantize_op and TorchAO dequantize_affine).
self.quantization_ops: Set[torch._ops.OpOverload] = set()
try:
# modelopt import ensures torch.ops.tensorrt.quantize_op.default is registered
@@ -123,6 +123,29 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
except Exception as e:
pass
+ try:
+ import torchao # noqa: F401
+
+ assert torch.ops.torchao.dequantize_affine.default
+ self.quantization_ops.add(torch.ops.torchao.dequantize_affine.default)
+ except Exception:
+ pass
+
+ try:
+ from torchao.quantization.quant_primitives import ( # noqa: F401
+ _dequantize_affine_float8_non_decomposed,
+ _quantize_affine_float8_non_decomposed,
+ )
+
+ self.quantization_ops.add(
+ torch.ops.torchao.quantize_affine_float8_non_decomposed.default
+ )
+ self.quantization_ops.add(
+ torch.ops.torchao.dequantize_affine_float8_non_decomposed.default
+ )
+ except Exception:
+ pass
+
# TODO: Update this function when quantization is added
def is_impure(self, node: torch.fx.node.Node) -> bool:
diff --git a/tests/py/dynamo/models/test_torchao_fp8_woq.py b/tests/py/dynamo/models/test_torchao_fp8_woq.py
new file mode 100644
index 0000000000..989bfc98cc
--- /dev/null
+++ b/tests/py/dynamo/models/test_torchao_fp8_woq.py
@@ -0,0 +1,172 @@
+# type: ignore
+import importlib
+import unittest
+
+import pytest
+import torch
+import torch_tensorrt as torchtrt
+from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
+
+assertions = unittest.TestCase()
+
+
+def _has_fp8_gpu() -> bool:
+ if not torch.cuda.is_available():
+ return False
+ major, _ = torch.cuda.get_device_capability()
+ return major >= 9
+
+
+def _promote_float8_woq(model: torch.nn.Module) -> torch.nn.Module:
+ from torchao.quantization import dequantize_affine
+ from torchao.quantization.quantize_.workflows import Float8Tensor
+
+ class Float8TensorNonDecomposed(Float8Tensor):
+ def dequantize(self, output_dtype=None):
+ if output_dtype is None:
+ output_dtype = torch.bfloat16
+ return dequantize_affine(
+ self.qdata,
+ self.block_size,
+ self.scale,
+ None,
+ self.qdata.dtype,
+ output_dtype=output_dtype,
+ )
+
+ for param in model.parameters():
+ if isinstance(param, Float8Tensor):
+ param.__class__ = Float8TensorNonDecomposed
+ param.requires_grad_(False)
+ return model
+
+
+@pytest.mark.unit
+@unittest.skipIf(importlib.util.find_spec("torchao") is None, "torchao not installed")
+@unittest.skipIf(not _has_fp8_gpu(), "FP8 GPU (compute capability >= 9.0) is required")
+def test_linear_fp8_woq():
+ from torch._inductor.constant_folding import (
+ _dont_constant_fold,
+ add_dont_constant_fold,
+ )
+ from torchao.quantization import Float8WeightOnlyConfig, quantize_
+
+ class LinearModel(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.linear = torch.nn.Linear(64, 128)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.linear(x)
+
+ model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda")
+ example_input = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda")
+ quantize_(model, Float8WeightOnlyConfig())
+ model = _promote_float8_woq(model)
+
+ op = torch.ops.torchao.dequantize_affine.default
+ add_dont_constant_fold(op)
+ try:
+ exp_program = torch.export.export(model, (example_input,), strict=True)
+ finally:
+ if op in _dont_constant_fold:
+ _dont_constant_fold.remove(op)
+
+ dq_nodes = [
+ n
+ for n in exp_program.graph.nodes
+ if n.target == torch.ops.torchao.dequantize_affine.default
+ ]
+ assertions.assertTrue(
+ len(dq_nodes) >= 1,
+ msg="Exported graph is missing torchao.dequantize_affine",
+ )
+
+ trt_mod = torchtrt.dynamo.compile(
+ exp_program,
+ inputs=[example_input],
+ min_block_size=1,
+ use_explicit_typing=True,
+ require_full_compilation=True,
+ cache_built_engines=False,
+ reuse_cached_engines=False,
+ )
+
+ with torch.no_grad():
+ eager_out = model(example_input)
+ trt_out = trt_mod(example_input)
+ if isinstance(trt_out, (list, tuple)):
+ trt_out = trt_out[0]
+ cos_sim = cosine_similarity(eager_out, trt_out)
+ assertions.assertTrue(
+ cos_sim > COSINE_THRESHOLD,
+ msg=(
+ "TorchAO FP8 WOQ TRT outputs don't match eager. "
+ f"Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}"
+ ),
+ )
+ torch._dynamo.reset()
+
+
+@pytest.mark.unit
+@unittest.skipIf(importlib.util.find_spec("torchao") is None, "torchao not installed")
+@unittest.skipIf(not _has_fp8_gpu(), "FP8 GPU (compute capability >= 9.0) is required")
+def test_linear_fp8_static():
+ import sys
+ from pathlib import Path
+
+ example_dir = (
+ Path(__file__).resolve().parents[4] / "examples" / "dynamo" / "torchao"
+ )
+ sys.path.insert(0, str(example_dir))
+ from static_fp8_utils import quantize_static_fp8
+
+ class LinearModel(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.linear1 = torch.nn.Linear(32, 64, bias=False)
+ self.linear2 = torch.nn.Linear(64, 16, 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(4, 32, dtype=torch.bfloat16, device="cuda")
+ quantize_static_fp8(model, (example_input,), calibration_steps=3)
+
+ exp_program = torch.export.export(model, (example_input,), strict=True)
+ graph_str = str(exp_program.graph)
+ assertions.assertIn(
+ "quantize_affine_float8_non_decomposed",
+ graph_str,
+ "Exported graph is missing static FP8 quantize",
+ )
+ assertions.assertIn(
+ "dequantize_affine_float8_non_decomposed",
+ graph_str,
+ "Exported graph is missing static FP8 dequantize",
+ )
+
+ trt_mod = torchtrt.dynamo.compile(
+ exp_program,
+ inputs=[example_input],
+ enabled_precisions={torch.float8_e4m3fn},
+ min_block_size=1,
+ require_full_compilation=True,
+ cache_built_engines=False,
+ reuse_cached_engines=False,
+ )
+ with torch.no_grad():
+ eager_out = model(example_input)
+ trt_out = trt_mod(example_input)
+ if isinstance(trt_out, (list, tuple)):
+ trt_out = trt_out[0]
+ cos_sim = cosine_similarity(eager_out, trt_out)
+ assertions.assertTrue(
+ cos_sim > COSINE_THRESHOLD,
+ msg=(
+ "TorchAO static FP8 TRT outputs don't match eager. "
+ f"Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}"
+ ),
+ )
+ torch._dynamo.reset()