Do not use autocast for mxfp8

Signed-off-by: ajrasane <[email protected]>

Author: ajrasane

.github/workflows/gpu_tests.yml

@@ -62,7 +62,7 @@ jobs:
     runs-on: linux-amd64-gpu-l4-latest-1
     timeout-minutes: 120
     container: &gpu_container
-      image: nvcr.io/nvidia/pytorch:25.08-py3
+      image: nvcr.io/nvidia/pytorch:25.06-py3
       env:
         GIT_DEPTH: 1000 # For correct version for tests/gpu/torch/quantization/plugins/test_megatron.py
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages

.gitlab/tests.yml

@@ -68,7 +68,7 @@ example-trtllm:
 
 example-onnx:
   extends: example-torch
-  image: nvcr.io/nvidia/tensorrt:25.08-py3
+  image: nvcr.io/nvidia/tensorrt:25.06-py3
   tags: [docker, linux, 2-gpu, sm>=89]
   parallel:
     matrix:

modelopt/onnx/export/mxfp8_exporter.py

@@ -166,4 +166,34 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
                         attr.s = b"tanh"
                         logger.debug(f"Updated GELU node {node.name} to use tanh approximation")
 
+        # Insert cast to fp16 after Sqrt nodes
+        cast_nodes_to_insert = []
+        for idx, node in enumerate(graph.node):
+            if node.op_type == "Sqrt":
+                sqrt_output = node.output[0]
+                cast_output = f"{sqrt_output}_cast_fp16"
+
+                # Create Cast node
+                cast_node = onnx.helper.make_node(
+                    "Cast",
+                    inputs=[sqrt_output],
+                    outputs=[cast_output],
+                    to=onnx_dtype_map["Half"],
+                    name=f"{node.name}_cast_fp16",
+                )
+                cast_nodes_to_insert.append((idx + 1, cast_node))
+
+                # Update consumers to use cast output
+                for consumer in graph.node:
+                    if consumer == node:
+                        continue
+                    for i, inp in enumerate(consumer.input):
+                        if inp == sqrt_output:
+                            consumer.input[i] = cast_output
+
+        # Insert Cast nodes in reverse order to preserve indices
+        for offset, (pos, cast_node) in enumerate(cast_nodes_to_insert):
+            graph.node.insert(pos + offset, cast_node)
+            logger.debug(f"Inserted Cast to FP16 after {cast_node.input[0]}")
+
         return onnx_model

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -580,7 +580,7 @@ def get_onnx_bytes_and_metadata(
     except StopIteration:
         param_dtype = torch.float32
     if weights_dtype in ["fp16", "bf16"] and param_dtype == torch.float32:
-        if is_int4_quantized(model):
+        if is_int4_quantized(model) or is_mxfp8_quantized(model):
             assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet"
             onnx_opt_graph = convert_float_to_float16(
                 onnx_opt_graph,