Fix cudagraphs, add rms norm pattern matcher

[suyoggupta]

• Add pattern matcher for rms norm + unit test
• Capture graphs from largest batch size to smallest
  AutoDeploy:

===========================================================                                                                       
= PERFORMANCE OVERVIEW                                                                                                            
===========================================================
Request Throughput (req/sec):                     109.4132
Total Output Throughput (tokens/sec):             14004.8949
Total Token Throughput (tokens/sec):              28009.7898
Total Latency (ms):                               9139.6616
Average request latency (ms):                     8749.7151
Per User Output Throughput [w/ ctx] (tps/user):   14.6394
Per GPU Output Throughput (tps/gpu):              14004.8949

Pytorch:

===========================================================                                                                       
= PERFORMANCE OVERVIEW                                                                                                            
===========================================================                                                                       
Request Throughput (req/sec):                     117.5372                                                                        
Total Output Throughput (tokens/sec):             15044.7646                                                                      
Total Token Throughput (tokens/sec):              30089.5293                                                                      
Total Latency (ms):                               8507.9430                                                                       
Average request latency (ms):                     8157.6348                                                                       
Per User Output Throughput [w/ ctx] (tps/user):   15.7019                                                                         
Per GPU Output Throughput (tps/gpu):              15044.7646

[Copilot]

Pull Request Overview

This PR adds a graph transformation to match and replace RMSNorm patterns, integrates the new matcher into the transformation pipeline, provides custom operators and unit tests for both FlashInfer and Triton backends, and also fixes cuDAGraph batch size ordering and extends the executor to handle new tokens.

• Introduce match_rms_norm pattern matcher, custom ops, and corresponding unit tests.
• Update the main transform pipeline and initialization (transform.py) to invoke RMSNorm matching.
• Fix cuDAGraph batch size sorting, add max_beam_width and new_tokens handling in ad_executor.py, and adjust the default overlap scheduler flag.

Reviewed Changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 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 across both backends.
tensorrt_llm/_torch/auto_deploy/transformations/transform.py	Import and invoke match_rms_norm in the transformation pipeline.
tensorrt_llm/_torch/auto_deploy/transformations/library/rms_norm.py	Implement search and replacement functions for RMSNorm optimization.
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py	Add max_beam_width, extend _prepare_inputs to accept new_tokens.
tensorrt_llm/_torch/auto_deploy/llm_args.py	Change default for disable_overlap_scheduler to False.
tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py	Define FlashInfer and Triton custom ops for RMSNorm.
tensorrt_llm/_torch/auto_deploy/custom_ops/init.py	Export rms_norm operators.
tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py	Sort cuda_graph_batch_sizes in descending order.

Comments suppressed due to low confidence (4)

tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py:9

• The function name flashinfer_rmsnorm (no underscore) differs from the op name flashinfer_rms_norm. Consider renaming the Python function to flashinfer_rms_norm for consistency.

@torch.library.custom_op("auto_deploy::flashinfer_rms_norm", mutates_args=())

tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py:42

• The function name triton_rmsnorm (no underscore) does not match the op name triton_rms_norm. Renaming it to triton_rms_norm would improve readability and maintain consistency.

@torch.library.custom_op("auto_deploy::triton_rms_norm", mutates_args=())

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py:157

• The docstring for _prepare_inputs was not updated to mention the new new_tokens parameter. Please extend the docstring to describe its purpose and expected behavior.

    ) -> bool:

tensorrt_llm/_torch/auto_deploy/transformations/transform.py:196

• [nitpick] The call to match_rms_norm uses the default backend. It may be clearer to surface the backend option here or pass it via configuration rather than relying on the default in the function signature.

        match_rms_norm(egm)

[suyoggupta]

@Fridah-nv : ptal.

[Fridah-nv]

LGTM for the pattern matcher part!

[nzmora-nvidia]

LGTM

[lucaslie]

Here is a rough summary of my investigation of hitting torch.compile's recompile limit during cuda graph capture and a suggestion on how to fix it

diff --git a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py
index 2f0bd3d85..29653e2c2 100644
--- a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py
+++ b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py
@@ -11,5 +11,14 @@ class TorchOptCompiler(TorchCudagraphCompiler):
     """Compiler that uses both torch.compile and CUDA graphs."""
 
     def _init_captured_graph(self, gm, in_spec, out_spec) -> CapturedGraph:
+        # TODO: every graph capture will trigger a recompile, hence we need a higher limit
+        # TODO: refine this fix:
+        #   1. rewrite CapturedGraph utility to always take cudagraph batch sizes as input
+        #       a. max batch size can be inferred from cudagraph batch sizes
+        #       b. move heuristic to batch sizes from max_batch_size to TorchCudagraphCompiler class
+        #   2. Set the recompile limit here to be max(len(cuda_graph_batch_sizes), torch._dynamo.config.recompile_limit)
+        #   3. Consider moving cache_size_limit config to here or TorchCompile backend as well
+        #      --> right now, it's hard-coded into build_and_run_ad.py
+        torch._dynamo.config.recompile_limit = 20  # will also set cache_size_limit to 20
         gm = torch.compile(gm, dynamic=True)
         return super()._init_captured_graph(gm, in_spec, out_spec)

It's not a blocker though - we can probably fix it a follow-up PR. @suyoggupta maybe you can just a comment to the torch-opt backend for this

@nzmora-nvidia, maybe you can take this on as a small task to get more familiar with our cudagraph capture?

[nzmora-nvidia]

I created NVIDIA#6142