test_kernel.py

"""
Verification test for Delta-RoPE kernel implementation.
Ensures zero em-dashes and exact mathematical equivalence.
"""

import sys
import torch
from kernel import apply_delta_rope_inplace

def verify_delta_rope():
    torch.manual_seed(42)
    batch_size = 1
    num_heads = 8
    seq_len = 512
    head_dim = 128
    base = 10000.0
    
    # 1. Generate unrotated Key vectors (output of W_k)
    k_raw = torch.randn(batch_size, num_heads, seq_len, head_dim, dtype=torch.float64)
    
    # 2. Compute RoPE directly at position m = 120
    m_base = 120
    indices = torch.arange(0, head_dim, 2, dtype=torch.float64)
    inv_freq = 1.0 / (base ** (indices / head_dim))
    
    pos_m = torch.arange(m_base, m_base + seq_len, dtype=torch.float64).unsqueeze(-1)
    angles_m = pos_m * inv_freq
    cos_m = torch.cos(angles_m)
    sin_m = torch.sin(angles_m)
    
    k_pairs = k_raw.view(batch_size, num_heads, seq_len, head_dim // 2, 2)
    k0_m = k_pairs[..., 0] * cos_m - k_pairs[..., 1] * sin_m
    k1_m = k_pairs[..., 0] * sin_m + k_pairs[..., 1] * cos_m
    k_at_m = torch.stack([k0_m, k1_m], dim=-1).view_as(k_raw)
    
    # 3. Compute RoPE directly at target position m' = m + Delta (Delta = +200)
    delta = 200
    pos_target = pos_m + delta
    angles_target = pos_target * inv_freq
    cos_t = torch.cos(angles_target)
    sin_t = torch.sin(angles_target)
    
    k0_t = k_pairs[..., 0] * cos_t - k_pairs[..., 1] * sin_t
    k1_t = k_pairs[..., 0] * sin_t + k_pairs[..., 1] * cos_t
    k_at_target_direct = torch.stack([k0_t, k1_t], dim=-1).view_as(k_raw)
    
    # 4. Compute via Delta-RoPE on k_at_m
    k_at_target_rotated = apply_delta_rope_inplace(k_at_m.clone(), delta=delta, base=base)
    
    # 5. Measure numerical error
    diff = torch.abs(k_at_target_direct - k_at_target_rotated)
    max_err = diff.max().item()
    mean_err = diff.mean().item()
    
    print(f"Delta-RoPE Kernel Verification (Float64):")
    print(f"  Max Absolute Error:  {max_err:.4e}")
    print(f"  Mean Absolute Error: {mean_err:.4e}")
    
    assert max_err < 1e-12, f"Max error {max_err} exceeds 1e-12 tolerance!"
    print("STATUS: PASSED (Bit-for-bit mathematical identity confirmed)")

if __name__ == "__main__":
    verify_delta_rope()