[dynamo][hops] Support non proxyable outputs in hops

[anijain2305]

Stack from ghstack (oldest at bottom):

• [dynamo] Rehaul the autograd.Function support #166788
• -> [dynamo][hops] Support non proxyable outputs in hops #167377

PR Description and code-comments generated by claude

Summary

Refactors speculate_subgraph to decouple graph outputs from output VariableTrackers (VTs), enabling better support for higher-order operators (HOPs) with non-proxyable outputs.

Motivation

Previously, speculate_subgraph unnecessarily entangled the graph output and the call_function VTs that needed to be sent for later tracing. In Dynamo (outside of speculate subgraph), VTs and graph are separate concepts: VTs can run ahead while the graph is just a side data structure of computation seen so far.

This entanglement caused issues with HOP support where subgraph outputs are not proxyable. The old implementation would call pytree flatten because it was unnecessarily coupling graph_output and output VTs, leading to complications with non-tensor/non-proxyable returns.

Changes

Core Refactoring:

• Removed the OutputSpec dataclass that tracked treespec, constant value masks, and filtering logic
• Introduced graph_output_vts - a separate tuple of VTs that should be graph outputs
• VTs are now collected by visiting the output structure and extracting only TensorVariable and SymNodeVariable instances
• Graph outputs and function return values are now properly separated concepts

Key Benefits:

1. Cleaner separation of concerns: Graph outputs (what goes in the FX graph) vs function outputs (what Dynamo continues tracing with)
2. Better non-proxyable output support: Enables HOPs to return custom objects (UDFs) that contain tensors without forcing them to be graph outputs
3. Simplified code: Removed complex logic for filtering constants and maintaining treespecs
4. Maintains proxy overwriting: Still correctly overwrites original VT proxies with subgraph output proxies for proper graph construction

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/167377

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 4325dd1 with merge base 47acdea ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor / unit-test / inductor-pallas-test / test (inductor-pallas, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[zou3519]

nit: Use aot_eager if you're not going to do anything with the graphs :P