Piecewise compilation

[nopperl]

This PR enables piecewise torch.compile support for PLaMo2. The compilation is only enabled for attention layers, because enabling torch.compile for mamba layers requires a redesign of the mamba cache in upstream vLLM.

TODO: add benchmark

Illustration

The compilation result is exemplified using the QK-norm in the Attention layer:

vllm/vllm/model_executor/models/plamo2.py

Lines 422 to 427 in 7bc2e28

	q_shape = q.shape
	q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape)
	q = self.q_norm.forward_native(q).reshape(q_shape)
	k_shape = k.shape
	k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape)
	k = self.k_norm.forward_native(k).reshape(k_shape)

This is the timeline of the RMSNorm(Q) and RMSNorm(K) ops:

This is the timeline of the same ops with compilation enabled:

The entries for the QK-norm are between cutlass::Kernel2(/triton_tem_fused_mm) and the one prior to flash::flash_fwd_splitkv_kernel. Without compilation, an individual kernel is launched for every torch op inside the QK-norm. The launch and memory transaction overhead for these kernels is significantly higher than the actual computation.

With compilation, these kernels are fused into a single Triton kernel, which drastically reduces the runtime of the QK-norm.

Example

Example inference script with compilation configuration: https://github.com/pfn-attic/vllm-plamo2-plugin/blob/piecewise-compile/example.py

Output:

Prompt: 'Hello, my name is', Generated text: ' Bobo and I would like to borrow money from you today. If I’m lucky, I’ll get $5'
Prompt: 'The president of the United States is', Generated text: ' the US President or US President. Democratic Republic of Congo was in 1996'
Prompt: 'The capital of France is', Generated text: ' Paris, but do you know which city is the capital of Italy? Where is Rome located in the world'
Prompt: 'The future of AI is', Generated text: ' in your hands, and yours alone\nArtificial Intelligence (AI) is changing the'

By setting compilation_config=0 in the above script, the compilation can be disabled. I have confirmed that the model outputs are identical with compilation enabled and disabled.

Note: This PR depends on/includes the fix in vllm-project#14913.

[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.

🚀