fix(runtime): keep the output structure under whole-graph CUDA graphs - #4906

Stacked on #4905. The base is narendasan/cudagraphs-kwargs-order.

Summary

• CudaGraphsTorchTensorRTModule (whole-graph CUDA graphs, used when a compiled module has PyTorch subgraphs) cloned its outputs out of the static buffers as a flat list: [output.clone() for output in output_buffers]. That assumes a tensor or a flat sequence of tensors.
• An exported Hugging Face model returns its ModelOutput (a pytree), so with any graph break every call failed:

  AttributeError: 'CausalLMOutputWithPast' object has no attribute 'clone'
  

  Nested tuples and dicts fail the same way.
• Fix: clone the tensor leaves with tree_map, keeping the output structure.
  • Tuples are now returned as tuples, and a one-element sequence as a sequence, matching what the module returns without CUDA graphs. Before, the module returned lists and unwrapped single elements.
  • As before (fix(runtime): record inputs updated in place as the caller's tensors under whole-graph CUDA graphs #4904), an output inside an input that is updated in place is returned as is.

Why it matters

Graph breaks cost the most when every engine boundary is a separate launch, and CUDA graphs are the remedy. For example, TensorRT 11.3's attention at head_dim 256 is several times slower than PyTorch's FlashAttention-2, so Gemma runs faster with its attention left in PyTorch, despite 26 graph breaks. For a 256-token forward of Gemma-3-1B (bf16, RTX 5090):

• 5.89 ms with everything in TensorRT.
• 4.55 ms with attention in PyTorch and whole-graph CUDA graphs (this fix), against 4.48 ms for torch.compile(mode="max-autotune").

Gemma-2-2B gets 8.69 ms, against 8.51 ms. Without this fix, enable_cudagraphs cannot run such a module at all.

Tests

• New test_cudagraphs_whole_graph_nested_outputs_cpp (a dict-with-tuple output, a graph break forced with torch_executed_ops) fails without the change with the error above.
• pytest tests/py/dynamo/runtime/test_002_cudagraphs_{cpp,py}.py test_aliased_io_cudagraphs.py test_empty_output.py test_input_lifetime.py test_output_allocator.py test_004_weight_streaming.py: 49 passed.
• Found by a Hugging Face architecture sweep. RTX 5090, TensorRT 11.3.

Stack

3 of 4. Base: #4905. Next: #4907.

1. fix(runtime): record inputs updated in place as the caller's tensors under whole-graph CUDA graphs #4904 fix(runtime): record inputs updated in place as the caller's tensors under whole-graph CUDA graphs
2. fix(runtime): flatten the inputs by the input spec under whole-graph CUDA graphs #4905 fix(runtime): flatten the inputs by the input spec under whole-graph CUDA graphs
3. fix(runtime): keep the output structure under whole-graph CUDA graphs #4906 fix(runtime): keep the output structure under whole-graph CUDA graphs
4. fix(runtime): create the warm-up inputs before unflattening them under whole-graph CUDA graphs #4907 fix(runtime): create the warm-up inputs before unflattening them under whole-graph CUDA graphs

🤖 Generated with Claude Code