ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: unify rotary embedding interface and update docs

Browse Source 
- Merge cos/sin into single freqs_cis tensor [batch, seq, dim/2, 2] throughout the pipeline: RotaryEmbedding buffer, forward return type, apply_rotary_emb signature, CUDA kernel interface
- CUDA kernel now takes freqs_cis directly and reads cos/sin via stride offset internally, eliminating Python-side slice/copy overhead
- Kernel interface: rotary_emb(x, freqs_cis) replaces rotary_emb(x, cos, sin)
- All call sites pass rotary_emb as Tensor (was tuple), type annotations consistent
- Update build threads from 8 to 16
- Fix all docs: get-started, inference, training, cuda_kernels, architecture, internals — reflect new rotary interface, KVCache fields, rotary backend dispatch, .so path, kernel registry count, file layout
This commit is contained in:
ViperEkura
2026-07-31 16:52:25 +08:00
parent 75411ce0cc
commit 7aa5ed09d9
11 changed files with 120 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/rotary_emb.cu
+15 -20
View File
@@ -3,8 +3,7 @@
__global__ void rotary_emb_kernel(
const __nv_bfloat16* __restrict__ x,
const float* __restrict__ cos,
const float* __restrict__ sin,
const float* __restrict__ freqs_cis,
__nv_bfloat16* __restrict__ out,
int batch,
int seq_len,
@@ -26,14 +25,14 @@ __global__ void rotary_emb_kernel(
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;
int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2;
__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 c = freqs_cis[cs_offset];
float s = freqs_cis[cs_offset + 1];
float out_even = x_even * c - x_odd * s;
float out_odd = x_even * s + x_odd * c;
@@ -45,23 +44,21 @@ __global__ void rotary_emb_kernel(
torch::Tensor rotary_emb(
torch::Tensor x,
torch::Tensor cos,
torch::Tensor sin
torch::Tensor freqs_cis
) {
TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis 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(freqs_cis.dim() == 4, "freqs_cis must be 4D [batch, seq_len, dim/2, 2]");
TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
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);
@@ -73,8 +70,7 @@ torch::Tensor rotary_emb(
rotary_emb_kernel<<<grid, block>>>(
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
cos.data_ptr<float>(),
sin.data_ptr<float>(),
freqs_cis.data_ptr<float>(),
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
batch, seq_len, n_heads, head_dim
);
@@ -85,8 +81,7 @@ torch::Tensor rotary_emb(
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)"
py::arg("freqs_cis"),
"Fused rotary embedding (bf16 x, f32 freqs_cis [b,s,d/2,2], bf16 out)"
);
}