From aabf36663310be4da2668516bcbd2a29b0600b9e Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Tue, 1 Sep 2026 15:06:08 +0800 Subject: [PATCH] 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 | 2 +- csrc/kernels/{rotary => }/rotary_emb.cu | 58 ++++++++++++++----------- docs/developer/cuda_kernels.md | 7 ++- 3 files changed, 36 insertions(+), 31 deletions(-) rename csrc/kernels/{rotary => }/rotary_emb.cu (55%) diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index eecf091..b08b91f 100644 --- a/csrc/CMakeLists.txt +++ b/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) diff --git a/csrc/kernels/rotary/rotary_emb.cu b/csrc/kernels/rotary_emb.cu similarity index 55% rename from csrc/kernels/rotary/rotary_emb.cu rename to csrc/kernels/rotary_emb.cu index db69e76..b1453af 100644 --- a/csrc/kernels/rotary/rotary_emb.cu +++ b/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(freqs_cis + f_off); + const uint2 xr = *reinterpret_cast(x + x_off); + __nv_bfloat162 p0 = *reinterpret_cast(&xr.x); + __nv_bfloat162 p1 = *reinterpret_cast(&xr.y); - __nv_bfloat162 x_pair = *reinterpret_cast(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(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<<>>( reinterpret_cast(x.data_ptr()), diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index eca745f..b28259b 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -10,7 +10,7 @@ AstrAI includes optional custom CUDA kernels for attention, rotary embedding, an | `attn_prefill` | `attention/prefill.cu` | GQA prefill attention (split-Q) | | `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention | | `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) | -| `rotary_emb` | `rotary/rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) | +| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) | | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) | Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist: @@ -27,7 +27,7 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac ### Rotary Embedding Kernel -The `rotary_emb` kernel (`csrc/kernels/rotary/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel: +The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel: - One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store - f32 cos/sin input, bf16 compute and output @@ -460,8 +460,7 @@ csrc/ │ │ ├── prefill.cu # → module attn_prefill │ │ ├── paged_decode.cu # → module attn_paged_decode │ │ └── paged_prefill.cu # → module attn_paged_prefill -│ ├── rotary/ -│ │ └── rotary_emb.cu # rotary embedding (kernel + binding in one file) → module rotary_emb +│ ├── rotary_emb.cu # rotary embedding (kernel + binding in one file) → module rotary_emb │ └── fp8/ # FP8 family (module name fp8_ops) │ ├── common.h # FP8Format enum, Fp8GemmTraits, FP8Params / FP8QuantizeParams PODs, layout tags (no torch) │ ├── quantize.cuh # quantize kernels: vectorized + 32×32-tile transpose (out_layout 0/1/2) (no torch)