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()) diff --git a/csrc/cuda/rope_sm90.cu b/csrc/cuda/rope_sm90.cu new file mode 100644 index 0000000..ce27a32 --- /dev/null +++ b/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 +#include +#include + +namespace { + +template +__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(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_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(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_sm90", [&] { + rope_apply_sm90_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/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3..6bf9f72 100644 --- a/csrc/ops.cpp +++ b/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) diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 55a4a20..ad49d1f 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-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 new file mode 100644 index 0000000..5cc003f --- /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 RoPESM90Op + +__all__ = ["RoPESM90Op"] 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..032fe45 --- /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 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_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``). +""" + +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_sm90(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_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) + # Inputs: x, positions, theta. + return grad_x, None, None + + +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_sm90`` CUDA kernel. + """ + + op_class = "elementwise" + + def __init__(self) -> None: + if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply_sm90"): + raise RuntimeError( + "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_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"RoPESM90Op 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..a0d7f4f 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_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" @@ -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_SM90, + 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 6f94e04..7156263 100644 --- a/setup.py +++ b/setup.py @@ -139,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)]