[CUDA] Enable full cudagraph for FlashMLA

[ProExpertProg]

Enable fullgraph CUDAGraph capture for the FlashMLA decode case.

Hacks:

• building the capture metadata
• prefill batch bypasses compiled code and manually calls eager code

Tested with:

python examples/offline_inference/basic/generate.py --model deepseek-ai/DeepSeek-V2-Lite --trust-remote-code -O {"full_cuda_graph":true}

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[hypnopump]

Could you give some details on speedup associated with this modification?

[mergify]

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

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

[ProExpertProg]

I'm thinking about setting a context to bypass Inductor backend directly here, but I guess we could also separately compile (and capture) prefill stage - any thoughts?
cc @youkaichao

[youkaichao]

this is hacky and depends on the fact that self.model.model is the underlying nn.Module we compile (which might not be the case if we are running multi-modality models).

I think we can have some fields in the forward context, like not_use_cudagraph, and set it to true only when we capture full cudagraph but get prefill data in the current batch. then, when we decide to replay the cudagraph, we can check this field.

it should just be changing this line

vllm/vllm/compilation/cuda_piecewise_backend.py

Line 212 in b9f61e1

entry.cudagraph.replay()

to return entry.runnable(*args)

[ProExpertProg]

Could you give some details on speedup associated with this modification?

I haven't necessarily profiled this but it's meant to enable the double-batch-overlap optimization (prototype in #18415)

[mergify]

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

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

[izhuhaoran]

Hi, any further progress on this pr?

[ProExpertProg]

Hi, any further progress on this pr?

Almost ready for review!

[ProExpertProg]

Currently experiencing some issues when batching (in unit test), need to investigate further.

[mergify]

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

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