[WIP] Use FlashInfer RoPE

[james-p-xu]

Motivation

NOTE: flashinfer.apply_rope_pos_ids does not exist in the prebuilt wheel, must build from source. Is this an issue?

We want to verify the correctness of flashinfer's RoPE against vLLM's RoPE, in preparation of replacing vLLM's get_rope with flashinfer's.

cc: @ByronHsu

Modifications

Added standalone python script for comparison.

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

[yzh119]

It's worth noting that flashinfer uses fp32 internally for sin/cos, we found there will be some non-trivial output difference if we use fp16 sin/cos.

[zhyncs]

@james-p-xu How is it going? Has it already run successfully after removing this dependency using flashinfer latest https://github.com/flashinfer-ai/flashinfer-nightly/releases

sglang/python/sglang/srt/models/llama.py

Line 25 in 60769be

from vllm.model_executor.layers.rotary_embedding import get_rope

[zhyncs]

I will first merge this PR into the sgl-project:rpoe branch today to run the CI with nightly FlashInfer. You can check it out later and submit a new PR. Thanks!

[james-p-xu]

Looks like there is another correctness issue with flashinfer apply_rope_pos_ids function. Will sync with Zihao offline on this.

[zhyncs]

also related with #2620

[zhyncs]

Hi James @james-p-xu, we've decided to adopt this #2964, and @ByronHsu will help rewrite the CUDA kernel. I'll close this PR for now. Thanks for your contribution!