TensorRT icon indicating copy to clipboard operation
TensorRT copied to clipboard
verified publisher icon pytorch

adding rotary embedding example, with graph rewrite for complex subgraph [WIP]

Open apbose opened this issue 7 months ago • 1 comments

This PR-

  1. Adds an example for parallel rotary embedding
  2. Adds logic for complex graph detection
  3. Adds a pass for complex graph rewrite in aten_lowering_pass Please note that this PR is currently for the single GPU case where there is no DTensor in the inputs of the torch module. Ideally this should not require runtime changes. This should avoid the graph breaks caused due to view_as_complex and view_as_real nodes.

apbose avatar Jun 13 '25 00:06 apbose

I see lint error here would reformat /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py which is unrelated to my change. Not sure why is it failing here

apbose avatar Jun 18 '25 04:06 apbose