[kernel] Fix position ids in rope (#3173)

2025-01-27 01:09:51 -08:00
parent 52c03f16b9
commit 514f37c32b
4 changed files with 7 additions and 3 deletions
--- a/sgl-kernel/tests/test_rotary_embedding.py
+++ b/sgl-kernel/tests/test_rotary_embedding.py
@@ -196,3 +196,7 @@ def test_correctness(
        query_ref_out, query_flashinfer_out, atol=1e-2, rtol=1e-2
    )
    torch.testing.assert_close(key_ref_out, key_flashinfer_out, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])