diff --git a/tests/unit/test_kv_quantize.py b/tests/unit/test_kv_quantize.py index b198d396..443b0304 100644 --- a/tests/unit/test_kv_quantize.py +++ b/tests/unit/test_kv_quantize.py @@ -64,13 +64,13 @@ for M in [1, 4, 32, 128]: # =========================================================================== print("\n--- Test 3: FP32 RoPE (production rope_dim=64, hd=512) ---") hd = 512; rope_dim = 64 -cos_cache = torch.randn(1024, rope_dim // 2, device=device, dtype=torch.float32) -sin_cache = torch.randn(1024, rope_dim // 2, device=device, dtype=torch.float32) -# Normalize for proper RoPE -for i in range(1024): - norm = (cos_cache[i] ** 2 + sin_cache[i] ** 2).sqrt().clamp(min=1e-8) - cos_cache[i] /= norm - sin_cache[i] /= norm +# Build proper RoPE cache (same as single_shot build_rope_cache) +theta = 10000.0 +freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim)) +angles = torch.outer(torch.arange(1024, dtype=torch.float32), freqs) +cos_cache = torch.cos(angles).to(device) # (1024, rope_dim/2) FP32 +sin_cache = torch.sin(angles).to(device) # (1024, rope_dim/2) FP32 +# cos²+sin²=1 by construction for M in [1, 4, 8]: data = torch.randn(M, hd, device=device, dtype=torch.float32) * 2.0