kernels.py

# 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)