[AutoDeploy] rms norm pattern matching

[suyoggupta]

This PR introduces RMSNorm support by adding a graph transformation that matches RMSNorm patterns and replaces them with optimized FlashInfer or Triton custom ops. It also integrates the new pattern matcher into the transformation pipeline and adds a unit test.

Add match_rms_norm_with_pm transform and register it in the main pipeline
Implement custom operators for FlashInfer and Triton RMSNorm
Introduce a new unit test for end-to-end validation

Test Coverage

Unit tests for the rms norm transformation

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]

Launch build/test pipelines. All previously running jobs will be killed.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-[Post-Merge]-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[Copilot]

Pull Request Overview

This PR introduces RMSNorm support by adding a graph transformation that matches RMSNorm patterns and replaces them with optimized FlashInfer or Triton custom ops. It also integrates the new pattern matcher into the transformation pipeline and adds a unit test.

• Add match_rms_norm_with_pm transform and register it in the main pipeline
• Implement custom operators for FlashInfer and Triton RMSNorm
• Introduce a new unit test for end-to-end validation

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rms_norm.py	Add unit tests for RMSNorm pattern matching and execution
tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher_utils.py	New utilities for FX-based pattern matchers with op-ignore support
tensorrt_llm/_torch/auto_deploy/transformations/transform.py	Invoke match_rms_norm_with_pm in the transform pipeline
tensorrt_llm/_torch/auto_deploy/transformations/library/rms_norm.py	Implement RMSNorm pattern registration and replacement logic
tensorrt_llm/_torch/auto_deploy/transformations/library/init.py	Export the new RMSNorm transform
tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py	Define FlashInfer and Triton custom RMSNorm operators
tensorrt_llm/_torch/auto_deploy/custom_ops/init.py	Export the new RMSNorm custom ops

Comments suppressed due to low confidence (2)

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rms_norm.py:44

• [nitpick] Rename the parameter variant to something more descriptive like backend to clarify its purpose in selecting the RMSNorm implementation.

@pytest.mark.parametrize("variant", ["flashinfer", "triton"])

tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rms_norm.py:43

• The tests currently only cover float16 inputs. Consider parameterizing over input dtypes (e.g., bfloat16 and float32) to validate pattern matching and custom ops across all supported precisions.

@pytest.mark.parametrize("eps", [1e-2, 1e-6])

[Copilot]

The current pattern registration hard-codes eps=1e-6 and does not ignore variations in eps; as a result, patterns with different eps values (e.g., 1e-2) won't match. Consider passing exclusive_arg_names=('eps',) to register_pattern so the eps argument is treated as a wildcard.

Suggested change

	scalar_workaround={"eps": 1e-6},
	exclusive_arg_names=('eps',),

[Copilot]

[nitpick] Remove or replace this print statement with a logging call (e.g., logger.debug) to avoid polluting test output.

Suggested change

	print(gm_transformed.graph)
	logger.debug(gm_transformed.graph)

[lucaslie]

We should probably wait until NVIDIA#5106 gets merged before we continue working on this one

[suyoggupta]

subsumed in #58