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

This commit is contained in:

ViperEkura

2026-09-01 15:30:07 +08:00

parent 1c17e80882

commit aabf366633

3 changed files with 36 additions and 31 deletions

									
										csrc/CMakeLists.txt
									
		+1
		-1
	
												View File
												
				@@ -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)