[Main][Ops] Make triton rope support index_selecting from cos_sin_cache (#5450)

### What this PR does / why we need it? This PR extends original `rope_triton_forward` and `split_qkv_rmsnorm_rope` to support `cos_sin_cache` && `positions` as inputs. This fully aligns to vLLM RoPE api interface. Compared with earlier implementation for RoPE, the benefits are: 1. avoiding pre-computation of `cos` `sin` before model execution, which helps to remove redundant codes. 2. allowing eagle3 draft model to have different rope parameters with main model (see #6612 ). This help to recover accept rate && accuracy in that case. In addition, this kernel change only introduces very small performance degradation. Those `index_select` or `chunk` operations are now changed into simple memory access in triton kernel (For example, https://github.com/vllm-project/vllm-ascend/pull/5450/changes#diff-a4c2d3071530df193b98f9bf38553874bc4d47571336711f116c26d019cfbb6aR77-R81). **Highlights** - **RoPE Cache Unification**: Replaced separate _sin and _cos global tensors with a unified cos_sin_cache and explicit positions tensor for Rotary Positional Embeddings (RoPE), streamlining data handling. - **Triton Kernel Integration**: Updated Triton kernels (split_qkv_rmsnorm_rope_kernel, _triton_rope) to directly consume the cos_sin_cache and positions for more efficient and integrated RoPE calculations. - **Custom Operation Registration**: Registered `rope_forward_oot` as a new custom operation, allowing its use in fused compilation passes and providing a dedicated entry point for the new RoPE implementation. - **Refactored RoPE Forward Pass**: Modified the rope_forward_oot function to accept the new cos_sin_cache and positions arguments, enabling a more flexible and integrated RoPE application within the system. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - vLLM version: v0.13.0 - vLLM main: 5326c89803 Additional test on Qwen3-235b accuracy: | Aime2024 | GSM8K | Livecodebench | | -------- | -------- | -------- | | 83.33 | 96.26 | 70.23 | --------- Signed-off-by: Angazenn <supperccell@163.com>
2026-02-11 21:20:53 +08:00
parent 6bc44bf49b
commit c0c2eb614e
13 changed files with 378 additions and 243 deletions
--- a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_rope.py
+++ b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_rope.py
@@ -8,6 +8,7 @@ from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton

 IS_NEOX_STYLE = [True, False]
 DTYPES = [torch.bfloat16, torch.float16]
+MAX_POSITION_EMBEDDINGS = [262144]
 HEAD_SIZES = [64, 128]
 ROTARY_DIMS = [32, 64]
 NUM_Q_HEADS = [64]
@@ -139,3 +140,83 @@ def test_rotary_embedding_triton_kernel(
    gc.collect()
    torch.npu.empty_cache()
    torch.npu.reset_peak_memory_stats()
+
+
+
+@pytest.mark.parametrize("max_position_embeddings", MAX_POSITION_EMBEDDINGS)
+@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("num_q_heads", NUM_Q_HEADS)
+@pytest.mark.parametrize("num_k_heads", NUM_K_HEADS)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("rotary_dim", ROTARY_DIMS)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.parametrize("device", DEVICES)
+@torch.inference_mode()
+def test_rotary_embedding_triton_kernel_with_cos_sin_cache(
+    max_position_embeddings: int,
+    is_neox_style: bool,
+    num_tokens: int,
+    num_q_heads: int,
+    num_k_heads: int,
+    head_size: int,
+    rotary_dim: int,
+    dtype: torch.dtype,
+    seed: int,
+    device: str,
+) -> None:
+    torch.manual_seed(seed)
+    torch.set_default_device(device)
+    init_device_properties_triton()
+    if rotary_dim == -1:
+        rotary_dim = head_size
+    cos_sin_cache = torch.randn(max_position_embeddings, rotary_dim, dtype=dtype, device=device)
+    positions = torch.randint(low=0, high=max_position_embeddings, size=(num_tokens,), dtype=torch.int64, device=device)
+    q_trt = torch.randn(num_tokens,
+                        num_q_heads,
+                        head_size,
+                        dtype=dtype,
+                        device=device)
+    k_trt = torch.randn(num_tokens,
+                        num_k_heads,
+                        head_size,
+                        dtype=dtype,
+                        device=device)
+    q_gold = torch.randn(num_tokens,
+                         num_q_heads,
+                         head_size,
+                         dtype=dtype,
+                         device=device)
+    k_gold = torch.randn(num_tokens,
+                         num_k_heads,
+                         head_size,
+                         dtype=dtype,
+                         device=device)
+    q_trt.copy_(q_gold)
+    k_trt.copy_(k_gold)
+    q_trt, k_trt = rope_forward_triton(q_trt,
+                                       k_trt,
+                                       cos_sin_cache=cos_sin_cache,
+                                       positions=positions,
+                                       rope_dim=rotary_dim,
+                                       is_neox_style=is_neox_style)
+    cos, sin = cos_sin_cache.index_select(0, positions).chunk(2, dim=-1)
+    q_gold, k_gold = _rope_pytorch_native(q_gold,
+                                          k_gold,
+                                          cos,
+                                          sin,
+                                          rope_dim=rotary_dim,
+                                          is_neox_style=is_neox_style)
+    # Compare the results.
+    torch.testing.assert_close(q_trt.view(q_gold.size()),
+                               q_gold,
+                               atol=DEFAULT_ATOL,
+                               rtol=DEFAULT_RTOL)
+    torch.testing.assert_close(k_trt.view(k_gold.size()),
+                               k_gold,
+                               atol=DEFAULT_ATOL,
+                               rtol=DEFAULT_RTOL)
+    gc.collect()
+    torch.npu.empty_cache()
+    torch.npu.reset_peak_memory_stats()
--- a/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_split_qkv_rmsnorm_rope.py
+++ b/tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_split_qkv_rmsnorm_rope.py
@@ -7,6 +7,7 @@ import torch
 import vllm_ascend.ops.register_custom_ops  # noqa
 from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton

+MAX_POSITION_EMBEDDINGS = [262144]
 NUM_TOKENS = [1, 4, 8, 16, 1024]
 NUM_QKV_HEADS = [(12, 1), (16, 1), (32, 4), (64, 4)]
 HEAD_SIZES = [128]
@@ -55,6 +56,7 @@ def rms_norm(
    return out


+@pytest.mark.parametrize("max_position_embeddings", MAX_POSITION_EMBEDDINGS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("num_q_heads, num_kv_heads", NUM_QKV_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -63,7 +65,7 @@ def rms_norm(
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", DEVICES)
@torch.inference_mode()
-def test_split_qkv_rmsnorm_rope(num_tokens, num_q_heads, num_kv_heads,
+def test_split_qkv_rmsnorm_rope(max_position_embeddings, num_tokens, num_q_heads, num_kv_heads,
                                head_size, eps, dtype, seed, device):
    torch.manual_seed(seed)
    torch.set_default_device(device)
@@ -77,12 +79,10 @@ def test_split_qkv_rmsnorm_rope(num_tokens, num_q_heads, num_kv_heads,
                      device=device)
    q_weight = torch.randn(head_size, dtype=dtype, device=device)
    k_weight = torch.randn(head_size, dtype=dtype, device=device)
-    sin = torch.from_numpy(
+    cos_sin_cache = torch.from_numpy(
        np.random.uniform(0, 1,
-                          [num_tokens, 1, 1, head_size])).to(dtype).npu()
-    cos = torch.from_numpy(
-        np.random.uniform(0, 1,
-                          [num_tokens, 1, 1, head_size])).to(dtype).npu()
+                          [max_position_embeddings, head_size])).to(dtype).npu()
+    positions = torch.randint(low=0, high=max_position_embeddings, size=(num_tokens,), dtype=torch.int64, device=device)
    # fused kernel
    q, k, v = torch.ops.vllm.qkv_rmsnorm_rope(input=qkv,
                                              q_weight=q_weight,
@@ -91,8 +91,12 @@ def test_split_qkv_rmsnorm_rope(num_tokens, num_q_heads, num_kv_heads,
                                              kv_hidden_size=kv_hidden_size,
                                              head_dim=head_size,
                                              eps=eps,
-                                              cos=cos,
-                                              sin=sin)
+                                              cos_sin_cache=cos_sin_cache,
+                                              positions=positions)
+    
+    cos, sin = cos_sin_cache.index_select(0, positions).view(num_tokens, 2, -1).repeat(1, 1, 2).chunk(2, dim=-2)
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)

    # split
    _q, _k, v_gold = qkv.cpu().split(
@@ -129,6 +133,7 @@ def test_split_qkv_rmsnorm_rope(num_tokens, num_q_heads, num_kv_heads,
    torch.npu.reset_peak_memory_stats()


+@pytest.mark.parametrize("max_position_embeddings", MAX_POSITION_EMBEDDINGS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("num_q_heads, num_kv_heads", NUM_QKV_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -137,7 +142,7 @@ def test_split_qkv_rmsnorm_rope(num_tokens, num_q_heads, num_kv_heads,
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", DEVICES)
@torch.inference_mode()
-def test_split_qkv_rmsnorm_rope_with_bias(num_tokens, num_q_heads,
+def test_split_qkv_rmsnorm_rope_with_bias(max_position_embeddings, num_tokens, num_q_heads,
                                          num_kv_heads, head_size, eps, dtype,
                                          seed, device):
    torch.manual_seed(seed)
@@ -154,12 +159,10 @@ def test_split_qkv_rmsnorm_rope_with_bias(num_tokens, num_q_heads,
    k_weight = torch.randn(head_size, dtype=dtype, device=device)
    q_bias = torch.randn(head_size, dtype=dtype, device=device)
    k_bias = torch.randn(head_size, dtype=dtype, device=device)
-    sin = torch.from_numpy(
+    cos_sin_cache = torch.from_numpy(
        np.random.uniform(0, 1,
-                          [num_tokens, 1, 1, head_size])).to(dtype).npu()
-    cos = torch.from_numpy(
-        np.random.uniform(0, 1,
-                          [num_tokens, 1, 1, head_size])).to(dtype).npu()
+                          [max_position_embeddings, head_size])).to(dtype).npu()
+    positions = torch.randint(low=0, high=max_position_embeddings, size=(num_tokens,), dtype=torch.int64, device=device)
    # fused kernel
    q, k, v = torch.ops.vllm.qkv_rmsnorm_rope(input=qkv,
                                              q_weight=q_weight,
@@ -170,8 +173,12 @@ def test_split_qkv_rmsnorm_rope_with_bias(num_tokens, num_q_heads,
                                              eps=eps,
                                              q_bias=q_bias,
                                              k_bias=k_bias,
-                                              cos=cos,
-                                              sin=sin)
+                                              cos_sin_cache=cos_sin_cache,
+                                              positions=positions)
+    
+    cos, sin = cos_sin_cache.index_select(0, positions).view(num_tokens, 2, -1).repeat(1, 1, 2).chunk(2, dim=-2)
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)

    # split
    _q, _k, v_gold = qkv.cpu().split(