EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
tests/ut/ops/test_comm_utils.py
View File

@@ -20,13 +20,14 @@ import torch
from pytest_mock import MockerFixture
from tests.ut.base import PytestBase
from vllm_ascend.ops.moe.comm_utils import (
_gather_along_first_dim, async_all_to_all,
gather_from_sequence_parallel_region)
from vllm_ascend.ops.fused_moe.comm_utils import (
_gather_along_first_dim,
async_all_to_all,
gather_from_sequence_parallel_region,
)
class TestDistributedCommunication(PytestBase):
@pytest.fixture(autouse=True)
def context(self, mocker: MockerFixture):
mocker.patch("torch.npu.current_device", return_value="cpu")
@@ -36,63 +37,52 @@ class TestDistributedCommunication(PytestBase):
@pytest.mark.parametrize(
"input_tensor, output_split_sizes, input_split_sizes",
[(torch.randn(8, 16), [2, 2, 2, 2], [2, 2, 2, 2]),
(torch.randn(16, 32), None, None)])
def test_async_all_to_all(self, input_tensor, output_split_sizes,
input_split_sizes, mocker: MockerFixture):
[(torch.randn(8, 16), [2, 2, 2, 2], [2, 2, 2, 2]), (torch.randn(16, 32), None, None)],
)
def test_async_all_to_all(self, input_tensor, output_split_sizes, input_split_sizes, mocker: MockerFixture):
"""Test async_all_to_all"""
mock_group = mocker.MagicMock()
mocker.patch("torch.distributed.all_to_all_single",
return_value=mocker.MagicMock())
mocker.patch("torch.distributed.all_to_all_single", return_value=mocker.MagicMock())
_, a2a_out, handle = async_all_to_all(input_tensor, output_split_sizes,
input_split_sizes, mock_group)
_, a2a_out, handle = async_all_to_all(input_tensor, output_split_sizes, input_split_sizes, mock_group)
# Check if the output tensor is created properly
if output_split_sizes is None:
assert a2a_out.shape == input_tensor.shape
else:
total_output_size = sum(output_split_sizes)
expected_shape = [total_output_size] + list(
input_tensor.size())[1:]
expected_shape = [total_output_size] + list(input_tensor.size())[1:]
assert a2a_out.shape == torch.Size(expected_shape)
# Ensure handle is returned from async operation
assert handle is not None
assert isinstance(handle, mocker.MagicMock)
@pytest.mark.parametrize("world_size, test_tensor, expected",
[(1, torch.randn(8, 16), (8, 16)),
(4, torch.randn(8, 16), (32, 16))])
def test_gather_along_first_dim(self, test_tensor, expected, world_size,
mocker: MockerFixture):
@pytest.mark.parametrize(
"world_size, test_tensor, expected", [(1, torch.randn(8, 16), (8, 16)), (4, torch.randn(8, 16), (32, 16))]
)
def test_gather_along_first_dim(self, test_tensor, expected, world_size, mocker: MockerFixture):
"""Test _gather_along_first_dim"""
mocker.patch("torch.distributed.get_world_size",
return_value=world_size)
mocker.patch("torch.distributed.get_world_size", return_value=world_size)
result = _gather_along_first_dim(test_tensor, mocker.MagicMock())
assert result.shape == expected
@pytest.mark.parametrize("input_tensor, output_split_sizes",
[(torch.randn(8, 16), None),
(torch.randn(8, 16), [2, 2, 2, 2])])
def test_gather_from_sequence_parallel_region(self, input_tensor,
output_split_sizes,
mocker: MockerFixture):
@pytest.mark.parametrize(
"input_tensor, output_split_sizes", [(torch.randn(8, 16), None), (torch.randn(8, 16), [2, 2, 2, 2])]
)
def test_gather_from_sequence_parallel_region(self, input_tensor, output_split_sizes, mocker: MockerFixture):
"""Test gather_from_sequence_parallel_region"""
mock_group = mocker.MagicMock()
result = gather_from_sequence_parallel_region(input_tensor, mock_group,
output_split_sizes)
result = gather_from_sequence_parallel_region(input_tensor, mock_group, output_split_sizes)
# If output_split_sizes is not provided, result should have expanded first dimension by world size
if output_split_sizes is None:
expected_shape = [input_tensor.shape[0] * 4] + list(
input_tensor.shape[1:])
expected_shape = [input_tensor.shape[0] * 4] + list(input_tensor.shape[1:])
assert result.shape == torch.Size(expected_shape)
else:
# If output_split_sizes is provided, result shape is dictated by sum of output_split_sizes
expected_shape = [sum(output_split_sizes)] + list(
input_tensor.shape[1:])
expected_shape = [sum(output_split_sizes)] + list(input_tensor.shape[1:])
assert result.shape == torch.Size(expected_shape)