[RFC]: Address piecewise graph splitting and attention fusion incompatibility

[ProExpertProg]

UPDATE: we chose approach 2, implemented in #24281 and pytorch/pytorch#162207.

Motivation.

I wanted to get a quick opinion from people on possible solutions on the issue where attention+quant fusion is incompatible with splitting_ops and hence piecewise cudagraphs. I've thought about it before but @nvpohanh brought it up in a DM this morning.

I am looking for feedback on both possible solutions and whether this is worth solving.

Problem:

There is an incompatibility between @fhl2000 ’s FULL_AND_PIECEWISE and the splitting_ops setting. In vLLM, we are currently splitting the FX graph before we start to apply graph fusion passes. However, we would like to enable Attn+Q fusions, so we clear the splitting_ops to keep the Attn op in the FX graph, but that breaks the FULL_AND_PIECEWISE mode because now we cannot capture piece-wise cuda graphs with full FX graph.

Proposed Change.

There are a few possible solutions:

1. Set splitting_ops=[] and use cudagraph_mode=FULL (current approach)
2. Split the graph in Inductor after custom passes
3. Split the graph after AotDispatcher (and custom passes) but before Inductor
4. While splitting the graph, make modifications to enable fusion

More details and the specific drawbacks of each approach are posted below.

1. Disable splitting and use cudagraph_mode=FULL

There are a few issues with this approach:

• Mixed decode-prefill batches can be faster with piecewise cudagraphs.
• Cascade attention requires piecewise cudagraphs.
• Some attention backends don't always support cudagraph_mode=FULL. But those backends might not support fp8 fusion either (MLA) so this might not be a big issue.

2. Split the graph in Inductor after custom passes

This approach was endorsed by @nvpohanh. I'm not sure if Inductor could handle a split graph during compilation (unsplit graph going into post_grad_custom_post_pass and split graph coming out)

3. Split the graph after AotDispatcher (and custom passes) but before Inductor

This is my current favorite approach. It would mean custom passes receive normalized and functionalized IR and can operate on the whole graph. However, this approach could increase compile time as we're not reusing piecewise graphs for the aot piece (minor) and if the Inductor piece reruns aot (major but avoidable).

4. While splitting the graph, make modifications to enable fusion

This means either performing attn+quant fusion manually on the Dynamo graph or moving ops around to make sure quant ops end in the same graph as attention. This is the dirtiest, least scalable, but also the quickest approach.

---

While approaches 2 and 3 would be best long-term, they might require more work to make sure caching of both the whole graph and the subgraphs works as intended. I am not 100% confident this wouldn't break something.

Feedback Period.

8/20-8/27

CC List.

@youkaichao @zou3519 @BoyuanFeng @nvpohanh @mgoin @ilmarkov

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[ProExpertProg]

#21031 enables SequenceParallelismPass for dynamic shapes but also only works on the full graph. So that's another pass that would benefit from this. Likely there will be other passes in the future that prefer to operate on the full graph

[BoyuanFeng]

Split the graph in Inductor after custom passes

Yes we can split graph in inductor. We can add tags=(torch._C.Tag.cudagraph_unsafe,) when we register a custom op. Then inductor will split the graph at the custom op. We can apply cudagraph to each graph partition.

[zou3519]

Still thinking, here's a quick comment.

2. Split the graph in Inductor after custom passes

As @BoyuanFeng mentioned, this is basically getting vLLM to use torch.compile's built-in CUDAGraph splitting. torch.compile's built-in CUDAGraph splitting happens after custom passes in Inductor.

[zou3519]

3. Split the graph after AotDispatcher (and custom passes) but before Inductor

I think my main comment is that this is possible, but may come with some large amount of pain, especially if we don't re-use torch._inductor.standalone_compile. We built torch._inductor.standalone_compile (it does AOTDispatcher+Inductor together) to stabilize how vLLM calls AOTDispatcher+Inductor and fix various caching issues. The pain that we may run into are: (1) unclear dependencies being passed between AOTDispatcher and Inductor and (2) coming up with a caching strategy and making sure it works.

[ProExpertProg]

Split the graph in Inductor after custom passes

Yes we can split graph in inductor. We can add tags=(torch._C.Tag.cudagraph_unsafe,) when we register a custom op. Then inductor will split the graph at the custom op. We can apply cudagraph to each graph partition.

One issue I see with that is we want to be able to achieve CUDAGraphWrapper functionality from the outside. So we want to be able to tell the compiled graph (at runtime) whether to replay (or record if not recorded) a graph or whether to run in eager mode. That's because we want to have the ability to dispatch between full cudagraphs and piecewise at runtime, ideally using a single FX graph.

Another feature that might be nice would be replaying different cudagraphs based on the batch composition or in the future expert load or different cudagraphs for different mamba layers. But we can cross that bridge when we get there, right now we only split on attention.

A workable approach here would be to be able to hook into the inductor CG mechanism or manually redo the splitting with custom wrappers if needed.

[ProExpertProg]

3. Split the graph after AotDispatcher (and custom passes) but before Inductor

I think my main comment is that this is possible, but may come with some large amount of pain, especially if we don't re-use torch._inductor.standalone_compile. We built torch._inductor.standalone_compile (it does AOTDispatcher+Inductor together) to stabilize how vLLM calls AOTDispatcher+Inductor and fix various caching issues. The pain that we may run into are: (1) unclear dependencies being passed between AOTDispatcher and Inductor and (2) coming up with a caching strategy and making sure it works.

Yeah I think that's fine, I guess we should measure how long Dynamo+Aot+standalone takes compared to Dynamo+standalone.

[zou3519]

2. Split the graph in Inductor after custom passes

@ProExpertProg I misunderstood @BoyuanFeng's comment. What Boyuan meant was that Inductor already has a mechanism to "split the graph based on specific custom operators". We can offer a user extension point to allow users to put their own wrapper functions around these graph splits (so that vLLM can put its CUDAGraph Wrapper around this).

[ProExpertProg]

We can offer a user extension point to allow users to put their own wrapper functions around these graph splits (so that vLLM can put its CUDAGraph Wrapper around this).

I think this (approach 2 + custom torch cudagraph hook) seems like the best way forward. If it doesn't work out, we can try the Dynamo+Aot+standalone route. @youkaichao does that sound good to you?