[Misc] Clean sgl-kernel test (#5216)

2025-04-11 02:28:41 +08:00
parent d050df368c
commit ed01b4515e
8 changed files with 63 additions and 76 deletions
--- a/sgl-kernel/tests/test_per_token_group_quant_8bit.py
+++ b/sgl-kernel/tests/test_per_token_group_quant_8bit.py
@@ -304,10 +304,10 @@ def test_per_token_group_quant_with_column_major(
        scale_tma_aligned=scale_tma_aligned,
    )

-    assert torch.allclose(
+    torch.testing.assert_close(
        x_q_triton.to(torch.float32), x_q_sglang.to(torch.float32), rtol=1e-3, atol=1e-5
    )
-    assert torch.allclose(
+    torch.testing.assert_close(
        x_s_triton.contiguous(), x_s_sglang.contiguous(), rtol=1e-3, atol=1e-5
    )