From 464fe1b77cbfe64ac56552c6ff299cb28a46984a Mon Sep 17 00:00:00 2001
From: Alison Shao <54658187+alisonshao@users.noreply.github.com>
Date: Mon, 13 Jul 2026 17:53:20 -0700
Subject: [PATCH] [CI] Use torch.testing.assert_close in custom-all-reduce test
 (~1400x faster compare) (#31080)

---
 test/registered/jit/test_custom_all_reduce.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/test/registered/jit/test_custom_all_reduce.py b/test/registered/jit/test_custom_all_reduce.py
index 0f3cc84f4..311521d59 100644
--- a/test/registered/jit/test_custom_all_reduce.py
+++ b/test/registered/jit/test_custom_all_reduce.py
@@ -28,7 +28,6 @@ from typing import List
 import pytest
 import torch
 import torch.distributed as dist
-import triton
 
 import sglang.srt.distributed.parallel_state as ps
 from sglang.jit_kernel.all_reduce import (
@@ -224,7 +223,7 @@ def test_custom_all_reduce(
         dist.all_reduce(out_ref, group=nccl_group)
         out_jit = run(inp)
         # Exact equality, since values are small integers within bf16 precision.
-        triton.testing.assert_close(out_ref, out_jit, atol=0, rtol=0)
+        torch.testing.assert_close(out_ref, out_jit, atol=0, rtol=0)
 
 
 if __name__ == "__main__":