From 61944dd49603696a15afb94fd50d240cef8ae0c2 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sat, 18 Jul 2026 03:43:14 +0000 Subject: [PATCH 1/3] add rope kernel on sm120, now switch to sm90 Signed-off-by: Zhang Jian --- csrc/ops.cpp | 5 + csrc/rope_kernel.cu | 102 ++++++++++++ rl_engine/kernels/gtest/operator_specs.py | 12 ++ .../ops/cuda/rotary_embedding/__init__.py | 6 + .../kernels/ops/cuda/rotary_embedding/rope.py | 90 +++++++++++ .../ops/triton/rotary_embedding/__init__.py | 6 + .../ops/triton/rotary_embedding/rope.py | 153 ++++++++++++++++++ rl_engine/kernels/registry.py | 10 +- setup.py | 1 + 9 files changed, 383 insertions(+), 2 deletions(-) create mode 100644 csrc/rope_kernel.cu create mode 100644 rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py create mode 100644 rl_engine/kernels/ops/cuda/rotary_embedding/rope.py create mode 100644 rl_engine/kernels/ops/triton/rotary_embedding/__init__.py create mode 100644 rl_engine/kernels/ops/triton/rotary_embedding/rope.py diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3..848b2ac 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -61,6 +61,8 @@ torch::Tensor linear_logp_logits_bf16_to_dlogits(torch::Tensor logits, #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) +// RoPE (rotate-half) apply; cos/sin precomputed fp32, sin_sign = +1 fwd / -1 bwd. +torch::Tensor rope_apply_cuda(torch::Tensor x, torch::Tensor cos, torch::Tensor sin, double sin_sign); torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output); torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor fused_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); @@ -167,6 +169,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); + // RoPE rotate-half apply (forward and backward share the kernel via sin_sign) + m.def("rope_apply", &rope_apply_cuda, "RoPE rotate-half apply (GPT-NeoX), CUDA"); + // registry Prefix-Shared Attention m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); #endif diff --git a/csrc/rope_kernel.cu b/csrc/rope_kernel.cu new file mode 100644 index 0000000..8d279fe --- /dev/null +++ b/csrc/rope_kernel.cu @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// CUDA RoPE kernel (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 +#include +#include + +namespace { + +template +__global__ void rope_apply_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(blockDim.x) + threadIdx.x; + const int64_t total = n_rows * static_cast(half); + if (idx >= total) { + return; + } + + const int64_t row = idx / half; + const int i = static_cast(idx % half); + const int seq = static_cast(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(x[base + i]); + const float x2 = static_cast(x[base + i + half]); + + out[base + i] = static_cast(x1 * c - x2 * sn); + out[base + i + half] = static_cast(x2 * c + x1 * sn); +} + +} // namespace + +// x: [n_rows, D] contiguous (any float dtype); cos/sin: [S, half] fp32 contiguous. +torch::Tensor rope_apply_cuda( + 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(D / 2); + const int S = static_cast(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(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_cuda", [&] { + rope_apply_kernel<<>>( + x.data_ptr(), + cos.data_ptr(), + sin.data_ptr(), + out.data_ptr(), + n_rows, + S, + half, + static_cast(sin_sign)); + }); + return out; +} diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 55a4a20..1aa9740 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/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": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPECudaOp", + }, + grad_input_names=("x",), + ), } diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py b/rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py new file mode 100644 index 0000000..60cc16f --- /dev/null +++ b/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 RoPECudaOp + +__all__ = ["RoPECudaOp"] diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py new file mode 100644 index 0000000..a05bba4 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Custom CUDA RoPE op (GPT-NeoX rotate-half), matching NativeRoPEOp. + +cos/sin are built in fp32 with the exact reference math and passed to a small +CUDA kernel (``_C.rope_apply``) that does the per-position rotation. Backward +reuses the same kernel with the sine negated (RoPE is an orthogonal rotation, so +``grad_x = grad_out * cos - rotate_half(grad_out) * sin``). +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.utils.logger import logger + + +def _build_cos_sin(positions: Tensor, half: int, theta: float, device: torch.device): + """fp32 cos/sin caches of shape [S, half], identical math to NativeRoPEOp.""" + inv_freq = 1.0 / ( + theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half) + ) + pos = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) + freqs = pos * inv_freq # [S, half] + return freqs.cos().contiguous(), freqs.sin().contiguous() + + +class _RoPEFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: + D = x.shape[-1] + if D % 2 != 0: + raise ValueError(f"RoPE head_dim must be even, got {D}") + if positions.dim() != 1: + raise NotImplementedError( + "CUDA RoPE currently supports 1-D positions [S] (shared across batch)." + ) + S = positions.shape[0] + x_2d = x.contiguous().reshape(-1, D) + n_rows = x_2d.shape[0] + if n_rows % S != 0: + raise ValueError( + f"row count {n_rows} not divisible by seq length {S}; " + "expected a [..., S, D] contiguous layout." + ) + cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) + ctx.save_for_backward(cos, sin) + out = _C.rope_apply(x_2d, cos, sin, 1.0) + return out.reshape(x.shape) + + @staticmethod + def backward(ctx, grad_out: Tensor): + cos, sin = ctx.saved_tensors + grad_x = None + if ctx.needs_input_grad[0]: + D = grad_out.shape[-1] + g_2d = grad_out.contiguous().reshape(-1, D) + # Inverse rotation: same kernel with the sine negated. + grad_x = _C.rope_apply(g_2d, cos, sin, -1.0).reshape(grad_out.shape) + # Inputs: x, positions, theta. + return grad_x, None, None + + +class RoPECudaOp: + """Custom CUDA RoPE op (GPT-NeoX rotate-half), differentiable w.r.t. ``x``. + + Qwen3 defaults: theta=1e6, head_dim=128, full-dimension rotation. cos/sin are + computed in fp32 from ``positions`` and ``theta`` (matching the reference); the + rotation runs in the precompiled ``_C.rope_apply`` CUDA kernel. + """ + + op_class = "elementwise" + + def __init__(self) -> None: + if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply"): + raise RuntimeError( + "CUDA RoPE kernel 'rope_apply' is not compiled into _C. " + "Rebuild the extension with 'pip install -e .'." + ) + logger.info("Successfully linked to precompiled _C.rope_apply kernel.") + + def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + return self.forward(x, positions, theta=theta) + + def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + if x.device.type != "cuda": + raise RuntimeError(f"RoPECudaOp requires a CUDA tensor, got device '{x.device}'.") + return _RoPEFunction.apply(x, positions, theta) diff --git a/rl_engine/kernels/ops/triton/rotary_embedding/__init__.py b/rl_engine/kernels/ops/triton/rotary_embedding/__init__.py new file mode 100644 index 0000000..abc9387 --- /dev/null +++ b/rl_engine/kernels/ops/triton/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.triton.rotary_embedding.rope import TritonRoPEOp + +__all__ = ["TritonRoPEOp"] diff --git a/rl_engine/kernels/ops/triton/rotary_embedding/rope.py b/rl_engine/kernels/ops/triton/rotary_embedding/rope.py new file mode 100644 index 0000000..317a09b --- /dev/null +++ b/rl_engine/kernels/ops/triton/rotary_embedding/rope.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Triton RoPE kernel (GPT-NeoX / HF rotate-half), matching NativeRoPEOp. + +For a token at absolute position ``p`` and head_dim index ``d`` (half = D // 2):: + + inv_freq[i] = theta ** (-i / half) i in [0, half) + angle = p * inv_freq[d % half] + out[d=half] = x[d] * cos(angle) + x[d-half] * sin(angle) + +cos/sin are built in fp32 with the *exact* reference math (``theta ** x``) so the +fp32 path stays bit-close to the gold even at large positions where cos/sin are +numerically sensitive; the Triton kernel does the elementwise rotation in fp32 +and rounds back to the input dtype on store. + +RoPE is a per-position orthogonal rotation, so the input gradient is the same +rotation with the sine negated:: + + grad_x = grad_out * cos - rotate_half(grad_out) * sin + +which the same kernel produces when called with ``sin_sign = -1``. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from torch import Tensor + + +@triton.jit +def _rope_kernel( + x_ptr, # [n_rows, D] flattened input + cos_ptr, # [S, HALF] fp32 cosine cache + sin_ptr, # [S, HALF] fp32 sine cache + out_ptr, # [n_rows, D] output + n_rows, + S, + SIN_SIGN: tl.constexpr, # +1.0 forward, -1.0 backward + HALF: tl.constexpr, # D // 2, power-of-two block width + stride_row, + stride_d, +): + """One program per row (a single [B, H, S] token vector of width D).""" + row = tl.program_id(0) + if row >= n_rows: + return + + # Row layout is [..., S, D] contiguous, so the sequence index is row % S. + seq_idx = row % S + d = tl.arange(0, HALF) + cos = tl.load(cos_ptr + seq_idx * HALF + d) + sin = tl.load(sin_ptr + seq_idx * HALF + d) * SIN_SIGN + + base = row * stride_row + x1 = tl.load(x_ptr + base + d * stride_d).to(tl.float32) + x2 = tl.load(x_ptr + base + (d + HALF) * stride_d).to(tl.float32) + + out1 = x1 * cos - x2 * sin + out2 = x2 * cos + x1 * sin + + out_dtype = out_ptr.dtype.element_ty + tl.store(out_ptr + base + d * stride_d, out1.to(out_dtype)) + tl.store(out_ptr + base + (d + HALF) * stride_d, out2.to(out_dtype)) + + +def _build_cos_sin(positions: Tensor, half: int, theta: float, device: torch.device): + """fp32 cos/sin caches of shape [S, half], identical math to NativeRoPEOp.""" + inv_freq = 1.0 / ( + theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half) + ) + pos = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) + freqs = pos * inv_freq # [S, half] + return freqs.cos().contiguous(), freqs.sin().contiguous() + + +def _launch_rope( + x: Tensor, cos: Tensor, sin: Tensor, S: int, sin_sign: float +) -> Tensor: + D = x.shape[-1] + half = D // 2 + x_2d = x.contiguous().reshape(-1, D) + n_rows = x_2d.shape[0] + + out = torch.empty_like(x_2d) + grid = (n_rows,) + _rope_kernel[grid]( + x_2d, + cos, + sin, + out, + n_rows, + S, + SIN_SIGN=float(sin_sign), + HALF=half, + stride_row=x_2d.stride(0), + stride_d=x_2d.stride(1), + ) + return out.reshape(x.shape) + + +class _RoPEFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: + D = x.shape[-1] + if D % 2 != 0: + raise ValueError(f"RoPE head_dim must be even, got {D}") + if positions.dim() != 1: + raise NotImplementedError( + "Triton RoPE currently supports 1-D positions [S] (shared across batch)." + ) + S = positions.shape[0] + n_rows = x.numel() // D + if n_rows % S != 0: + raise ValueError( + f"row count {n_rows} not divisible by seq length {S}; " + "expected a [..., S, D] contiguous layout." + ) + cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) + ctx.save_for_backward(cos, sin) + ctx.seq_len = S + return _launch_rope(x, cos, sin, S, sin_sign=1.0) + + @staticmethod + def backward(ctx, grad_out: Tensor): + cos, sin = ctx.saved_tensors + grad_x = None + if ctx.needs_input_grad[0]: + # Inverse rotation: same kernel with the sine negated. + grad_x = _launch_rope(grad_out, cos, sin, ctx.seq_len, sin_sign=-1.0) + # Inputs: x, positions, theta. + return grad_x, None, None + + +class TritonRoPEOp: + """Triton RoPE op (GPT-NeoX rotate-half), differentiable w.r.t. ``x``. + + Qwen3 defaults: theta=1e6, head_dim=128, full-dimension rotation. cos/sin are + computed in fp32 from ``positions`` and ``theta`` (matching the reference) and + the rotation runs in a Triton kernel -- no external cos/sin cache is accepted. + """ + + op_class = "elementwise" + + def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + return self.forward(x, positions, theta=theta) + + def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + if x.device.type not in ("cuda", "hip", "xpu"): + raise RuntimeError(f"TritonRoPEOp requires a GPU tensor, got device '{x.device}'.") + return _RoPEFunction.apply(x, positions, theta) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e..62baf26 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -59,6 +59,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp" PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" + TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" + CUDA_ROPE = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPECudaOp" PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" @@ -176,7 +178,11 @@ def __init__(self): "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], # Default dispatch logic for new operators "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], - "rope": [OpBackend.PYTORCH_NATIVE_ROPE], + "rope": [ + OpBackend.CUDA_ROPE, + OpBackend.TRITON_ROPE, + OpBackend.PYTORCH_NATIVE_ROPE, + ], }, "rocm": { "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], @@ -190,7 +196,7 @@ def __init__(self): "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], - "rope": [OpBackend.PYTORCH_NATIVE_ROPE], + "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], diff --git a/setup.py b/setup.py index a4bd3b6..513859c 100644 --- a/setup.py +++ b/setup.py @@ -78,6 +78,7 @@ def get_extensions(): "csrc/ops.cpp", "csrc/fused_logp_kernel.cu", "csrc/deterministic_logp_kernel.cu", + "csrc/rope_kernel.cu", "csrc/cuda/attention/prefix_shared_attention.cu", ] From 252edd5ba172aaa7565c9d8a6ce2c6742d759392 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sat, 18 Jul 2026 04:44:53 +0000 Subject: [PATCH 2/3] tested on sm90 Signed-off-by: Zhang Jian --- csrc/{rope_kernel.cu => cuda/rope_sm90.cu} | 10 ++++---- csrc/ops.cpp | 10 ++++---- rl_engine/kernels/gtest/operator_specs.py | 2 +- .../ops/cuda/rotary_embedding/__init__.py | 4 ++-- .../kernels/ops/cuda/rotary_embedding/rope.py | 24 +++++++++---------- rl_engine/kernels/registry.py | 4 ++-- setup.py | 2 +- 7 files changed, 28 insertions(+), 28 deletions(-) rename csrc/{rope_kernel.cu => cuda/rope_sm90.cu} (93%) diff --git a/csrc/rope_kernel.cu b/csrc/cuda/rope_sm90.cu similarity index 93% rename from csrc/rope_kernel.cu rename to csrc/cuda/rope_sm90.cu index 8d279fe..ce27a32 100644 --- a/csrc/rope_kernel.cu +++ b/csrc/cuda/rope_sm90.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright (c) 2026 RL-Kernel Contributors // -// CUDA RoPE kernel (GPT-NeoX / HF rotate-half), matching NativeRoPEOp. +// 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): @@ -22,7 +22,7 @@ namespace { template -__global__ void rope_apply_kernel( +__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] @@ -55,7 +55,7 @@ __global__ void rope_apply_kernel( } // namespace // x: [n_rows, D] contiguous (any float dtype); cos/sin: [S, half] fp32 contiguous. -torch::Tensor rope_apply_cuda( +torch::Tensor rope_apply_sm90( torch::Tensor x, torch::Tensor cos, torch::Tensor sin, @@ -87,8 +87,8 @@ torch::Tensor rope_apply_cuda( auto stream = at::cuda::getCurrentCUDAStream(); AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "rope_apply_cuda", [&] { - rope_apply_kernel<<>>( + at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "rope_apply_sm90", [&] { + rope_apply_sm90_kernel<<>>( x.data_ptr(), cos.data_ptr(), sin.data_ptr(), diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 848b2ac..6bf9f72 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -58,11 +58,11 @@ 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) -// RoPE (rotate-half) apply; cos/sin precomputed fp32, sin_sign = +1 fwd / -1 bwd. -torch::Tensor rope_apply_cuda(torch::Tensor x, torch::Tensor cos, torch::Tensor sin, double sin_sign); torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output); torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor fused_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); @@ -152,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) @@ -169,9 +172,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); - // RoPE rotate-half apply (forward and backward share the kernel via sin_sign) - m.def("rope_apply", &rope_apply_cuda, "RoPE rotate-half apply (GPT-NeoX), CUDA"); - // registry Prefix-Shared Attention m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); #endif diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 1aa9740..ad49d1f 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -65,7 +65,7 @@ def _load_object(path: str) -> Any: candidate_paths={ "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp", - "cuda": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPECudaOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op", }, grad_input_names=("x",), ), diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py b/rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py index 60cc16f..5cc003f 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPECudaOp +from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op -__all__ = ["RoPECudaOp"] +__all__ = ["RoPESM90Op"] diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py index a05bba4..032fe45 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Custom CUDA RoPE op (GPT-NeoX rotate-half), matching NativeRoPEOp. +"""Custom CUDA RoPE op for SM90 (GPT-NeoX rotate-half), matching NativeRoPEOp. cos/sin are built in fp32 with the exact reference math and passed to a small -CUDA kernel (``_C.rope_apply``) that does the per-position rotation. Backward +CUDA kernel (``_C.rope_apply_sm90``) that does the per-position rotation. Backward reuses the same kernel with the sine negated (RoPE is an orthogonal rotation, so ``grad_x = grad_out * cos - rotate_half(grad_out) * sin``). """ @@ -47,7 +47,7 @@ def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: ) cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) ctx.save_for_backward(cos, sin) - out = _C.rope_apply(x_2d, cos, sin, 1.0) + out = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) return out.reshape(x.shape) @staticmethod @@ -58,33 +58,33 @@ def backward(ctx, grad_out: Tensor): D = grad_out.shape[-1] g_2d = grad_out.contiguous().reshape(-1, D) # Inverse rotation: same kernel with the sine negated. - grad_x = _C.rope_apply(g_2d, cos, sin, -1.0).reshape(grad_out.shape) + grad_x = _C.rope_apply_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) # Inputs: x, positions, theta. return grad_x, None, None -class RoPECudaOp: - """Custom CUDA RoPE op (GPT-NeoX rotate-half), differentiable w.r.t. ``x``. +class RoPESM90Op: + """Custom CUDA RoPE op for SM90 (GPT-NeoX rotate-half), differentiable w.r.t. ``x``. Qwen3 defaults: theta=1e6, head_dim=128, full-dimension rotation. cos/sin are computed in fp32 from ``positions`` and ``theta`` (matching the reference); the - rotation runs in the precompiled ``_C.rope_apply`` CUDA kernel. + rotation runs in the precompiled ``_C.rope_apply_sm90`` CUDA kernel. """ op_class = "elementwise" def __init__(self) -> None: - if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply"): + if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply_sm90"): raise RuntimeError( - "CUDA RoPE kernel 'rope_apply' is not compiled into _C. " - "Rebuild the extension with 'pip install -e .'." + "CUDA RoPE kernel 'rope_apply_sm90' is not compiled into _C. " + "Rebuild the extension with 'KERNEL_ALIGN_FORCE_SM90=1 pip install -e .'." ) - logger.info("Successfully linked to precompiled _C.rope_apply kernel.") + logger.info("Successfully linked to precompiled _C.rope_apply_sm90 kernel.") def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: return self.forward(x, positions, theta=theta) def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: if x.device.type != "cuda": - raise RuntimeError(f"RoPECudaOp requires a CUDA tensor, got device '{x.device}'.") + raise RuntimeError(f"RoPESM90Op requires a CUDA tensor, got device '{x.device}'.") return _RoPEFunction.apply(x, positions, theta) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 62baf26..a0d7f4f 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -60,7 +60,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp" PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" - CUDA_ROPE = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPECudaOp" + CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" @@ -179,7 +179,7 @@ def __init__(self): # Default dispatch logic for new operators "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rope": [ - OpBackend.CUDA_ROPE, + OpBackend.CUDA_ROPE_SM90, OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE, ], diff --git a/setup.py b/setup.py index 513859c..f4568da 100644 --- a/setup.py +++ b/setup.py @@ -78,7 +78,6 @@ def get_extensions(): "csrc/ops.cpp", "csrc/fused_logp_kernel.cu", "csrc/deterministic_logp_kernel.cu", - "csrc/rope_kernel.cu", "csrc/cuda/attention/prefix_shared_attention.cu", ] @@ -140,6 +139,7 @@ def get_extensions(): sm90_srcs = [ "csrc/cuda/fused_logp_sm90.cu", "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build ] enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] From 4d74b4e462f7f2b40cf20cf324df1be18f72cb15 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sat, 18 Jul 2026 04:59:19 +0000 Subject: [PATCH 3/3] add dedicated benchmark Signed-off-by: Zhang Jian --- benchmarks/benchmark_rope.py | 191 +++++++++++++++++++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 benchmarks/benchmark_rope.py diff --git a/benchmarks/benchmark_rope.py b/benchmarks/benchmark_rope.py new file mode 100644 index 0000000..6fd395e --- /dev/null +++ b/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())