perf: vectorize rotary kernel loads and halve index math

- walk exact 2-pair chunks (8B x access, 16B cos/sin float4) and decompose the flat index per chunk instead of per pair, halving integer div/mod work
- enforce head_dim % 4 == 0 at the binding instead of carrying a scalar fallback path
- raise the grid-stride block cap from 1024 to 2048 for full SM coverage on streaming shapes
- hoist kernels/rotary/rotary_emb.cu to kernels/rotary_emb.cu (single-file directory)

Benchmark: NVIDIA L20 (sm_89, shared GPU), interleaved A/B of old and new module, 500-iter means
- (32768 tokens, 8 heads, D=64): 72.7 -> 37.5 us (1.94x)
- (32768 tokens, 32 heads, D=256): 4420 -> 3630 us (1.22x)
- (32768 tokens, 32 heads, D=128): 2120 -> 1824 us (1.16x)
- (32 tokens, 32 heads, D=128) decode size: unchanged at ~1.9 us

csrc/CMakeLists.txt

@@ -68,7 +68,7 @@ set(KERNEL_SRCS
     attention/prefill.cu
     attention/paged_decode.cu
     attention/paged_prefill.cu
-    rotary/rotary_emb.cu
+    rotary_emb.cu
 )
 
 if(ASTRAI_CUDA_ARCH GREATER_EQUAL 89)

csrc/kernels/rotary_emb.cu

@@ -11,34 +11,41 @@ __global__ void rotary_emb_kernel(
     int n_heads,
     int head_dim
 ) {
-    const int half_dim = head_dim >> 1;
-    const int total = n_tokens * n_heads * half_dim;
+    // Each head tiles into exact 2-pair chunks: one 8B x access and one 16B
+    // cos/sin access per chunk (head_dim % 4 == 0 is enforced on the host).
+    const int chunks = head_dim >> 2;
+    const int total = n_tokens * n_heads * chunks;
 
-    for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
-         idx < total;
-         idx += gridDim.x * blockDim.x) {
+    for (int c = blockIdx.x * blockDim.x + threadIdx.x;
+         c < total;
+         c += gridDim.x * blockDim.x) {
+        const int chunk = c % chunks;
+        const int tmp = c / chunks;
+        const int head = tmp % n_heads;
+        const int token = tmp / n_heads;
 
-        int pair = idx % half_dim;
-        int tmp = idx / half_dim;
-        int head = tmp % n_heads;
-        tmp /= n_heads;
-        int token = tmp;
+        const int x_off = (tmp * head_dim) + (chunk << 2);
+        const int f_off = ((token * chunks) + chunk) << 2;
 
-        int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
-        int cs_offset = (token * half_dim + pair) * 2;
+        const float4 f = *reinterpret_cast<const float4*>(freqs_cis + f_off);
+        const uint2 xr = *reinterpret_cast<const uint2*>(x + x_off);
+        __nv_bfloat162 p0 = *reinterpret_cast<const __nv_bfloat162*>(&xr.x);
+        __nv_bfloat162 p1 = *reinterpret_cast<const __nv_bfloat162*>(&xr.y);
 
-        __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));
+        const float e0 = __bfloat162float(__low2bfloat16(p0));
+        const float o0 = __bfloat162float(__high2bfloat16(p0));
+        const float e1 = __bfloat162float(__low2bfloat16(p1));
+        const float o1 = __bfloat162float(__high2bfloat16(p1));
 
-        float c = freqs_cis[cs_offset];
-        float s = freqs_cis[cs_offset + 1];
+        __nv_bfloat162 r0 = __floats2bfloat162_rn(
+            e0 * f.x - o0 * f.y, e0 * f.y + o0 * f.x);
+        __nv_bfloat162 r1 = __floats2bfloat162_rn(
+            e1 * f.z - o1 * f.w, e1 * f.w + o1 * f.z);
 
-        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;
+        uint2 oraw;
+        *reinterpret_cast<__nv_bfloat162*>(&oraw.x) = r0;
+        *reinterpret_cast<__nv_bfloat162*>(&oraw.y) = r1;
+        *reinterpret_cast<uint2*>(out + x_off) = oraw;
     }
 }
 
@@ -64,7 +71,7 @@ torch::Tensor rotary_emb(
     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(head_dim % 4 == 0, "head_dim must be a multiple of 4");
     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");
@@ -72,10 +79,9 @@ torch::Tensor rotary_emb(
 
     auto out = torch::empty_like(x);
 
-    int half_dim = head_dim / 2;
-    int total = n_tokens * n_heads * half_dim;
+    int work = n_tokens * n_heads * (head_dim / 4);
     int block = 256;
-    int grid = std::min((total + block - 1) / block, 1024);
+    int grid = std::min((work + block - 1) / block, 2048);
 
     rotary_emb_kernel<<<grid, block, 0, stream>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),