# kernels.py — Numba JIT kernels for fast CPU inference import numpy as np import torch from numba import jit, prange import math from typing import Optional, Tuple # ───────────────────────────────────────────────────────────────────────────── # ROPE PRECOMPUTATION (runs once at startup) # ───────────────────────────────────────────────────────────────────────────── @jit(nopython=True, parallel=True, cache=True) def precompute_rope_numba(head_dim: int, max_len: int, theta: float) -> Tuple[np.ndarray, np.ndarray]: """Precompute RoPE frequencies — fully vectorized Numba""" cos = np.zeros((max_len, head_dim // 2), dtype=np.float32) sin = np.zeros((max_len, head_dim // 2), dtype=np.float32) for i in prange(head_dim // 2): inv_freq = 1.0 / (theta ** (2.0 * i / head_dim)) for pos in prange(max_len): angle = pos * inv_freq cos[pos, i] = math.cos(angle) sin[pos, i] = math.sin(angle) return cos, sin # ───────────────────────────────────────────────────────────────────────────── # FUSED ROPE APPLICATION (rotate query/key in-place) # ───────────────────────────────────────────────────────────────────────────── @jit(nopython=True, parallel=True, cache=True) def apply_rope_numba(x: np.ndarray, cos: np.ndarray, sin: np.ndarray) -> np.ndarray: """ Apply RoPE rotation (fully in-place, optimized) x: (batch, heads, seq_len, head_dim) cos, sin: (seq_len, head_dim // 2) Rotates: [x0, x1] -> [x0*cos - x1*sin, x1*cos + x0*sin] """ B, H, L, D = x.shape half_d = D // 2 for b in prange(B): for h in prange(H): for l in prange(L): for d in prange(half_d): x0 = x[b, h, l, d] x1 = x[b, h, l, d + half_d] c = cos[l, d] s = sin[l, d] x[b, h, l, d] = x0 * c - x1 * s x[b, h, l, d + half_d] = x1 * c + x0 * s return x # ───────────────────────────────────────────────────────────────────────────── # FUSED SOFTMAX (numerically stable, masked) # ───────────────────────────────────────────────────────────────────────────── @jit(nopython=True, parallel=True, cache=True) def fused_softmax_numba( scores: np.ndarray, scale: float, mask: Optional[np.ndarray] = None ) -> np.ndarray: """ Fused softmax with scale and mask scores: (batch, heads, seq_len, seq_len) mask: (seq_len, seq_len) with -inf for masked positions Returns: attention weights (same shape) """ B, H, L, _ = scores.shape out = np.zeros_like(scores) for b in prange(B): for h in prange(H): for i in prange(L): # Find max for numerical stability max_val = -1e10 for j in range(L): if mask is None or mask[i, j] > -1e8: val = scores[b, h, i, j] * scale if val > max_val: max_val = val # Compute exp and sum sum_exp = 0.0 for j in range(L): if mask is None or mask[i, j] > -1e8: exp_val = math.exp(scores[b, h, i, j] * scale - max_val) out[b, h, i, j] = exp_val sum_exp += exp_val else: out[b, h, i, j] = 0.0 # Normalize if sum_exp > 1e-9: for j in range(L): out[b, h, i, j] /= sum_exp return out # ───────────────────────────────────────────────────────────────────────────── # FUSED ATTENTION (Q @ K^T -> softmax -> @ V, all in one kernel) # ───────────────────────────────────────────────────────────────────────────── @jit(nopython=True, parallel=True, cache=True) def fused_attention_numba( Q: np.ndarray, K: np.ndarray, V: np.ndarray, scale: float, mask: Optional[np.ndarray] = None ) -> np.ndarray: """ Full attention in one fused kernel Q, K, V: (batch, heads, seq_len, head_dim) scale: 1/sqrt(head_dim) mask: (seq_len, seq_len) or None Returns: (batch, heads, seq_len, head_dim) """ B, H, L, D = Q.shape out = np.zeros((B, H, L, D), dtype=np.float32) for b in prange(B): for h in prange(H): for i in prange(L): # Step 1: Compute scores[i, :] = Q[i] @ K[:].T scores = np.zeros(L, dtype=np.float32) max_score = -1e10 for j in range(L): dot = 0.0 for d in range(D): dot += Q[b, h, i, d] * K[b, h, j, d] scaled = dot * scale scores[j] = scaled if mask is None or mask[i, j] > -1e8: if scaled > max_score: max_score = scaled # Step 2: Softmax (numerically stable) sum_exp = 0.0 for j in range(L): if mask is None or mask[i, j] > -1e8: exp_val = math.exp(scores[j] - max_score) scores[j] = exp_val sum_exp += exp_val else: scores[j] = 0.0 if sum_exp > 1e-9: for j in range(L): scores[j] /= sum_exp # Step 3: Apply to values: out[i] = sum_j(scores[j] * V[j]) for d in range(D): val = 0.0 for j in range(L): val += scores[j] * V[b, h, j, d] out[b, h, i, d] = val return out # ───────────────────────────────────────────────────────────────────────────── # TORCH WRAPPERS (handle CPU ↔ Numba conversions) # ───────────────────────────────────────────────────────────────────────────── def apply_rope_fused(x_torch: torch.Tensor, cos_torch: torch.Tensor, sin_torch: torch.Tensor) -> torch.Tensor: """ Torch wrapper for apply_rope_numba Converts to numpy, runs Numba kernel, converts back """ B, H, L, D = x_torch.shape x_np = x_torch.detach().cpu().numpy().astype(np.float32) cos_np = cos_torch.cpu().numpy().astype(np.float32) sin_np = sin_torch.cpu().numpy().astype(np.float32) x_out = apply_rope_numba(x_np, cos_np, sin_np) return torch.from_numpy(x_out).to(x_torch.device).to(x_torch.dtype) def fused_attention( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, scale: float, mask: Optional[torch.Tensor] = None ) -> torch.Tensor: """ Torch wrapper for fused_attention_numba """ Q_np = Q.detach().cpu().numpy().astype(np.float32) K_np = K.detach().cpu().numpy().astype(np.float32) V_np = V.detach().cpu().numpy().astype(np.float32) mask_np = mask.cpu().numpy().astype(np.float32) if mask is not None else None out_np = fused_attention_numba(Q_np, K_np, V_np, scale, mask_np) return torch.from_numpy(out_np).to(Q.device).to(Q.dtype) def fused_softmax( scores: torch.Tensor, scale: float, mask: Optional[torch.Tensor] = None ) -> torch.Tensor: """ Torch wrapper for fused_softmax_numba """ scores_np = scores.detach().cpu().numpy().astype(np.float32) mask_np = mask.cpu().numpy().astype(np.float32) if mask is not None else None out_np = fused_softmax_numba(scores_np, scale, mask_np) return torch.from_numpy(out_np).to(scores.device).to(scores.dtype)