feat: [AutoDeploy] generalizing cudagraph to multiple dynamic inputs

[lucaslie]

Previously, our cudagraph implementation assumed that only the first input is batch-size dependent. This update generalizes the implementation to handle multiple inputs to the cudagraph that are batch-size dependent.

Unit tests and integration tests are also updated for improved coverage

---

Copilot Summary

This pull request includes multiple changes to improve the functionality of the auto-deploy system and enhance the testing framework. The key changes include modifications to the benchmarking logic, updates to the CompiledGraph class, and enhancements to unit tests for better coverage and performance.

Improvements to Benchmarking Logic:

• Updated main function in examples/auto_deploy/build_and_run_ad.py to add a condition that prevents benchmarking when the runtime is "demollm" and provides a warning message.

Enhancements to CompiledGraph Class:

• Added support for multiple batched inputs in CompiledGraph by introducing num_batched_inputs parameter and updating related methods to handle multiple input buffers. [1] [2] [3] [4] [5]
• Modified TorchOptCompiler to pass num_batched_inputs to CompiledGraph.

Testing Framework Enhancements:

• Added a new model ModelWithMultipleInputs and updated tests to handle multiple input scenarios, ensuring comprehensive coverage for the new functionality. [1] [2]
• Skipped a slow unit test in test_deepseek_patches.py to improve test suite performance.

Codebase Simplification:

• Simplified the _flatten_args function in compiler.py by removing unnecessary assertions and returning a list directly.
• Updated the initialization logic in compiler.py to handle dynamic shapes more efficiently.

These changes collectively enhance the robustness and performance of the auto-deploy system, while also improving the maintainability and efficiency of the codebase.

[lucaslie]

/bot run

[Copilot]

Copilot reviewed 8 out of 8 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (1)

tensorrt_llm/_torch/auto_deploy/compile/compiler.py:21

• The previous assertion that ensured the main input tensor was at least 2D has been removed. To prevent potential runtime issues with inputs of insufficient dimensions, consider adding an alternative check or a clear error message.

return tree_flatten_spec(all_args, in_spec)

[tensorrt-cicd]

PR_Github #2377 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #2377 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1711 completed with status: 'FAILURE'

[lucaslie]

/bot run

[tensorrt-cicd]

PR_Github #2490 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #2490 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #1786 completed with status: 'FAILURE'

[Fridah-nv]

This update is also important to support models like Llama4ForConditionalGeneration that take both text inputs and image inputs. Should we add a similar test model in test_torch_opt.py ?

[lucaslie]

This update is also important to support models like Llama4ForConditionalGeneration that take both text inputs and image inputs. Should we add a similar test model in test_torch_opt.py ?

@Fridah-nv, I have a general test that tests for a variable number of inputs with batch_dim. I think that's sufficient for now. IF there is a different way we need to handle Llama-4 let's do that later

[Fridah-nv]

This update is also important to support models like Llama4ForConditionalGeneration that take both text inputs and image inputs. Should we add a similar test model in test_torch_opt.py ?

@Fridah-nv, I have a general test that tests for a variable number of inputs with batch_dim. I think that's sufficient for now. IF there is a different way we need to handle Llama-4 let's do that later

I see, thanks for the reply

[Fridah-nv]

The changes LGTM. @suyoggupta could you approve it since I don't seem to have rights to approve PR

[lucaslie]

/bot run

[tensorrt-cicd]

PR_Github #2972 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #2972 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #2084 completed with status: 'FAILURE'

[lucaslie]

/bot run

[tensorrt-cicd]

PR_Github #3081 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #3081 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #2150 completed with status: 'FAILURE'

[lucaslie]

/bot skip --comment "Unrelated test failure; all other tests pass"

[tensorrt-cicd]

PR_Github #3091 [ skip ] triggered by Bot

[tensorrt-cicd]

PR_Github #3091 [ skip ] completed with state SUCCESS
Skipping testing for commit a59456d