ViperEkura
3e67b4f88d
- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call) - RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor - apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit) - Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out - Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16) - Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)
93 lines
2.9 KiB
Plaintext
93 lines
2.9 KiB
Plaintext
#include <torch/extension.h>
|
|
#include <cuda_bf16.h>
|
|
|
|
__global__ void rotary_emb_kernel(
|
|
const __nv_bfloat16* __restrict__ x,
|
|
const float* __restrict__ cos,
|
|
const float* __restrict__ sin,
|
|
__nv_bfloat16* __restrict__ out,
|
|
int batch,
|
|
int seq_len,
|
|
int n_heads,
|
|
int head_dim
|
|
) {
|
|
const int half_dim = head_dim >> 1;
|
|
const int total = batch * seq_len * n_heads * half_dim;
|
|
|
|
for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
idx < total;
|
|
idx += gridDim.x * blockDim.x) {
|
|
|
|
int pair = idx % half_dim;
|
|
int tmp = idx / half_dim;
|
|
int head = tmp % n_heads;
|
|
tmp /= n_heads;
|
|
int seq = tmp % seq_len;
|
|
int b = tmp / seq_len;
|
|
|
|
int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1);
|
|
int cs_offset = (b * seq_len + seq) * half_dim + pair;
|
|
|
|
__nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
|
|
float x_even = __bfloat162float(__low2bfloat16(x_pair));
|
|
float x_odd = __bfloat162float(__high2bfloat16(x_pair));
|
|
|
|
float c = cos[cs_offset];
|
|
float s = sin[cs_offset];
|
|
|
|
float out_even = x_even * c - x_odd * s;
|
|
float out_odd = x_even * s + x_odd * c;
|
|
|
|
__nv_bfloat162 out_pair = __floats2bfloat162_rn(out_even, out_odd);
|
|
*reinterpret_cast<__nv_bfloat162*>(out + x_offset) = out_pair;
|
|
}
|
|
}
|
|
|
|
torch::Tensor rotary_emb(
|
|
torch::Tensor x,
|
|
torch::Tensor cos,
|
|
torch::Tensor sin
|
|
) {
|
|
|
|
int batch = x.size(0);
|
|
int seq_len = x.size(1);
|
|
int n_heads = x.size(2);
|
|
int head_dim = x.size(3);
|
|
|
|
TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
|
|
TORCH_CHECK(cos.is_cuda(), "cos must be on CUDA");
|
|
TORCH_CHECK(sin.is_cuda(), "sin must be on CUDA");
|
|
TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
|
|
TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]");
|
|
TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
|
|
TORCH_CHECK(cos.dim() == 3, "cos must be 3D [batch, seq_len, head_dim/2]");
|
|
TORCH_CHECK(sin.dim() == 3, "sin must be 3D [batch, seq_len, head_dim/2]");
|
|
TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
|
|
|
|
auto out = torch::empty_like(x);
|
|
|
|
int half_dim = head_dim / 2;
|
|
int total = batch * seq_len * n_heads * half_dim;
|
|
int block = 256;
|
|
int grid = std::min((total + block - 1) / block, 1024);
|
|
|
|
rotary_emb_kernel<<<grid, block>>>(
|
|
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
|
|
cos.data_ptr<float>(),
|
|
sin.data_ptr<float>(),
|
|
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
|
|
batch, seq_len, n_heads, head_dim
|
|
);
|
|
|
|
return out;
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("rotary_emb", &rotary_emb,
|
|
py::arg("x"),
|
|
py::arg("cos"),
|
|
py::arg("sin"),
|
|
"Fused rotary embedding (bf16 x, f32 cos/sin, bf16 out)"
|
|
);
|
|
}
|