[None][feat] add ReLU2 NVFP4 fusion for AutoDeploy with tests - #11957
Merged
tcherckez-nvidia merged 4 commits intoMar 8, 2026
Conversation
Add the ReLU2+NVFP4 fusion transform and required AutoDeploy wrapper ops, and cover the path with transform and custom-op unit tests. Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com> Made-with: Cursor
tcherckez-nvidia
force-pushed
the
relu2-nvfp4-fusion
branch
from
March 5, 2026 18:56
627dfb3 to
800a712
Compare
Contributor
📝 WalkthroughWalkthroughThis pull request introduces new NVFP4 quantization capabilities with ReLU2 fusion support. It adds a new transform configuration, implements fused quantization operations with fake variants, creates a graph transformation to detect and optimize ReLU2-quantization patterns, and provides comprehensive unit and integration tests to validate the functionality. Changes
Sequence DiagramsequenceDiagram
participant Graph as Graph Module
participant Detect as ReLU2 Pattern<br/>Detection
participant Fused as Fused Quant Op<br/>Insertion
participant PreQuant as PreQuant Linear Op<br/>Insertion
participant Output as Transformed<br/>Graph
Graph->>Detect: Iterate over<br/>torch_quant_nvfp4_linear nodes
Detect->>Detect: Check for ReLU²<br/>chain or call_module
alt Pattern Matched
Detect->>Fused: Extract ReLU input<br/>and metadata
Fused->>Fused: Create trtllm_fused_relu2_<br/>quant_nvfp4 node
Fused->>PreQuant: Extract fp4 and<br/>scale factors
PreQuant->>PreQuant: Create trtllm_nvfp4_<br/>prequant_linear node
PreQuant->>Output: Replace original node<br/>and erase obsolete paths
else Pattern Not Matched
Detect->>Output: Leave node unchanged
end
Output->>Output: Recompile GraphModule
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
Author
|
@coderabbitai help |
Contributor
ChatThere are 3 ways to chat with CodeRabbit:
CodeRabbit commands
Other keywords and placeholders
Status, support, documentation and community
|
Collaborator
Author
|
@coderabbitai review |
Contributor
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py (1)
65-133: Add coverage for thecall_module(ReLUSquaredActivation)match branch.Current tests validate the aten route and a non-match route, but not the
call_moduleroute that_is_relu2_chainexplicitly supports.Possible test addition
+@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_fuse_relu2_quant_nvfp4_matches_call_module_relu2(): + class ReLUSquaredActivation(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = torch.nn.functional.relu(x) + return y * y + + class CallModuleRelu2Model(nn.Module): + def __init__(self): + super().__init__() + self.impl = TinyRelu2NVFP4Linear() + self.act = ReLUSquaredActivation() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.act(x) + return torch.ops.auto_deploy.torch_quant_nvfp4_linear( + x, self.impl.weight_fp4, self.impl.bias, + self.impl.input_scale, self.impl.weight_scale, self.impl.alpha + ) + + model = CallModuleRelu2Model().to("cuda") + x = torch.rand(3, 64, dtype=torch.float16, device="cuda") + gm = torch_export_to_gm(model, args=(x,), clone=True) + gm_transformed = InferenceOptimizer(None, {"fuse_relu2_quant_nvfp4": {"stage": "post_load_fusion"}})(None, gm) + + assert _count_op(gm_transformed, torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4.default) == 1 + assert _count_op(gm_transformed, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear.default) == 1🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py` around lines 65 - 133, Tests currently miss the code path where the relu2 chain is expressed as a module (call_module(ReLUSquaredActivation)); add a new unit test in the same file that constructs a small nn.Module which uses an actual ReLUSquaredActivation submodule (instead of aten ops) before calling torch.ops.auto_deploy.torch_quant_nvfp4_linear, export it with torch_export_to_gm, run InferenceOptimizer with "fuse_relu2_quant_nvfp4" and assert that _count_op(gm_transformed, torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4.default) == 1 and that torch.ops.auto_deploy.torch_quant_nvfp4_linear and trtllm_nvfp4_prequant_linear counts are as expected (mirroring test_fuse_relu2_quant_nvfp4_rewrite_and_numerics); ensure you reference and exercise the _is_relu2_chain path by including the ReLUSquaredActivation submodule and use TinyRelu2NVFP4Linear fields (weight_fp4, bias, input_scale, weight_scale, alpha) like the existing tests.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.py (1)
322-347: Add non-2D test shapes to validate wrapper reshape semantics.Both new wrapper tests currently validate only 2D inputs, while production wrappers flatten/restore leading dimensions. Adding one 3D shape would harden regression coverage.
Example parameterization pattern
+@pytest.mark.parametrize("x_shape", [(8, 64), (2, 4, 64)]) def test_fused_relu2_quant_nvfp4_wrapper_matches_trtllm_op(): - x = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda") + x = torch.randn(*x_shape, dtype=torch.bfloat16, device="cuda") input_scale = fp4_global_scale(x).to(torch.float32) fp4_wrapped, sf_wrapped = torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4(x, input_scale) - fp4_ref, sf_ref = torch.ops.trtllm.fused_relu2_quantize(x, input_scale, 16) + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + fp4_ref, sf_ref = torch.ops.trtllm.fused_relu2_quantize(x_2d, input_scale, 16) + fp4_ref = fp4_ref.reshape(*x.shape[:-1], fp4_ref.shape[-1])Also applies to: 354-382
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.py` around lines 322 - 347, Update the test_fused_relu2_quant_nvfp4_wrapper_matches_trtllm_op test to also exercise a non-2D input (e.g. add a 3D shape like x.shape = (2, 8, 64)) so the wrapper's flatten/restore semantics are validated; for the new shape, call torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4 and the reference torch.ops.trtllm.fused_relu2_quantize the same way, ensure fp4_wrapped/ fp4_ref and sf_wrapped/sf_ref shapes and dtypes match after reshape, and run the same nvfp4_gemm functional equivalence check (using w, w_fp4, w_scale, alpha) to compare out_wrapped vs out_ref — keep the same assertions and tolerances but iterate or parametrize over the 2D and 3D input shapes so wrapper reshape behavior is covered.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fuse_relu2_quant_nvfp4.py`:
- Around line 71-83: The try/except blocks in fuse_relu2_quant_nvfp4.py
currently swallow all exceptions (the block around inspecting input_node.target
and the later block that logs but still catches Exception); change both to catch
only AttributeError, RuntimeError, and TypeError explicitly instead of a bare
except Exception, e.g., replace "except Exception:" with "except
(AttributeError, RuntimeError, TypeError):" for the try that reads
input_node.target/get_submodule/type(submod).__name__ (references:
input_node.target, gm.get_submodule, type(submod).__name__) and for the later
try that logs the error—keep the existing logging there but narrow the caught
exception types to the same tuple so unexpected errors still surface.
---
Nitpick comments:
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.py`:
- Around line 322-347: Update the
test_fused_relu2_quant_nvfp4_wrapper_matches_trtllm_op test to also exercise a
non-2D input (e.g. add a 3D shape like x.shape = (2, 8, 64)) so the wrapper's
flatten/restore semantics are validated; for the new shape, call
torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4 and the reference
torch.ops.trtllm.fused_relu2_quantize the same way, ensure fp4_wrapped/ fp4_ref
and sf_wrapped/sf_ref shapes and dtypes match after reshape, and run the same
nvfp4_gemm functional equivalence check (using w, w_fp4, w_scale, alpha) to
compare out_wrapped vs out_ref — keep the same assertions and tolerances but
iterate or parametrize over the 2D and 3D input shapes so wrapper reshape
behavior is covered.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py`:
- Around line 65-133: Tests currently miss the code path where the relu2 chain
is expressed as a module (call_module(ReLUSquaredActivation)); add a new unit
test in the same file that constructs a small nn.Module which uses an actual
ReLUSquaredActivation submodule (instead of aten ops) before calling
torch.ops.auto_deploy.torch_quant_nvfp4_linear, export it with
torch_export_to_gm, run InferenceOptimizer with "fuse_relu2_quant_nvfp4" and
assert that _count_op(gm_transformed,
torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4.default) == 1 and that
torch.ops.auto_deploy.torch_quant_nvfp4_linear and trtllm_nvfp4_prequant_linear
counts are as expected (mirroring
test_fuse_relu2_quant_nvfp4_rewrite_and_numerics); ensure you reference and
exercise the _is_relu2_chain path by including the ReLUSquaredActivation
submodule and use TinyRelu2NVFP4Linear fields (weight_fp4, bias, input_scale,
weight_scale, alpha) like the existing tests.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 5357e1ec-275a-497b-b6b4-79e42a424cad
📒 Files selected for processing (5)
tensorrt_llm/_torch/auto_deploy/config/default.yamltensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.pytensorrt_llm/_torch/auto_deploy/transform/library/fuse_relu2_quant_nvfp4.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
tensorrt_llm/_torch/auto_deploy/transform/library/fuse_relu2_quant_nvfp4.py (1)
71-83:⚠️ Potential issue | 🟠 MajorNarrow broad exception handlers and stop silently swallowing failures.
Line 81 and Line 174 catch
Exception, and the first block suppresses errors withpass. This can mask real graph-inspection failures and violates project exception-handling rules.Proposed fix
- except Exception: - pass + except (AttributeError, RuntimeError, TypeError) as exc: + ad_logger.debug( + "fuse_relu2_quant_nvfp4: failed call_module relu2 inspection for target=%s: %s", + getattr(input_node, "target", None), + exc, + ) return None @@ - except Exception as e: + except (AttributeError, RuntimeError, TypeError) as e: ad_logger.info( "fuse_relu2_quant_nvfp4: linear input is call_module(%s), " "get_submodule failed: %s", input_arg.target, e, )As per coding guidelines:
**/*.{py,cpp,h,hpp,cuh}: Avoid broad exception handling — catch specific exceptions, not bare except:and**/*.py: When using try-except blocks in Python, limit the except to the smallest set of errors possible.Also applies to: 166-180
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/transform/library/fuse_relu2_quant_nvfp4.py` around lines 71 - 83, The try/except is too broad and currently swallows all errors; replace the bare except with specific exceptions likely from attribute access or dict lookups (for example catch AttributeError, KeyError, and TypeError) and handle them explicitly (return None or re-raise after logging) instead of pass; update the block around reading input_node.target / gm.get_submodule(...) / type(submod) and the similar block at lines ~166-180 to catch only those specific exceptions (e.g., except (AttributeError, KeyError, TypeError) as e:) and either return None or log the error context (include target and type_name) before returning so real failures aren’t silently discarded.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py (1)
397-407: Minor lint cleanup: markbiasas intentionally unused in fake prequant op.Line 403 is currently unused; explicitly discarding it keeps fake-kernel intent clear and quiets ARG001 warnings.
Proposed fix
def trtllm_nvfp4_prequant_linear_fake( @@ ) -> torch.Tensor: - del input_sf, weight_scale, alpha + del input_sf, weight_scale, alpha, bias out_features = weight_fp4.shape[0]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py` around lines 397 - 407, The bias parameter in trtllm_nvfp4_prequant_linear_fake is intentionally unused but not explicitly discarded; to silence ARG001 and clarify intent, explicitly mark it as unused (e.g., add bias to the existing deletion of unused args in trtllm_nvfp4_prequant_linear_fake) so the function body reads as discarding input_sf, weight_scale, alpha, and bias.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py (1)
95-133: Strengthen the non-match test with a numerical equivalence assertion.It already validates no rewrite happened; adding
y_refvsy_newcomparison would guard accidental semantic drift in the non-fused path.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py` around lines 95 - 133, Add a numerical equivalence check comparing the model outputs before and after the transformation to ensure the non-fused path remains semantically identical: run the original scripted/graph model (from torch_export_to_gm or a forward call on model) to produce y_ref using input x, then run the transformed graph gm_transformed to produce y_new and assert they are numerically equal (e.g., torch.allclose or torch.testing.assert_close with appropriate tolerances). Place this check in test_fuse_relu2_quant_nvfp4_does_not_match_non_relu2 after gm_transformed is created and before the _count_op assertions, referencing NonRelu2Model, torch_quant_nvfp4_linear, torch_export_to_gm, and InferenceOptimizer so the intent and location are clear.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.py (1)
322-347: Consider adding a non-2D input case to exercise reshape/restore paths.Both wrappers flatten internal inputs; adding one 3D/4D test case would directly validate that path (not just 2D behavior).
Also applies to: 354-382
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.py` around lines 322 - 347, Add a new subcase in test_fused_relu2_quant_nvfp4_wrapper_matches_trtllm_op (and mirror in the sibling test at 354-382) that uses a non-2D input (e.g., a 3D or 4D tensor on CUDA with dtype torch.bfloat16) to force the wrapper's reshape/restore path; run the same sequence of ops (fp4_global_scale, torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4, compare shapes/dtypes, then validate functional equivalence via nvfp4_gemm using a matching quantized weight) and assert equality the same way as the existing 2D case so the reshape/restore logic is exercised and verified.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fuse_relu2_quant_nvfp4.py`:
- Around line 71-83: The try/except is too broad and currently swallows all
errors; replace the bare except with specific exceptions likely from attribute
access or dict lookups (for example catch AttributeError, KeyError, and
TypeError) and handle them explicitly (return None or re-raise after logging)
instead of pass; update the block around reading input_node.target /
gm.get_submodule(...) / type(submod) and the similar block at lines ~166-180 to
catch only those specific exceptions (e.g., except (AttributeError, KeyError,
TypeError) as e:) and either return None or log the error context (include
target and type_name) before returning so real failures aren’t silently
discarded.
---
Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py`:
- Around line 397-407: The bias parameter in trtllm_nvfp4_prequant_linear_fake
is intentionally unused but not explicitly discarded; to silence ARG001 and
clarify intent, explicitly mark it as unused (e.g., add bias to the existing
deletion of unused args in trtllm_nvfp4_prequant_linear_fake) so the function
body reads as discarding input_sf, weight_scale, alpha, and bias.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.py`:
- Around line 322-347: Add a new subcase in
test_fused_relu2_quant_nvfp4_wrapper_matches_trtllm_op (and mirror in the
sibling test at 354-382) that uses a non-2D input (e.g., a 3D or 4D tensor on
CUDA with dtype torch.bfloat16) to force the wrapper's reshape/restore path; run
the same sequence of ops (fp4_global_scale,
torch.ops.auto_deploy.trtllm_fused_relu2_quant_nvfp4, compare shapes/dtypes,
then validate functional equivalence via nvfp4_gemm using a matching quantized
weight) and assert equality the same way as the existing 2D case so the
reshape/restore logic is exercised and verified.
In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py`:
- Around line 95-133: Add a numerical equivalence check comparing the model
outputs before and after the transformation to ensure the non-fused path remains
semantically identical: run the original scripted/graph model (from
torch_export_to_gm or a forward call on model) to produce y_ref using input x,
then run the transformed graph gm_transformed to produce y_new and assert they
are numerically equal (e.g., torch.allclose or torch.testing.assert_close with
appropriate tolerances). Place this check in
test_fuse_relu2_quant_nvfp4_does_not_match_non_relu2 after gm_transformed is
created and before the _count_op assertions, referencing NonRelu2Model,
torch_quant_nvfp4_linear, torch_export_to_gm, and InferenceOptimizer so the
intent and location are clear.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ac31459a-c70c-4180-8084-50fd0934958a
📒 Files selected for processing (5)
tensorrt_llm/_torch/auto_deploy/config/default.yamltensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.pytensorrt_llm/_torch/auto_deploy/transform/library/fuse_relu2_quant_nvfp4.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/quantization/test_quant.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py
lucaslie
approved these changes
Mar 5, 2026
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
Collaborator
Author
|
/bot run |
tcherckez-nvidia
enabled auto-merge (squash)
March 8, 2026 07:16
Signed-off-by: tcherckez-nvidia <127761168+tcherckez-nvidia@users.noreply.github.com>
Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #38131 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #38131 [ run ] completed with state |
tianyuz-nv
pushed a commit
to wanqian-nv/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
…#11957) Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
limin2021
pushed a commit
to limin2021/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
…#11957) Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
longcheng-nv
pushed a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Mar 31, 2026
…#11957) Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add the ReLU2+NVFP4 fusion transform and required AutoDeploy wrapper ops, and cover the path with transform and custom-op unit tests.
Made-with: Cursor
Summary by CodeRabbit
New Features
Tests
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.