diff --git a/dsv4/kernels/cuda/rope_cuda.cu b/dsv4/kernels/cuda/rope_cuda.cu new file mode 100644 index 00000000..c4599c0f --- /dev/null +++ b/dsv4/kernels/cuda/rope_cuda.cu @@ -0,0 +1,92 @@ +/* + * rope_cuda.cu + * + * Fused forward/inverse partial RoPE kernel for DeepSeek V4. + * GPT-J style (interleaved) RoPE on last rope_dim=64 dims of each head. + * + * Replaces 5-6 PyTorch kernel launches per RoPE call with 1 CUDA kernel. + * Total savings: ~1000 launches/token → 183 launches/token (~0.8ms at 2µs/launch). + * + * C API for ctypes loading (no ATen/pybind11). + */ + +#include +#include +#include +#include + +__global__ void apply_rope_kernel( + __nv_bfloat16* __restrict__ x, // (T, n_h, hd) — modified in-place + const int64_t* __restrict__ positions, // (T,) — token positions + const float* __restrict__ cos_cache, // (max_pos, rope_dim//2) + const float* __restrict__ sin_cache, // (max_pos, rope_dim//2) + const int T, + const int n_h, + const int hd, + const int nope_dim, // hd - rope_dim = 448 + const int rope_dim, // 64 + const bool inverse // true = inverse RoPE +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int half_rope = rope_dim / 2; + const int total_pairs = T * n_h * half_rope; + + if (idx >= total_pairs) return; + + const int pair_idx = idx % half_rope; + const int head_idx = (idx / half_rope) % n_h; + const int token_idx = idx / (half_rope * n_h); + + // Get position and cos/sin values + int64_t pos = positions[token_idx]; + float c = cos_cache[pos * half_rope + pair_idx]; + float s = sin_cache[pos * half_rope + pair_idx]; + + // Compute pointer to the two elements of the pair + const int even_offset = token_idx * n_h * hd + head_idx * hd + nope_dim + 2 * pair_idx; + const int odd_offset = even_offset + 1; + + // Load BF16 values, convert to FP32 + float x_even = __bfloat162float(x[even_offset]); + float x_odd = __bfloat162float(x[odd_offset]); + + // Apply rotation + float rot_even, rot_odd; + if (inverse) { + rot_even = x_even * c + x_odd * s; + rot_odd = -x_even * s + x_odd * c; + } else { + rot_even = x_even * c - x_odd * s; + rot_odd = x_even * s + x_odd * c; + } + + // Store back as BF16 + x[even_offset] = __float2bfloat16(rot_even); + x[odd_offset] = __float2bfloat16(rot_odd); +} + +// C API for ctypes +extern "C" { + +void apply_rope_launch( + void* x_ptr, + const int64_t* positions_ptr, + const float* cos_ptr, + const float* sin_ptr, + int T, int n_h, int hd, + int nope_dim, int rope_dim, + bool inverse, + int grid_size, int block_size, + void* stream_ptr +) { + cudaStream_t stream = static_cast(stream_ptr); + apply_rope_kernel<<>>( + static_cast<__nv_bfloat16*>(x_ptr), + positions_ptr, + cos_ptr, + sin_ptr, + T, n_h, hd, nope_dim, rope_dim, inverse + ); +} + +} // extern "C" diff --git a/dsv4/ops/rope_cuda.py b/dsv4/ops/rope_cuda.py new file mode 100644 index 00000000..292b0f84 --- /dev/null +++ b/dsv4/ops/rope_cuda.py @@ -0,0 +1,93 @@ +"""CUDA RoPE kernel — 1 kernel launch per call instead of 5-6 PyTorch ops. + +Uses ctypes to call the compiled kernel directly (no ATen/pybind11). +Same pattern as fmha_multitile_op.py and other production kernels. +""" + +import torch +import ctypes +import subprocess +from pathlib import Path + +_LIB = None + +def _compile_and_load(): + global _LIB + if _LIB is not None: + return _LIB + + cu_path = Path(__file__).parent / "cuda" / "rope_cuda.cu" + assert cu_path.exists(), f"rope_cuda.cu not found at {cu_path}" + + # Compile to shared library + build_dir = Path(__file__).parent / "cuda" / "_build_cache" + build_dir.mkdir(parents=True, exist_ok=True) + so_path = build_dir / "librope_cuda.so" + + if not so_path.exists() or cu_path.stat().st_mtime > so_path.stat().st_mtime: + nvcc = "/usr/local/cuda/bin/nvcc" + cmd = [ + nvcc, "-shared", "-o", str(so_path), str(cu_path), + "-arch=sm_100a", + "--generate-code=arch=compute_100a,code=[sm_100a,compute_100a]", + "-use_fast_math", "-O3", + "-Xcompiler", "-fPIC", + ] + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + raise RuntimeError(f"rope_cuda.cu compilation failed:\n{result.stderr}") + + _LIB = ctypes.CDLL(str(so_path)) + return _LIB + + +def apply_rope(x, positions, cos_cache, sin_cache, rope_dim, inverse=False): + """Apply forward or inverse RoPE in-place using a single CUDA kernel. + + Args: + x: (T, n_h, hd) BF16 — modified in-place + positions: (T,) int64 — token positions + cos_cache: (max_pos, rope_dim//2) float32 + sin_cache: (max_pos, rope_dim//2) float32 + rope_dim: 64 + inverse: True for inverse RoPE + + Returns: + x (modified in-place) + """ + lib = _compile_and_load() + T, n_h, hd = x.shape + nope_dim = hd - rope_dim + half_rope = rope_dim // 2 + + # Ensure types and devices + pos = positions.to(device=x.device, dtype=torch.int64) + assert x.dtype == torch.bfloat16 + assert cos_cache.dtype == torch.float32 + assert sin_cache.dtype == torch.float32 + + # Launch parameters + total_pairs = T * n_h * half_rope + threads = 256 + blocks = (total_pairs + threads - 1) // threads + + # Get raw CUDA stream + stream = torch.cuda.current_stream().cuda_stream + + # Call the kernel + lib.apply_rope_launch( + ctypes.c_void_p(x.data_ptr()), + ctypes.c_void_p(pos.data_ptr()), + ctypes.c_void_p(cos_cache.data_ptr()), + ctypes.c_void_p(sin_cache.data_ptr()), + ctypes.c_int(T), + ctypes.c_int(n_h), + ctypes.c_int(hd), + ctypes.c_int(nope_dim), + ctypes.c_int(rope_dim), + ctypes.c_bool(inverse), + ctypes.c_int(blocks), + ctypes.c_int(threads), + ctypes.c_void_p(stream), + ) + return x diff --git a/tests/unit/test_rope_cuda.py b/tests/unit/test_rope_cuda.py new file mode 100644 index 00000000..b99b575b --- /dev/null +++ b/tests/unit/test_rope_cuda.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Test CUDA RoPE kernel correctness. + +Compare CUDA kernel output vs PyTorch reference. +Must achieve cos >= 0.999998 for production. +""" +import torch +import math +import sys + +def build_rope_cache(max_pos, rope_dim, device, theta=10000.): + freqs = 1. / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim)) + angles = torch.outer(torch.arange(max_pos, dtype=torch.float32), freqs) + return torch.cos(angles).to(device), torch.sin(angles).to(device) + +def apply_rope_ref(x, pos, cos, sin, rope_dim, inverse=False): + """PyTorch reference — the current _apply_rope implementation.""" + T, nh, hd = x.shape + nope = hd - rope_dim + c, s = cos[pos].unsqueeze(1), sin[pos].unsqueeze(1) + xr = x[:, :, nope:] # view + ev = xr[..., 0::2].clone() + od = xr[..., 1::2] + if inverse: + xr[..., 0::2] = (ev * c + od * s).bfloat16() + xr[..., 1::2] = (-ev * s + od * c).bfloat16() + else: + xr[..., 0::2] = (ev * c - od * s).bfloat16() + xr[..., 1::2] = (ev * s + od * c).bfloat16() + return x + +def test_rope_cuda(): + from dsv4.ops.rope_cuda import apply_rope + + device = "cuda:0" + rope_dim = 64 + hd = 512 + n_h = 128 + T = 1 # decode + max_pos = 4096 + + cos, sin = build_rope_cache(max_pos, rope_dim, device) + + # Test forward RoPE + torch.manual_seed(42) + x_ref = torch.randn(T, n_h, hd, dtype=torch.bfloat16, device=device) + x_cuda = x_ref.clone() + positions = torch.tensor([100], dtype=torch.long, device=device) + + apply_rope_ref(x_ref, positions, cos, sin, rope_dim, inverse=False) + apply_rope(x_cuda, positions, cos, sin, rope_dim, inverse=False) + + cos_sim = torch.nn.functional.cosine_similarity( + x_ref.flatten().float(), x_cuda.flatten().float(), dim=0 + ).item() + max_diff = (x_ref.float() - x_cuda.float()).abs().max().item() + + print(f"Forward RoPE (T=1, n_h=128, hd=512):") + print(f" Cosine: {cos_sim:.8f}") + print(f" Max diff: {max_diff:.8f}") + + if cos_sim < 0.999998: + print(f" ❌ FAIL: cosine < 0.999998") + return False + print(f" ✅ PASS") + + # Test inverse RoPE + x_ref_inv = x_ref.clone() + x_cuda_inv = x_cuda.clone() + + apply_rope_ref(x_ref_inv, positions, cos, sin, rope_dim, inverse=True) + apply_rope(x_cuda_inv, positions, cos, sin, rope_dim, inverse=True) + + cos_sim_inv = torch.nn.functional.cosine_similarity( + x_ref_inv.flatten().float(), x_cuda_inv.flatten().float(), dim=0 + ).item() + + print(f"\nInverse RoPE (T=1, n_h=128, hd=512):") + print(f" Cosine: {cos_sim_inv:.8f}") + + if cos_sim_inv < 0.999998: + print(f" ❌ FAIL") + return False + print(f" ✅ PASS") + + # Test round-trip (forward + inverse should be identity) + x_rt = torch.randn(T, n_h, hd, dtype=torch.bfloat16, device=device) + x_orig = x_rt.clone() + apply_rope(x_rt, positions, cos, sin, rope_dim, inverse=False) + apply_rope(x_rt, positions, cos, sin, rope_dim, inverse=True) + + rt_cos = torch.nn.functional.cosine_similarity( + x_orig.flatten().float(), x_rt.flatten().float(), dim=0 + ).item() + print(f"\nRound-trip (forward + inverse):") + print(f" Cosine: {rt_cos:.8f}") + if rt_cos < 0.9999: + print(f" ❌ FAIL: round-trip error too large") + return False + print(f" ✅ PASS") + + # Test multi-token + T2 = 8 + x_ref2 = torch.randn(T2, n_h, hd, dtype=torch.bfloat16, device=device) + x_cuda2 = x_ref2.clone() + pos2 = torch.arange(T2, dtype=torch.long, device=device) + + apply_rope_ref(x_ref2, pos2, cos, sin, rope_dim, inverse=False) + apply_rope(x_cuda2, pos2, cos, sin, rope_dim, inverse=False) + + cos_sim2 = torch.nn.functional.cosine_similarity( + x_ref2.flatten().float(), x_cuda2.flatten().float(), dim=0 + ).item() + print(f"\nMulti-token forward (T=8, n_h=128, hd=512):") + print(f" Cosine: {cos_sim2:.8f}") + if cos_sim2 < 0.999998: + print(f" ❌ FAIL") + return False + print(f" ✅ PASS") + + return True + +if __name__ == "__main__": + torch.manual_seed(42) + success = test_rope_cuda() + sys.exit(0 if success else 1)