[Bugfix] Eliminate tuple inputs to submodules in graph partitioning

[gmagogsfm]

Move getitem operations to the same subgraph as their input to prevent submodules from receiving tuple arguments.

This is to fix issues like #24915 where split_graph partitioned graph at a node that produces a tuple, which is then passed to subsequent submodules as input. This unfortunately does not conform to the calling convention of AoT Autograd modules, which requires flattened arguments.

Signed-off-by: Yanan Cao [email protected]

[gemini-code-assist]

Code Review

This pull request effectively addresses the issue of tuple inputs to submodules in graph partitioning by moving getitem operations to the producer's subgraph. The changes in vllm/compilation/backends.py correctly implement this logic, and the new test file tests/compile/test_graph_partition.py provides good coverage for both single and multiple consumer scenarios, ensuring the fix works as intended. The addition of the test to the .buildkite/test-pipeline.yaml ensures continuous validation. The overall change improves the robustness of the graph partitioning logic for AoT Autograd modules. No critical or high-severity issues were found.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gmagogsfm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ProExpertProg]

This should always be true, right?

Suggested change

	if input_node in node_to_subgraph_id:
	node_to_subgraph_id[node] = node_to_subgraph_id[input_node]
	continue
	assert input_node in node_to_subgraph_id
	node_to_subgraph_id[node] = node_to_subgraph_id[input_node]
	continue

[gmagogsfm]

I think placeholder nodes would not be in node_to_subgraph_id map.

[zou3519]

hmm I think the claim is that we should never be getitem on a placeholder node or an output node so we can just assert that the input_node is in node_to_subgraph id?

Because Dynamo produces a graph where all placeholders should be Tensors or symints or scriptobjects

So I agree with Luka's suggestion

[gmagogsfm]

I see, changed.

[gmagogsfm]

@zou3519 @ProExpertProg I had to relax the assert to only when producer node is not a placeholder, because getitem is also used to fetch item/slice from a tensor as well. In that case, there can be legitimate getitem calls on placeholders. This is seen in qwen2.

PTAL.

[zou3519]

code reasonable, I think Luka's comment makes sense