forked from EngineX-Cambricon/enginex-mlu370-vllm
add qwen3
This commit is contained in:
15
vllm-v0.6.2/tests/kernels/test_permute_cols.py
Normal file
15
vllm-v0.6.2/tests/kernels/test_permute_cols.py
Normal file
@@ -0,0 +1,15 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm._custom_ops import permute_cols
|
||||
|
||||
|
||||
@pytest.mark.parametrize('shape', [(1, 512), (544, 4096), (67, 8192)])
|
||||
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16])
|
||||
def test_permute_cols(shape, dtype):
|
||||
x = torch.randn(shape, dtype=dtype).cuda()
|
||||
perm = torch.randperm(x.shape[1]).to(torch.int).cuda()
|
||||
opcheck(torch.ops._C.permute_cols, (x, perm))
|
||||
y = permute_cols(x, perm)
|
||||
torch.testing.assert_close(y, x[:, perm])
|
||||
Reference in New Issue
Block a user