95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.srt.debug_utils.comparator.utils import (
|
|
argmax_coord,
|
|
calc_rel_diff,
|
|
compute_smaller_dtype,
|
|
try_unify_shape,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
|
|
|
|
|
class TestCalcRelDiff:
|
|
def test_identical_tensors(self):
|
|
x = torch.randn(10, 10)
|
|
assert calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
|
|
|
|
def test_orthogonal_tensors(self):
|
|
result = calc_rel_diff(
|
|
torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])
|
|
).item()
|
|
assert result == pytest.approx(1.0, abs=1e-5)
|
|
|
|
def test_similar_tensors(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0])
|
|
y = torch.tensor([1.01, 2.01, 3.01])
|
|
result = calc_rel_diff(x, y).item()
|
|
assert 0.0 < result < 0.01
|
|
|
|
def test_negated_tensors(self):
|
|
x = torch.tensor([1.0, 2.0])
|
|
result = calc_rel_diff(x, -x).item()
|
|
assert result == pytest.approx(2.0, abs=1e-5)
|
|
|
|
|
|
class TestArgmaxCoord:
|
|
def test_1d_tensor(self):
|
|
x = torch.tensor([0.0, 0.0, 5.0, 0.0])
|
|
assert argmax_coord(x) == (2,)
|
|
|
|
def test_2d_tensor(self):
|
|
x = torch.zeros(3, 4)
|
|
x[1, 2] = 10.0
|
|
assert argmax_coord(x) == (1, 2)
|
|
|
|
def test_3d_tensor(self):
|
|
x = torch.zeros(2, 3, 4)
|
|
x[1, 2, 3] = 10.0
|
|
assert argmax_coord(x) == (1, 2, 3)
|
|
|
|
|
|
class TestTryUnifyShape:
|
|
def test_squeeze_leading_ones(self):
|
|
target = torch.Size([3, 4])
|
|
assert try_unify_shape(torch.randn(1, 1, 3, 4), target).shape == target
|
|
|
|
def test_no_squeeze_when_leading_dim_not_one(self):
|
|
target = torch.Size([3, 4])
|
|
assert try_unify_shape(torch.randn(2, 3, 4), target).shape == (2, 3, 4)
|
|
|
|
def test_same_shape_noop(self):
|
|
target = torch.Size([3, 4])
|
|
x = torch.randn(3, 4)
|
|
result = try_unify_shape(x, target)
|
|
assert result.shape == target
|
|
assert result.data_ptr() == x.data_ptr()
|
|
|
|
def test_trailing_dims_mismatch(self):
|
|
target = torch.Size([5, 6])
|
|
x = torch.randn(1, 3, 4)
|
|
result = try_unify_shape(x, target)
|
|
assert result.shape == (1, 3, 4)
|
|
|
|
|
|
class TestComputeSmallerDtype:
|
|
def test_float32_bfloat16(self):
|
|
assert compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16
|
|
|
|
def test_reverse_order(self):
|
|
assert compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16
|
|
|
|
def test_same_dtype_returns_none(self):
|
|
assert compute_smaller_dtype(torch.float32, torch.float32) is None
|
|
|
|
def test_unknown_pair_returns_none(self):
|
|
assert compute_smaller_dtype(torch.int32, torch.int64) is None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__]))
|