[Feature]: Optimize RoPE

[robertgshaw2-redhat]

🚀 The feature, motivation and pitch

Recently, we have noticed that rope is relatively slow with vLLM based on profiling.

goal is to optimize rope by:

• investigate why torch.compile's result is not optimal
• investigate alternative kernels (e.g. FlashInfer's rope)

Alternatives

No response

Additional context

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.

[tjtanaa]

@robertgshaw2-redhat We are working on integrating a triton mrope feature

[mgoin]

We should work on integrating the flashinfer rope kernel as a backend replacement for our custom op #21126

[TryinCoding]

@tjtanaa Hey, so is this issue still open? because I saw your PR on mrope and that accomplishes this right?

[mgoin]

Yes @TryinCoding, this is still open since that PR was just for mrope. We still are interested in improving the basic rope performance

[TryinCoding]

@mgoin I see. would implementing my own rope triton kernel be an optimal way to solve this?

[BoyuanFeng]

Did some benchmark on torch.compile vs hand-tuned kernels from Liger. torch.compile is faster on larger sequence lengths (>=512) but slower on smaller sequence lengths. WIP to improve torch.compile for smaller shapes.

For shape

q: [1, S, 32, D]
kv: [1, S, 32, D]

For shape

q: [1, S, 32, D]
kv: [1, S, 8, D]

[ProExpertProg]

@BoyuanFeng is Inductor generating a single triton kernel in the microbenchmark? Because in full model compilation Inductor produces 3 triton kernels for RoPE