ViperEkura
3e67b4f88d
- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call) - RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor - apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit) - Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out - Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16) - Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)
77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
from typing import Dict, Optional, Tuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
|
|
|
|
def get_rotary_emb(
|
|
dim: int,
|
|
max_len: int,
|
|
base: float = 10000,
|
|
device: Optional[torch.device] = None,
|
|
) -> Tuple[Tensor, Tensor]:
|
|
"""Precompute cos/sin tables for rotary embedding.
|
|
|
|
Returns:
|
|
(cos, sin) each of shape [max_len, dim/2] (f32)
|
|
"""
|
|
theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
|
|
t = torch.arange(0, max_len, dtype=torch.float64, device=device)
|
|
freqs = torch.outer(t, theta).float()
|
|
return torch.cos(freqs), torch.sin(freqs)
|
|
|
|
|
|
def ntk_base(base: float, dim: int, factor: float) -> float:
|
|
return base * (factor ** (dim / (dim - 2)))
|
|
|
|
|
|
class RotaryEmbedding(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
max_len: int,
|
|
base: float = 10000,
|
|
rope_scaling: Optional[Dict] = None,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.max_len = max_len
|
|
self.base = base
|
|
self.rope_scaling = rope_scaling
|
|
|
|
if rope_scaling is not None:
|
|
scaling_type = rope_scaling.get("type", "ntk")
|
|
factor = rope_scaling.get("factor", 1.0)
|
|
if scaling_type == "ntk":
|
|
self.base = ntk_base(base, dim, factor)
|
|
|
|
self._set_rotary_buffer(self.max_len)
|
|
|
|
def _set_rotary_buffer(self, max_len: int):
|
|
cos, sin = get_rotary_emb(self.dim, max_len, self.base)
|
|
self.register_buffer("cos_table", cos, persistent=False)
|
|
self.register_buffer("sin_table", sin, persistent=False)
|
|
|
|
def forward(
|
|
self, x: Tensor, position_ids: Optional[Tensor] = None
|
|
) -> Tuple[Tensor, Tensor]:
|
|
"""Lookup cos/sin for the given positions.
|
|
|
|
Args:
|
|
x: [batch, seq_len, ...] — only batch and seq_len are used.
|
|
position_ids: [batch, seq_len] optional position indices.
|
|
|
|
Returns:
|
|
(cos, sin) each of shape [batch, seq_len, dim/2] (f32)
|
|
"""
|
|
if position_ids is None:
|
|
position_ids = (
|
|
torch.arange(x.size(1), device=x.device)
|
|
.unsqueeze(0)
|
|
.expand(x.size(0), -1)
|
|
)
|
|
cos = self.cos_table[position_ids].float()
|
|
sin = self.sin_table[position_ids].float()
|
|
return cos, sin
|