[Ops][Triton] Add a triton kernel supporting partial rope.

[whx-sjtu]

What this PR does / why we need it?

This PR adds a triton rope kernel witch supports scenarios of rope_dim != head_dim. This can save the split op before rope and the concat op after rope. Profiling shows improvement.

Original Implementation(2 split+2 rope+ 2 slice +2 concat):

Because currently we only support piecewise aclgraph for DS 3.2, so there are plenty of free bubbles. You can see the computing time of all rope related kernels: 35us

New Triton Rope: 12us

Does this PR introduce any user-facing change?

None

How was this patch tested?

I will add related ut after ci integrated with triton.

• vLLM version: v0.11.2
• vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.2

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Code Review

This pull request introduces a Triton kernel for Rotary Positional Embedding (RoPE) to support partial RoPE, where the RoPE dimension is not equal to the head dimension. This is a valuable performance optimization as it avoids explicit split and concat operations. The implementation includes a new Triton kernel _triton_rope and a wrapper function rope_forward_triton. The changes in sfa_v1.py correctly use this new kernel when Triton is available, with a fallback to the existing implementation. The Triton kernel itself appears to be well-written, handling both NEOX and non-NEOX styles, and correctly deals with padding and masking for variable dimensions. The upcasting to float32 for intermediate computations is a good practice for maintaining precision. I have one comment regarding a docstring that could be improved for clarity. Overall, the changes are logical and well-structured.

[whx-sjtu]

For prefill stage with larger bs, the performance gains remains substantial. Here is an example of bs=1001 with DS V3.2:

Without triton rope: 109us

With triton rope: 38us