P3: CUDA RoPE kernel — single launch per call (vs 5-6 PyTorch ops)
New files: - dsv4/kernels/cuda/rope_cuda.cu: GPT-J interleaved RoPE kernel (forward+inverse) - dsv4/ops/rope_cuda.py: Python bridge with ctypes loading - tests/unit/test_rope_cuda.py: correctness test (cos >= 0.999998) Savings: ~915 launches/token → 183 launches/token
This commit is contained in:
92
dsv4/kernels/cuda/rope_cuda.cu
Normal file
92
dsv4/kernels/cuda/rope_cuda.cu
Normal file
@@ -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 <cuda.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
|
||||
__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<cudaStream_t>(stream_ptr);
|
||||
apply_rope_kernel<<<grid_size, block_size, 0, stream>>>(
|
||||
static_cast<__nv_bfloat16*>(x_ptr),
|
||||
positions_ptr,
|
||||
cos_ptr,
|
||||
sin_ptr,
|
||||
T, n_h, hd, nope_dim, rope_dim, inverse
|
||||
);
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
93
dsv4/ops/rope_cuda.py
Normal file
93
dsv4/ops/rope_cuda.py
Normal file
@@ -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
|
||||
126
tests/unit/test_rope_cuda.py
Normal file
126
tests/unit/test_rope_cuda.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user