refactor test tolerance

Author: airMeng

python/sgl_kernel/moe.py

@@ -355,7 +355,7 @@ def fused_experts(
     torch.ops.sgl_kernel.silu_and_mul(intermediate_cache2, intermediate_cache1)
 
     torch.ops.sgl_kernel.moe_grouped_mm_nt(
-        intermediate_cache3, intermediate_cache2.contiguous(), w2, offset, E
+        intermediate_cache3, intermediate_cache2, w2, offset, E
     )
 
     flat_weights = topk_weights.to(intermediate_cache3.dtype).flatten()[idxs]  # [N]

tests/test_moe_gemm.py

@@ -23,7 +23,7 @@ def create_random_xpu_tensor(shape, dtype, mean=0, std=0.01):
     Returns:
         torch.Tensor: Randomly initialized xpu tensor
     """
-    return torch.randn(shape, device="xpu").to(dtype)
+    return torch.empty(shape, dtype=dtype, device="xpu").normal_(mean, std)
 
 
 def torch_naive_moe(
@@ -65,7 +65,7 @@ def torch_naive_moe(
     ),
 )
 def test_moe_gemm(num_tokens, topk, num_experts, hidden_size, intermediate_size):
-    rtol, atol = 2e-2, 2e-1
+    rtol, atol = 1e-1, 1e-2
     a = create_random_xpu_tensor((num_tokens, hidden_size), torch.bfloat16)
     w1 = create_random_xpu_tensor(
         (num_experts, 2 * intermediate_size, hidden_size), torch.bfloat16
@@ -93,9 +93,7 @@ def test_moe_gemm(num_tokens, topk, num_experts, hidden_size, intermediate_size)
         topk_ids,
     )
     # import pdb; pdb.set_trace()
-    assert torch.allclose(
-        torch_output, sglang_output, rtol=rtol, atol=atol * hidden_size
-    )
+    torch.testing.assert_close(torch_output, sglang_output, rtol=rtol, atol=atol)
 
 
 if __name__ == "__main__":