Add FP32 dtype support for RoPE - Part1 (#13181)

This commit is contained in:
iLeGend
2025-11-15 11:37:18 -08:00
committed by GitHub
parent 0d116b9a0b
commit 20e59f9510
2 changed files with 13 additions and 2 deletions
@@ -47,6 +47,12 @@ from sgl_kernel.testing.rotary_embedding import (
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 32, 8, False),
(128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 16, 4, False),
(512, 128, 311, 10000, False, torch.bfloat16, "cuda", 3, 39, 4, 2, False),
(64, 64, 32, 8000, True, torch.float32, "cuda", 32, 32, 1, 1, False),
(256, 128, 4096, 10000, True, torch.float32, "cuda", 2, 512, 4, 2, False),
(512, 128, 311, 10000, True, torch.float32, "cuda", 3, 39, 4, 2, False),
(128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 32, 8, False),
(128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 16, 4, False),
(512, 128, 311, 10000, False, torch.float32, "cuda", 3, 39, 4, 2, False),
],
)
def test_correctness(