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

refactor: standardize packed 3d inference

Browse Source 
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch
This commit is contained in:
ViperEkura
2026-08-16 13:24:02 +08:00
parent 0dd9a417b7
commit 3406157431
15 changed files with 301 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/rotary_emb.cu
+19 -20
View File
@@ -7,13 +7,12 @@ __global__ void rotary_emb_kernel(
const __nv_bfloat16* __restrict__ x,
const float* __restrict__ freqs_cis,
__nv_bfloat16* __restrict__ out,
int batch,
int seq_len,
int n_tokens,
int n_heads,
int head_dim
) {
const int half_dim = head_dim >> 1;
const int total = batch * seq_len * n_heads * half_dim;
const int total = n_tokens * n_heads * half_dim;
for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
idx < total;
@@ -23,11 +22,10 @@ __global__ void rotary_emb_kernel(
int tmp = idx / half_dim;
int head = tmp % n_heads;
tmp /= n_heads;
int seq = tmp % seq_len;
int b = tmp / seq_len;
int token = tmp;
int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1);
int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2;
int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
int cs_offset = (token * half_dim + pair) * 2;
__nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
float x_even = __bfloat162float(__low2bfloat16(x_pair));
@@ -54,27 +52,28 @@ torch::Tensor rotary_emb(
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.dim() == 3 || x.dim() == 4,
"x must be [tokens, n_heads, head_dim] or "
"[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.dim() == x.dim(), "freqs_cis rank must match x rank");
TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
TORCH_CHECK(freqs_cis.scalar_type() == torch::kFloat32, "freqs_cis must be f32");
int batch = x.size(0);
int seq_len = x.size(1);
int n_heads = x.size(2);
int head_dim = x.size(3);
int n_tokens = x.dim() == 3 ? x.size(0) : x.size(0) * x.size(1);
int n_heads = x.size(x.dim() - 2);
int head_dim = x.size(x.dim() - 1);
TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
TORCH_CHECK(freqs_cis.size(0) == batch, "freqs_cis batch mismatch");
TORCH_CHECK(freqs_cis.size(1) == seq_len, "freqs_cis seq_len mismatch");
TORCH_CHECK(freqs_cis.size(2) == head_dim / 2, "freqs_cis dim/2 mismatch");
TORCH_CHECK(freqs_cis.size(3) == 2, "freqs_cis last dim must be 2 [cos, sin]");
TORCH_CHECK(freqs_cis.numel() == (int64_t)n_tokens * head_dim,
"freqs_cis token or rotary dimension mismatch");
TORCH_CHECK(freqs_cis.size(-2) == head_dim / 2, "freqs_cis dim/2 mismatch");
TORCH_CHECK(freqs_cis.size(-1) == 2, "freqs_cis last dim must be 2 [cos, sin]");
auto out = torch::empty_like(x);
int half_dim = head_dim / 2;
int total = batch * seq_len * n_heads * half_dim;
int total = n_tokens * n_heads * half_dim;
int block = 256;
int grid = std::min((total + block - 1) / block, 1024);
@@ -82,7 +81,7 @@ torch::Tensor rotary_emb(
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
freqs_cis.data_ptr<float>(),
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
batch, seq_len, n_heads, head_dim
n_tokens, n_heads, head_dim
);
C10_CUDA_CHECK(cudaGetLastError());
@@ -93,6 +92,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rotary_emb", &rotary_emb,
py::arg("x"),
py::arg("freqs_cis"),
"Fused rotary embedding (bf16 x, f32 freqs_cis [b,s,d/2,2], bf16 out)"
"Fused rotary embedding for packed 3D or dense 4D tensors"
);
}