From 78dd40918f8057a90cd10beb056de85c80d24b95 Mon Sep 17 00:00:00 2001
From: vishalpandya1990 <vishalpandya1990@gmail.com>
Date: Wed, 19 Nov 2025 11:29:21 +0530
Subject: [PATCH] Use ONNX DQ node instead of DQ custom-op for activation
 dequantization in nvfp4 (#536)

## What does this PR do?

**Type of change:** Minor update

**Overview:**

- We are currently using trt::DequantizeLinear custom op for
activation's dequantization path in NVFP4 model.
- With this change, substituting trt::DequantizeLinear custom op with
DequantizeLinear (ONNX node) since official ONNX DQ node should already
support it.

## Testing

_System details:_ Windows 11 22621, RTX 5090, TRT 10.10.0.31
_Model:_ SD3.5-Medium's FP4 transformer ONNX model (with quant-mha)
_Command:_
`trtexec --builderOptimizationLevel=4 --onnx=<.onnx file path>
--minShapes=hidden_states:2x16x64x64,timestep:2,encoder_hidden_states:2x77x4096,pooled_projections:2x2048
--optShapes=hidden_states:16x16x64x64,timestep:16,encoder_hidden_states:16x154x4096,pooled_projections:16x2048
--maxShapes=hidden_states:16x16x128x128,timestep:16,encoder_hidden_states:16x333x4096,pooled_projections:16x2048
--stronglyTyped`

_With change (i.e. with ONNX DQ node for activation):_ Throughput:
12.2501 qps, GPU Compute Time: min = 65.9408 ms, max = 71.7161 ms, mean
= 67.0051 ms, median = 66.4348 ms, percentile(90%) = 69.2656 ms,
percentile(95%) = 70.3727 ms, percentile(99%) = 71.7161 ms

_Without change (i.e. with custom DQ node for activation):_ Throughput:
12.0729 qps, GPU Compute Time: min = 66.5288 ms, max = 80.5145 ms, mean
= 68.2194 ms, median = 67.043 ms, percentile(90%) = 71.2693 ms,
percentile(95%) = 72.0937 ms, percentile(99%) = 80.5145 ms

Attached the trtexec log for reference (both with and without change)

[trtexec_log_fp4_custom_op_A_removal_with_without_change.txt](https://github.com/user-attachments/files/23475935/trtexec_log_fp4_custom_op_A_removal_with_without_change.txt)


## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->

Signed-off-by: vipandya <vipandya@nvidia.com>
---
 modelopt/torch/quantization/export_onnx.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py
index e8a38b162..d49d65f6e 100644
--- a/modelopt/torch/quantization/export_onnx.py
+++ b/modelopt/torch/quantization/export_onnx.py
@@ -557,7 +557,7 @@ def _fp4_dequantize(
             "Constant",
             value_t=torch.tensor(scale, dtype=torch_dtype_map["Float"]),
         )
-    return g.op("trt::DequantizeLinear", inputs, scale)
+    return g.op("DequantizeLinear", inputs, scale)
 
 
 def _fp4_dequantize_2(
@@ -568,7 +568,7 @@ def _fp4_dequantize_2(
     axis: int = -1,
 ):
     """Helper Function for Dequantization."""
-    return g.op("trt::DequantizeLinear", inputs, dyn_scale, axis_i=axis, block_size_i=block_size)
+    return g.op("DequantizeLinear", inputs, dyn_scale, axis_i=axis, block_size_i=block_size)
 
 
 def _mxfp8_dynamic_quantize(