53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
import random
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.ops.rand import seeded_uniform
|
|
from vllm.model_executor.utils import set_random_seed
|
|
|
|
|
|
@pytest.mark.parametrize("dtype",
|
|
[torch.float32, torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("use_3d", [True, False])
|
|
def test_seeded_uniform(dtype: torch.dtype, use_3d: bool):
|
|
device = "cuda"
|
|
for seed in range(512):
|
|
set_random_seed(seed)
|
|
rows = random.randint(1, 512)
|
|
cols = random.randint(1, 64000)
|
|
if use_3d:
|
|
third_dim = random.randint(2, 10)
|
|
dims = [rows, third_dim, cols]
|
|
else:
|
|
dims = [rows, cols]
|
|
seeds = torch.randint(torch.iinfo(torch.long).min,
|
|
torch.iinfo(torch.long).max, (rows, ),
|
|
device=device)
|
|
|
|
# Test that the same seed produces the same output
|
|
out = seeded_uniform(*dims, seeds=seeds, dtype=dtype, device=device)
|
|
out2 = seeded_uniform(*dims, seeds=seeds, dtype=dtype, device=device)
|
|
torch.testing.assert_close(out, out2)
|
|
# del to save memory
|
|
del out2
|
|
|
|
out3 = seeded_uniform(*dims, seeds=seeds, dtype=dtype, device=device)
|
|
torch.testing.assert_close(out, out3)
|
|
# del to save memory
|
|
del out3
|
|
|
|
# Initialize out tensor with garbage to ensure that it is overwritten
|
|
out_with_tensor = seeded_uniform(
|
|
*dims,
|
|
out=torch.full(
|
|
(*dims, ),
|
|
-1,
|
|
dtype=dtype,
|
|
device=device,
|
|
),
|
|
seeds=seeds,
|
|
dtype=dtype,
|
|
)
|
|
torch.testing.assert_close(out, out_with_tensor)
|