biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix RoPE test: use proper cos/sin cache

Browse Source
This commit is contained in:
biondizzle
2026-06-02 10:04:01 +00:00
parent f566b9b748
commit 12b6365b42
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/unit/test_kv_quantize.py
View File

@@ -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