Test Only fp4: Lluo/fp4 try out

.pre-commit-config.yaml

@@ -21,14 +21,14 @@ repos:
           - id: clang-format
             types_or: [c++, c, cuda]
     - repo: https://github.com/keith/pre-commit-buildifier
-      rev: 6.4.0
+      rev: 8.0.3
       hooks:
           - id: buildifier
             args:
                 - --warnings=all
           - id: buildifier-lint
     - repo: https://github.com/abravalheri/validate-pyproject
-      rev: v0.23
+      rev: v0.24.1
       hooks:
           - id: validate-pyproject
     - repo: https://github.com/pycqa/isort
@@ -37,17 +37,17 @@ repos:
           - id: isort
             name: isort (python)
     - repo: https://github.com/pre-commit/mirrors-mypy
-      rev: "v1.9.0"
+      rev: "v1.15.0"
       hooks:
           - id: mypy
             exclude: "^py/torch_tensorrt/fx|^examples|^tests|^py/torch_tensorrt/dynamo/_experimental|^tools|^docs|noxfile.py|setup.py|versions.py"
     - repo: https://github.com/astral-sh/ruff-pre-commit
       # Ruff version.
-      rev: v0.3.3
+      rev: v0.11.7
       hooks:
           - id: ruff
     - repo: https://github.com/psf/black
-      rev: 24.3.0
+      rev: 25.1.0
       hooks:
           - id: black
             exclude: ^examples/custom_converters/elu_converter/setup.py|^docs
@@ -57,7 +57,7 @@ repos:
           - id: typos
     - repo: https://github.com/astral-sh/uv-pre-commit
       # uv version.
-      rev: 0.5.5
+      rev: 0.7.1
       hooks:
           # Update the uv lockfile
           - id: uv-lock

MODULE.bazel

@@ -94,9 +94,9 @@ http_archive(
 http_archive(
     name = "tensorrt",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    strip_prefix = "TensorRT-10.9.0.34",
+    strip_prefix = "TensorRT-10.10.0.31",
     urls = [
-        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz",
+        "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.10.0/tars/TensorRT-10.10.0.31.Linux.x86_64-gnu.cuda-12.9.tar.gz",
     ],
 )
 

examples/dynamo/vgg16_ptq.py

@@ -169,7 +169,6 @@ def vgg16(num_classes=1000, init_weights=False):
 
 data = iter(training_dataloader)
 images, _ = next(data)
-
 crit = nn.CrossEntropyLoss()
 
 # %%
@@ -200,8 +199,11 @@ def calibrate_loop(model):
     quant_cfg = mtq.INT8_DEFAULT_CFG
 elif args.quantize_type == "fp8":
     quant_cfg = mtq.FP8_DEFAULT_CFG
+elif args.quantize_type == "fp4":
+    quant_cfg = mtq.NVFP4_DEFAULT_CFG
 # PTQ with in-place replacement to quantized modules
 mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
+
 # model has FP8 qdq nodes at this point
 
 # %%
@@ -233,12 +235,20 @@ def calibrate_loop(model):
     with export_torch_mode():
         # Compile the model with Torch-TensorRT Dynamo backend
         input_tensor = images.cuda()
+        torch.onnx.export(model, input_tensor, "mtq_vgg16_model.onnx")
 
         exp_program = torch.export.export(model, (input_tensor,), strict=False)
         if args.quantize_type == "int8":
             enabled_precisions = {torch.int8}
         elif args.quantize_type == "fp8":
             enabled_precisions = {torch.float8_e4m3fn}
+        elif args.quantize_type == "fp4":
+            enabled_precisions = {
+                torch.float4_e2m1fn_x2,
+                torch.float8_e4m3fn,
+                torch.float16,
+                torch.float32,
+            }
         trt_model = torchtrt.dynamo.compile(
             exp_program,
             inputs=[input_tensor],

py/torch_tensorrt/_enums.py

@@ -80,6 +80,12 @@ class dtype(Enum):
     :meta hide-value:
     """
 
+    f4 = auto()
+    """4 bit floating-point number, equivalent to ``dtype.fp4`` and ``dtype.float4``
+
+    :meta hide-value:
+    """
+
     uint8 = u8
     int8 = i8
 
@@ -91,6 +97,9 @@ class dtype(Enum):
     float8 = f8
     fp8 = f8
 
+    float4 = f4
+    fp4 = f4
+
     half = f16
     fp16 = f16
     float16 = f16
@@ -162,6 +171,8 @@ def _from(
                 return dtype.i32
             elif t == torch.float8_e4m3fn:
                 return dtype.f8
+            elif t == torch.float4_e2m1fn_x2:
+                return dtype.f4
             elif t == torch.half:
                 return dtype.f16
             elif t == torch.float:
@@ -188,6 +199,8 @@ def _from(
                 return dtype.i8
             elif t == trt.DataType.FP8:
                 return dtype.f8
+            elif t == trt.DataType.FP4:
+                return dtype.fp4
             elif t == trt.DataType.INT32:
                 return dtype.i32
             elif t == trt.DataType.INT64:
@@ -357,6 +370,8 @@ def to(
                 return torch.long
             elif self == dtype.f8:
                 return torch.float8_e4m3fn
+            elif self == dtype.f4:
+                return torch.float4_e2m1fn_x2
             elif self == dtype.f16:
                 return torch.half
             elif self == dtype.f32:
@@ -394,6 +409,8 @@ def to(
                 return trt.DataType.BOOL
             elif self == dtype.bf16:
                 return trt.DataType.BF16
+            elif self == dtype.f4:
+                return trt.DataType.FP4
             elif use_default:
                 return trt.DataType.FLOAT
             else:
@@ -410,6 +427,8 @@ def to(
                 return np.int64
             elif self == dtype.f16:
                 return np.float16
+            elif self == dtype.f4:
+                return np.float4_e2m1fn_x2
             elif self == dtype.f32:
                 return np.float32
             elif self == dtype.f64:

py/torch_tensorrt/dynamo/_compiler.py

@@ -581,13 +581,13 @@ def compile(
             "\nThis feature is unimplemented in Torch-TRT Dynamo currently."
         )
 
-    if use_explicit_typing:
-        if len(enabled_precisions) != 1 or not any(
-            x in enabled_precisions for x in {torch.float32, dtype.f32}
-        ):
-            raise AssertionError(
-                f"When use_explicit_typing is enabled, only torch.float32 is allowed in the enabled_precisions but found {enabled_precisions}"
-            )
+    # if use_explicit_typing:
+        # if len(enabled_precisions) != 1 or not any(
+        #     x in enabled_precisions for x in {torch.float32, dtype.f32}
+        # ):
+        #     raise AssertionError(
+        #         f"When use_explicit_typing is enabled, only torch.float32 is allowed in the enabled_precisions but found {enabled_precisions}"
+        #     )
 
     if use_fp32_acc:
         logger.debug(

py/torch_tensorrt/dynamo/_defaults.py

@@ -29,7 +29,14 @@
 REQUIRE_FULL_COMPILATION = False
 DRYRUN = False
 HARDWARE_COMPATIBLE = False
-SUPPORTED_KERNEL_PRECISIONS = {dtype.f32, dtype.f16, dtype.bf16, dtype.i8, dtype.f8}
+SUPPORTED_KERNEL_PRECISIONS = {
+    dtype.f32,
+    dtype.f16,
+    dtype.bf16,
+    dtype.i8,
+    dtype.f8,
+    dtype.f4,
+}
 TIMING_CACHE_PATH = os.path.join(
     tempfile.gettempdir(), "torch_tensorrt_engine_cache", "timing_cache.bin"
 )

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -87,6 +87,7 @@ def __init__(
 
         flag = 0
         if compilation_settings.use_explicit_typing:
+            _LOGGER.info("Using strongly typed network definition")
             STRONGLY_TYPED = 1 << (int)(
                 trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED
             )
@@ -274,17 +275,19 @@ def _populate_trt_builder_config(
                 self.compilation_settings.dla_global_dram_size,
             )
 
-        if dtype.float16 in self.compilation_settings.enabled_precisions:
-            builder_config.set_flag(trt.BuilderFlag.FP16)
+        # if dtype.float16 in self.compilation_settings.enabled_precisions:
+        #     builder_config.set_flag(trt.BuilderFlag.FP16)
 
         if dtype.int8 in self.compilation_settings.enabled_precisions:
             builder_config.set_flag(trt.BuilderFlag.INT8)
 
-        if dtype.fp8 in self.compilation_settings.enabled_precisions:
-            builder_config.set_flag(trt.BuilderFlag.FP8)
+        # if dtype.fp8 in self.compilation_settings.enabled_precisions:
+        #     builder_config.set_flag(trt.BuilderFlag.FP8)
 
-        if dtype.bfloat16 in self.compilation_settings.enabled_precisions:
-            builder_config.set_flag(trt.BuilderFlag.BF16)
+        # if dtype.fp4 in self.compilation_settings.enabled_precisions:
+        #     builder_config.set_flag(trt.BuilderFlag.FP4)
+        # if dtype.bfloat16 in self.compilation_settings.enabled_precisions:
+        #     builder_config.set_flag(trt.BuilderFlag.BF16)
 
         if self.compilation_settings.sparse_weights:
             builder_config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
@@ -351,7 +354,7 @@ def _populate_trt_builder_config(
                 builder_config.l2_limit_for_tiling = (
                     self.compilation_settings.l2_limit_for_tiling
                 )
-
+        print(f"lan added builder_config:{builder_config=}")
         return builder_config
 
     def _create_timing_cache(

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -617,6 +617,39 @@ def aten_ops_quantize_op(
         )
 
 
+try:
+    import modelopt.torch.quantization as mtq  # noqa: F401
+
+    assert torch.ops.tensorrt.dynamic_block_quantize_op.default
+except Exception as e:
+    _LOGGER.warning(
+        "Unable to import dynamic block quantize op. Please install modelopt library (https://github.com/NVIDIA/TensorRT-Model-Optimizer?tab=readme-ov-file#installation) to add support for compiling dynamic blockquantized models"
+    )
+else:
+
+    @dynamo_tensorrt_converter(torch.ops.tensorrt.dynamic_block_quantize_op.default)
+    def aten_ops_dynamic_block_quantize_op(
+        ctx: ConversionContext,
+        target: Target,
+        args: Tuple[Argument, ...],
+        kwargs: Dict[str, Argument],
+        name: str,
+    ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+        return impl.nvfp4_quantize.nvfp4_quantize(
+            ctx,
+            target,
+            SourceIR.ATEN,
+            name,
+            args[0],
+            args[1],
+            args[2],
+            args[3],
+            args[4],
+            args[5],
+            args[6],
+        )
+
+
 @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(

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -361,12 +361,37 @@ def create_constant(
             shape = list(torch_value.shape)
 
         if torch_value is not None:
+            if torch_value.dtype == torch.float8_e4m3fn:
+                weights = trt.Weights(
+                    type=trt.DataType.FP8,
+                    ptr=torch_value.data_ptr(),
+                    count=torch_value.numel(),
+                )
+                constant = ctx.net.add_constant(
+                    shape,
+                    weights,
+                )
+                constant.name = name
+                return constant.get_output(0)
+            # Iconstant layer does not support Uint8, it only support that FP4 data packed in uint8
+            if torch_value.dtype == torch.uint8:
+                weights = trt.Weights(
+                    type=trt.DataType.FP4,
+                    ptr=torch_value.data_ptr(),
+                    count=torch_value.numel() * 2,
+                )
+                shape[-1] = shape[-1] * 2
+                constant = ctx.net.add_constant(
+                    shape,
+                    weights,
+                )
+                constant.name = name
+                return constant.get_output(0)
             if torch_value.dtype == torch.bfloat16:
                 torch_value_fp32 = torch_value.to(torch.float32)
                 numpy_value = torch_value_fp32.numpy()
             else:
                 numpy_value = torch_value.numpy()
-
             ctx.mapping[name + " CONSTANT"] = numpy_value.reshape(-1)
             constant = ctx.net.add_constant(
                 shape,
@@ -381,7 +406,6 @@ def create_constant(
                     trt.DataType.BF16,
                     name + "_bf16_cast",
                 )
-
             return constant.get_output(0)
         else:
             raise ValueError(

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -14,6 +14,7 @@
     matmul,
     nccl_ops,
     normalization,
+    nvfp4_quantize,
     pad,
     permutation,
     pool,

py/torch_tensorrt/dynamo/conversion/impl/addmm.py

@@ -7,7 +7,7 @@
 from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.fx.types import TRTTensor
-
+import os
 
 def addmm(
     ctx: ConversionContext,
@@ -21,6 +21,10 @@ def addmm(
     beta: Union[float, int],
     alpha: Union[float, int],
 ) -> TRTTensor:
+    if os.getenv("DISABLE_GEMM", "false").lower() == "true":
+        print("lan added disable_gemm is set, skip addmm and returning mat2")
+        return mat2
+    print("lan added disable_gemm is not set, doing addmm")
     mm = impl.matmul.matrix_multiply(ctx, target, source_ir, f"{name}_mm", mat1, mat2)
     if alpha != 1:
         mm = impl.elementwise.mul(