add gatherep select. (#2740)

### What this PR does / why we need it?
add gatherep select.

- vLLM version: v0.10.1.1
- vLLM main:
e599e2c65e

Signed-off-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>
Co-authored-by: wangxiaoxin-sherie <wangxiaoxin7@huawei.com>
This commit is contained in:
sherie
2025-09-08 09:15:50 +08:00
committed by GitHub
parent 6666e5265d
commit 2693196ef8
3 changed files with 31 additions and 6 deletions

View File

@@ -0,0 +1,22 @@
import os
import unittest
from unittest import mock
from vllm_ascend.ascend_forward_context import get_dispatcher_name
class TestGetDispatcherName(unittest.TestCase):
def test_get_dispatcher_name(self):
result = get_dispatcher_name(1, False)
assert result == "TokenDispatcherWithAllGather"
result = get_dispatcher_name(4, False)
assert result == "TokenDispatcherWithAll2AllV"
result = get_dispatcher_name(16, True)
assert result == "TokenDispatcherWithAll2AllV"
result = get_dispatcher_name(16, False)
assert result == "TokenDispatcherWithMC2"
with mock.patch.dict(os.environ,
{"VLLM_ENABLE_FUSED_EXPERTS_ALLGATHER_EP": "1"}):
result = get_dispatcher_name(16, False)
assert result == "TokenDispatcherWithAllGather"