minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CI] Use torch.testing.assert_close in custom-all-reduce test (~1400x faster compare) (#31080)

Browse Source
This commit is contained in:
Alison Shao
2026-07-14 08:53:20 +08:00
committed by GitHub
parent 33f83011e0
commit 464fe1b77c
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/test_custom_all_reduce.py
+1 -2
View File
@@ -28,7 +28,6 @@ from typing import List
import pytest import pytest
import torch import torch
import torch.distributed as dist import torch.distributed as dist
import triton
import sglang.srt.distributed.parallel_state as ps import sglang.srt.distributed.parallel_state as ps
from sglang.jit_kernel.all_reduce import ( from sglang.jit_kernel.all_reduce import (
@@ -224,7 +223,7 @@ def test_custom_all_reduce(
dist.all_reduce(out_ref, group=nccl_group) dist.all_reduce(out_ref, group=nccl_group)
out_jit = run(inp) out_jit = run(inp)
# Exact equality, since values are small integers within bf16 precision. # 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__": if __name__ == "__main__":