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
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
"""Rotary embedding with auto-dispatch to CUDA kernel.
|
|
|
|
Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
|
|
CUDA kernel when available, falls back to torch complex multiply otherwise.
|
|
|
|
Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
|
|
freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
|
|
"""
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from astrai.extension.loader import is_available
|
|
|
|
_cache = {"available": None}
|
|
|
|
|
|
def _cuda_available() -> bool:
|
|
if _cache["available"] is None:
|
|
_cache["available"] = is_available("rotary_emb")
|
|
return _cache["available"]
|
|
|
|
|
|
def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
|
cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
|
|
dtype = x.dtype
|
|
x_ = x.float().reshape(*x.shape[:-1], -1, 2)
|
|
x_complex = torch.view_as_complex(x_)
|
|
freqs_cis_complex = torch.complex(cos, sin).unsqueeze(2)
|
|
x_rotated = x_complex * freqs_cis_complex
|
|
x_out = torch.view_as_real(x_rotated).flatten(-2)
|
|
return x_out.to(dtype)
|
|
|
|
|
|
def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
|
|
"""Apply rotary embedding to x.
|
|
|
|
Args:
|
|
x: [batch, seq_len, n_heads, head_dim] (bf16)
|
|
freqs_cis: [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs
|
|
|
|
Returns:
|
|
[batch, seq_len, n_heads, head_dim] (bf16)
|
|
"""
|
|
if (
|
|
_cuda_available()
|
|
and not torch.is_grad_enabled()
|
|
and x.is_cuda
|
|
and x.dtype == torch.bfloat16
|
|
):
|
|
from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
|
|
|
|
return _cuda_rotary(x, freqs_cis)
|
|
return _torch_apply(x, freqs_cis)
|