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

Fix test: cos_sin_cache on CUDA device

Browse Source
This commit is contained in:
biondizzle
2026-05-19 02:37:50 +00:00
parent 882d4996ff
commit ae233ab648
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_wo_a.py
View File

@@ -47,11 +47,11 @@ def test_inverse_rope():
rope_dim = ROPE_DIM
half_rope = rope_dim // 2
base = 10000.0
inv_freq = 1.0 / (base ** (torch.arange(0, half_rope, dtype=torch.float32) / half_rope))
inv_freq = 1.0 / (base ** (torch.arange(0, half_rope, dtype=torch.float32, device=DEVICE) / half_rope))
pos = torch.arange(max_pos, dtype=torch.float32)
pos = torch.arange(max_pos, dtype=torch.float32, device=DEVICE)
freqs = torch.outer(pos, inv_freq) # (max_pos, half_rope)
cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) # (max_pos, rope_dim)
cos_sin_cache = torch.cat([freqs.cos(), freqs.sin(), ], dim=-1) # (max_pos, rope_dim)
# Random attention output
o = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0