[WS1][kernels] RoPE (Triton, CUDA on SM90)

benchmarks/benchmark_rope.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Benchmark NativeRoPEOp vs TritonRoPEOp vs RoPESM90Op.
+
+RoPE is an elementwise per-position rotation. The native path builds broadcast
+cos/sin caches and does the rotate-half in PyTorch; the Triton and CUDA (SM90)
+kernels fuse the rotation and round back to the input dtype on store, so they
+touch less memory and are faster. Latency (forward and forward+backward) and peak
+forward VRAM are reported, swept over (batch, seq) on the Qwen3-8B rotary config
+(n_heads=32, head_dim=128, theta=1e6).
+
+Usage:
+    python benchmarks/benchmark_rope.py
+    python benchmarks/benchmark_rope.py --configs "8,512;16,4096"
+"""
+
+import argparse
+
+import torch
+from tabulate import tabulate
+
+from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp
+from rl_engine.kernels.ops.triton.rotary_embedding.rope import TritonRoPEOp
+from rl_engine.platforms.device import device_ctx
+from rl_engine.utils.logger import logger
+
+# Qwen3-8B rotary config.
+N_HEADS = 32
+HEAD_DIM = 128
+THETA = 1.0e6
+
+
+def _maybe_sm90_op():
+    """The Hopper (SM90) RoPE op, or None when unavailable (non-Hopper / not built)."""
+    from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE
+
+    if not (
+        torch.cuda.is_available()
+        and torch.cuda.get_device_capability()[0] == 9
+        and _EXT_AVAILABLE
+        and hasattr(_C, "rope_apply_sm90")
+    ):
+        return None
+    from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op
+
+    return RoPESM90Op()
+
+
+# (batch, seq)
+DEFAULT_CONFIGS = [
+    (8, 512),
+    (8, 2048),
+    (16, 4096),
+    (8, 8192),
+]
+
+
+def _make_inputs(batch, seq, device, dtype):
+    x = torch.randn(batch, N_HEADS, seq, HEAD_DIM, device=device, dtype=dtype)
+    positions = torch.arange(seq, device=device, dtype=torch.long)
+    return x, positions
+
+
+def _time_ms(fn, warmup, iters):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def _peak_vram_gb(fn, warmup=3, iters=5):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    baseline = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+    for _ in range(iters):
+        fn()
+    torch.cuda.synchronize()
+    return (torch.cuda.max_memory_allocated() - baseline) / (1024**3)
+
+
+def run_benchmark(args):
+    if device_ctx.device_type not in ["cuda", "xpu", "hip"]:
+        raise RuntimeError("rope benchmark requires a compatible GPU device.")
+
+    device = device_ctx.device
+    dtype = torch.bfloat16
+    native = NativeRoPEOp()
+    triton_op = TritonRoPEOp()
+    sm90_op = _maybe_sm90_op()
+
+    logger.info(
+        f"rope benchmark on {device} (dtype={dtype}); "
+        f"SM90 backend {'enabled' if sm90_op is not None else 'unavailable'}"
+    )
+
+    rows = []
+    for batch, seq in args.configs:
+        x, positions = _make_inputs(batch, seq, device, dtype)
+
+        def fwd(op, x=x, pos=positions):
+            with torch.no_grad():
+                op(x, pos, theta=THETA)
+
+        def fwd_bwd(op, x_src=x, pos=positions):
+            x_in = x_src.clone().requires_grad_(True)
+            op(x_in, pos, theta=THETA).sum().backward()
+
+        n_fwd = _time_ms(lambda: fwd(native), args.warmup, args.iters)
+        t_fwd = _time_ms(lambda: fwd(triton_op), args.warmup, args.iters)
+        n_fb = _time_ms(lambda: fwd_bwd(native), args.warmup, args.iters)
+        t_fb = _time_ms(lambda: fwd_bwd(triton_op), args.warmup, args.iters)
+        n_vram = _peak_vram_gb(lambda: fwd(native))
+        t_vram = _peak_vram_gb(lambda: fwd(triton_op))
+
+        row = [
+            f"{batch}x{seq}",
+            f"{n_fwd:.3f}",
+            f"{t_fwd:.3f}",
+            f"{n_fwd/t_fwd:.2f}x",
+            f"{n_fb:.3f}",
+            f"{t_fb:.3f}",
+            f"{n_fb/t_fb:.2f}x",
+            f"{n_vram*1024:.0f}",
+            f"{t_vram*1024:.0f}",
+        ]
+        if sm90_op is not None:
+            s_fwd = _time_ms(lambda: fwd(sm90_op), args.warmup, args.iters)
+            s_fb = _time_ms(lambda: fwd_bwd(sm90_op), args.warmup, args.iters)
+            s_vram = _peak_vram_gb(lambda: fwd(sm90_op))
+            row += [
+                f"{s_fwd:.3f}",
+                f"{n_fwd/s_fwd:.2f}x",
+                f"{t_fwd/s_fwd:.2f}x",
+                f"{s_fb:.3f}",
+                f"{s_vram*1024:.0f}",
+            ]
+        rows.append(row)
+
+    headers = [
+        "shape (B x S)",
+        "native fwd ms",
+        "triton fwd ms",
+        "fwd speedup",
+        "native f+b ms",
+        "triton f+b ms",
+        "f+b speedup",
+        "native fwd MB",
+        "triton fwd MB",
+    ]
+    if sm90_op is not None:
+        headers += [
+            "sm90 fwd ms",
+            "sm90 vs native",
+            "sm90 vs triton",
+            "sm90 f+b ms",
+            "sm90 fwd MB",
+        ]
+    print(tabulate(rows, headers=headers, tablefmt="github"))
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--iters", type=int, default=20)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument(
+        "--configs",
+        type=str,
+        default=None,
+        help="Semicolon-separated 'batch,seq' tuples, e.g. '8,512;16,4096'.",
+    )
+    args = parser.parse_args()
+    if args.configs:
+        args.configs = [tuple(int(x) for x in tup.split(",")) for tup in args.configs.split(";")]
+    else:
+        args.configs = DEFAULT_CONFIGS
+    return args
+
+
+if __name__ == "__main__":
+    run_benchmark(parse_args())

csrc/cuda/rope_sm90.cu

@@ -0,0 +1,102 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// CUDA RoPE kernel for SM90 (GPT-NeoX / HF rotate-half), matching NativeRoPEOp.
+//
+// For a row (one [B, H, S] token vector of width D, half = D / 2) at sequence
+// index s = row % S and pair index i in [0, half):
+//
+//   c = cos[s, i]              (fp32, precomputed to match the reference math)
+//   sn = sin[s, i] * sin_sign  (sin_sign = +1 forward, -1 backward)
+//   out[i]        = x[i] * c - x[i + half] * sn
+//   out[i + half] = x[i + half] * c + x[i] * sn
+//
+// The elementwise rotation is done in fp32 and rounded back to the input dtype.
+// Backward reuses the same kernel with sin_sign = -1 (RoPE is an orthogonal
+// per-position rotation, so grad_x = grad_out * cos - rotate_half(grad_out) * sin).
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+namespace {
+
+template <typename scalar_t>
+__global__ void rope_apply_sm90_kernel(
+    const scalar_t* __restrict__ x,   // [n_rows, D]
+    const float* __restrict__ cos,    // [S, half]
+    const float* __restrict__ sin,    // [S, half]
+    scalar_t* __restrict__ out,       // [n_rows, D]
+    const int64_t n_rows,
+    const int S,
+    const int half,
+    const float sin_sign) {
+  const int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
+  const int64_t total = n_rows * static_cast<int64_t>(half);
+  if (idx >= total) {
+    return;
+  }
+
+  const int64_t row = idx / half;
+  const int i = static_cast<int>(idx % half);
+  const int seq = static_cast<int>(row % S);
+
+  const float c = cos[seq * half + i];
+  const float sn = sin[seq * half + i] * sin_sign;
+
+  const int64_t base = row * (2LL * half);
+  const float x1 = static_cast<float>(x[base + i]);
+  const float x2 = static_cast<float>(x[base + i + half]);
+
+  out[base + i] = static_cast<scalar_t>(x1 * c - x2 * sn);
+  out[base + i + half] = static_cast<scalar_t>(x2 * c + x1 * sn);
+}
+
+}  // namespace
+
+// x: [n_rows, D] contiguous (any float dtype); cos/sin: [S, half] fp32 contiguous.
+torch::Tensor rope_apply_sm90(
+    torch::Tensor x,
+    torch::Tensor cos,
+    torch::Tensor sin,
+    double sin_sign) {
+  TORCH_CHECK(x.is_cuda(), "rope: x must be a CUDA tensor");
+  TORCH_CHECK(x.dim() == 2, "rope: x must be 2-D [n_rows, D]");
+  TORCH_CHECK(x.is_contiguous(), "rope: x must be contiguous");
+  TORCH_CHECK(cos.is_cuda() && sin.is_cuda(), "rope: cos/sin must be CUDA tensors");
+  TORCH_CHECK(cos.scalar_type() == torch::kFloat32 && sin.scalar_type() == torch::kFloat32,
+              "rope: cos/sin must be fp32");
+  TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), "rope: cos/sin must be contiguous");
+
+  const int64_t n_rows = x.size(0);
+  const int64_t D = x.size(1);
+  TORCH_CHECK(D % 2 == 0, "rope: head_dim must be even");
+  const int half = static_cast<int>(D / 2);
+  const int S = static_cast<int>(cos.size(0));
+  TORCH_CHECK(cos.size(1) == half && sin.size(1) == half,
+              "rope: cos/sin last dim must equal head_dim/2");
+  TORCH_CHECK(S > 0 && n_rows % S == 0,
+              "rope: n_rows must be divisible by seq length S");
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+  auto out = torch::empty_like(x);
+
+  const int64_t total = n_rows * static_cast<int64_t>(half);
+  const int threads = 256;
+  const int64_t blocks = (total + threads - 1) / threads;
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "rope_apply_sm90", [&] {
+        rope_apply_sm90_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
+            x.data_ptr<scalar_t>(),
+            cos.data_ptr<float>(),
+            sin.data_ptr<float>(),
+            out.data_ptr<scalar_t>(),
+            n_rows,
+            S,
+            half,
+            static_cast<float>(sin_sign));
+      });
+  return out;
+}

csrc/ops.cpp

@@ -58,6 +58,8 @@ torch::Tensor linear_logp_logits_bf16_to_dlogits(torch::Tensor logits,
                                                  torch::Tensor grad_logp,
                                                  torch::Tensor lse,
                                                  int64_t vocab_start_index);
+// RoPE (rotate-half) apply for SM90; cos/sin precomputed fp32, sin_sign = +1 fwd / -1 bwd.
+torch::Tensor rope_apply_sm90(torch::Tensor x, torch::Tensor cos, torch::Tensor sin, double sin_sign);
 #endif
 
 #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)
@@ -150,6 +152,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           "In-place local bf16 probs -> TP dlogits for selected log-prob backward");
     m.def("linear_logp_logits_bf16_to_dlogits", &linear_logp_logits_bf16_to_dlogits,
           "Build bf16 dlogits from bf16 logits and fp32 lse");
+
+    // RoPE rotate-half apply, SM90 (forward and backward share the kernel via sin_sign)
+    m.def("rope_apply_sm90", &rope_apply_sm90, "RoPE rotate-half apply (GPT-NeoX), SM90");
 #endif
 
 #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)

rl_engine/kernels/gtest/operator_specs.py

@@ -57,6 +57,18 @@ def _load_object(path: str) -> Any:
         },
         grad_input_names=("hidden", "lm_head_weight"),
     ),
+    "rope": OperatorSpec(
+        name="rope",
+        op_class="elementwise",
+        gold_path="rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp",
+            "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp",
+            "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op",
+        },
+        grad_input_names=("x",),
+    ),
 }
 
 

rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op
+
+__all__ = ["RoPESM90Op"]