ViperEkura
7aa5ed09d9
- 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
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
"""Rotary embedding CUDA kernel wrapper.
|
|
|
|
Calls the compiled CUDA kernel directly. If the kernel is not available,
|
|
raises ``RuntimeError``. Fallback to torch complex multiply is the
|
|
responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
|
|
|
|
Layout: x is [batch, seq_len, n_heads, head_dim] (bf16, contiguous).
|
|
freqs_cis is [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from astrai.extension.loader import _available, _modules
|
|
|
|
|
|
def _check_available():
|
|
if not _available.get("rotary_emb"):
|
|
raise RuntimeError(
|
|
"CUDA kernel 'rotary_emb' is not available. "
|
|
"Build with CSRC_KERNELS=true or use the torch fallback."
|
|
)
|
|
|
|
|
|
def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
|
|
"""Fused rotary embedding kernel.
|
|
|
|
Args:
|
|
x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous)
|
|
freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs
|
|
|
|
Returns:
|
|
[batch, seq_len, n_heads, head_dim] (bf16)
|
|
"""
|
|
_check_available()
|
|
if not x.is_contiguous():
|
|
x = x.contiguous()
|
|
if not freqs_cis.is_contiguous():
|
|
freqs_cis = freqs_cis.contiguous()
|
|
return _modules["rotary_emb"].rotary_emb(x, freqs_cis)
|