906e461ed6 · tduncan/Nvidia-Optimization-Stack-VLLM

.buildkite/test-pipeline.yaml

+1

@@ -390,6 +390,7 @@ steps:

390 - csrc/moe/

391 - tests/kernels/moe

392 - vllm/model_executor/layers/fused_moe/

393 + - vllm/distributed/device_communicators/

394 commands:

395 - pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

396 parallelism: 2

tests/distributed/test_comm_ops.py

+5 -7

@@ -18,7 +18,8 @@ from vllm.distributed import (broadcast_tensor_dict, get_pp_group,

18 tensor_model_parallel_all_reduce,

19 tensor_model_parallel_reduce_scatter)

20

21 -from ..utils import init_test_distributed_environment, multi_process_parallel

21 +from ..utils import (init_test_distributed_environment, multi_gpu_test,

22 + multi_process_parallel)

23

24

25 @ray.remote(num_gpus=1, max_calls=1)

@@ -226,8 +227,7 @@ def send_recv_test_worker(

227 torch.testing.assert_close(test_tensor, recv_tensor)

228

229

229 -@pytest.mark.skipif(torch.cuda.device_count() < 2,

230 - reason="Need at least 2 GPUs to run the test.")

230 +@multi_gpu_test(num_gpus=2)

231 @pytest.mark.parametrize("tp_size", [2])

232 @pytest.mark.parametrize("test_target", [

233 all_reduce_test_worker, all_gather_test_worker,

@@ -241,8 +241,7 @@ def test_multi_process_tensor_parallel(

241 multi_process_parallel(monkeypatch, tp_size, 1, test_target)

242

243

244 -@pytest.mark.skipif(torch.cuda.device_count() < 2,

245 - reason="Need at least 2 GPUs to run the test.")

244 +@multi_gpu_test(num_gpus=2)

245 @pytest.mark.parametrize("pp_size", [2])

246 @pytest.mark.parametrize(

247 "test_target", [send_recv_test_worker, send_recv_tensor_dict_test_worker])

@@ -254,8 +253,7 @@ def test_multi_process_pipeline_parallel(

253 multi_process_parallel(monkeypatch, 1, pp_size, test_target)

254

255

257 -@pytest.mark.skipif(torch.cuda.device_count() < 4,

258 - reason="Need at least 4 GPUs to run the test.")

256 +@multi_gpu_test(num_gpus=4)

257 @pytest.mark.parametrize("tp_size", [2])

258 @pytest.mark.parametrize("pp_size", [2])

259 @pytest.mark.parametrize("test_target", [

tests/kernels/moe/test_deepep_deepgemm_moe.py

+3

@@ -23,6 +23,7 @@ from vllm.utils import has_deep_ep, has_deep_gemm

23 from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_e8m0_used,

24 is_deep_gemm_supported)

25

26 +from ...utils import multi_gpu_test

27 from .parallel_utils import ProcessGroupInfo, parallel_launch

28 from .utils import make_test_weights

29

@@ -370,6 +371,7 @@ NUM_EXPERTS = [32]

371 @pytest.mark.parametrize("num_experts", NUM_EXPERTS)

372 @pytest.mark.parametrize("topk", TOPKS)

373 @pytest.mark.parametrize("world_dp_size", [(2, 1)])

374 +@multi_gpu_test(num_gpus=2)

375 @requires_deep_ep

376 @requires_deep_gemm

377 @pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

@@ -427,6 +429,7 @@ USE_FP8_DISPATCH = [False]

429 @pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH)

430 @pytest.mark.parametrize("block_size", [[128, 128]])

431 @pytest.mark.parametrize("world_dp_size", [(2, 1)])

432 +@multi_gpu_test(num_gpus=2)

433 @requires_deep_ep

434 @requires_deep_gemm

435 @pytest.mark.skipif(is_blackwell_deep_gemm_e8m0_used(),

tests/kernels/moe/test_deepep_moe.py

+3

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

24 from vllm.platforms import current_platform

25 from vllm.utils import has_deep_ep

26

27 +from ...utils import multi_gpu_test

28 from .parallel_utils import ProcessGroupInfo, parallel_launch

29

30 if has_deep_ep():

@@ -411,6 +412,7 @@ DTYPES = [torch.bfloat16, torch.float8_e4m3fn]

412 @pytest.mark.parametrize("topk", [6])

413 @pytest.mark.parametrize("world_dp_size", [(2, 1)])

414 @pytest.mark.parametrize("per_act_token_quant", [False, True])

415 +@multi_gpu_test(num_gpus=2)

416 @requires_deep_ep

417 def test_deep_ep_moe(

418 dtype: torch.dtype,

@@ -459,6 +461,7 @@ USE_FP8_DISPATCH = [True, False]

461 @pytest.mark.parametrize("topk", [6])

462 @pytest.mark.parametrize("world_dp_size", [(2, 1)])

463 @pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH)

464 +@multi_gpu_test(num_gpus=2)

465 @requires_deep_ep

466 def test_low_latency_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],

467 num_experts: int, topk: int,

tests/kernels/moe/test_modular_kernel_combinations.py

+2

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

16 from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx

17 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

18

19 +from ...utils import multi_gpu_test

20 from .modular_kernel_tools.common import (Config, RankTensors, WeightTensors,

21 reference_moe_impl,

22 run_modular_kernel)

@@ -162,6 +163,7 @@ def is_nyi_config(config: Config) -> bool:

163 product(MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))

164 @pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)

165 @pytest.mark.parametrize("world_size", [2])

166 +@multi_gpu_test(num_gpus=2)

167 @meets_multi_gpu_requirements

168 def test_modular_kernel_combinations_multigpu(

169 k: int, n: int, e: int, dtype: torch.dtype,

tests/kernels/moe/test_pplx_cutlass_moe.py

+2

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

17 from vllm.platforms import current_platform

18 from vllm.utils import cdiv

19

20 +from ...utils import multi_gpu_test

21 from .parallel_utils import ProcessGroupInfo, parallel_launch

22

23 try:

@@ -247,6 +248,7 @@ def _pplx_moe(

248 @pytest.mark.parametrize("per_out_ch", [True, False])

249 @pytest.mark.parametrize("world_dp_size", [[2, 1]]) #, [4, 2]])

250 @pytest.mark.parametrize("use_internode", [False])

251 +@multi_gpu_test(num_gpus=2)

252 @pytest.mark.skipif(

253 (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(

254 current_platform.get_device_capability()),

tests/kernels/moe/test_pplx_moe.py

+5

@@ -37,6 +37,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

37 from vllm.platforms import current_platform

38 from vllm.utils import round_up

39

40 +from ...utils import multi_gpu_test

41 from .parallel_utils import ProcessGroupInfo, parallel_launch

42

43 requires_pplx = pytest.mark.skipif(

@@ -452,6 +453,7 @@ def _pplx_prepare_finalize(

453 @pytest.mark.parametrize("use_internode", [False])

454 @pytest.mark.optional

455 @requires_pplx

456 +@multi_gpu_test(num_gpus=2)

457 def test_pplx_prepare_finalize_slow(

458 mnk: tuple[int, int, int],

459 e: int,

@@ -740,6 +742,7 @@ def _pplx_moe(

742 @pytest.mark.parametrize("use_internode", [False])

743 @pytest.mark.optional

744 @requires_pplx

745 +@multi_gpu_test(num_gpus=2)

746 def test_pplx_moe_slow(

747 mnk: tuple[int, int, int],

748 e: int,

@@ -880,6 +883,7 @@ def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,

883 @pytest.mark.parametrize("world_dp_size", [[2, 1]])

884 @pytest.mark.parametrize("use_internode", [False])

885 @requires_pplx

886 +@multi_gpu_test(num_gpus=2)

887 def test_pplx_prepare_finalize(

888 world_dp_size: tuple[int, int],

889 use_internode: bool,

@@ -893,6 +897,7 @@ def test_pplx_prepare_finalize(

897 @pytest.mark.parametrize("world_dp_size", [[2, 1]])

898 @pytest.mark.parametrize("use_internode", [False])

899 @requires_pplx

900 +@multi_gpu_test(num_gpus=2)

901 def test_pplx_moe(

902 world_dp_size: tuple[int, int],

903 use_internode: bool,

tests/utils.py

+6 -3

@@ -696,9 +696,12 @@ def multi_process_parallel(

696 os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1"

697 ray.init(

698 runtime_env={

699 - "working_dir": VLLM_PATH,

700 - "excludes":

701 - ["build", ".git", "cmake-build-*", "shellcheck", "dist"]

699 + "working_dir":

700 + VLLM_PATH,

701 + "excludes": [

702 + "build", ".git", "cmake-build-*", "shellcheck", "dist",

703 + "ep_kernels_workspace"

704 + ]

705 })

706

707 distributed_init_port = get_open_port()

tools/ep_kernels/install_python_libraries.sh

+13 -2

@@ -77,6 +77,7 @@ clone_repo() {

77 local repo_url=$1

78 local dir_name=$2

79 local key_file=$3

80 + local commit_hash=$4

81

82 if [ -d "$dir_name" ]; then

83 # Check if directory has uncommitted changes (dirty)

@@ -87,17 +88,27 @@ clone_repo() {

88 echo "$dir_name directory exists but clone appears incomplete, cleaning up and re-cloning"

89 rm -rf "$dir_name"

90 git clone "$repo_url"

91 + if [ -n "$commit_hash" ]; then

92 + cd "$dir_name"

93 + git checkout "$commit_hash"

94 + cd ..

95 + fi

96 else

97 echo "$dir_name directory exists and appears complete; manually update if needed"

98 fi

99 else

100 git clone "$repo_url"

101 + if [ -n "$commit_hash" ]; then

102 + cd "$dir_name"

103 + git checkout "$commit_hash"

104 + cd ..

105 + fi

106 fi

107 }

108

109 # build and install pplx, require pytorch installed

110 pushd $WORKSPACE

100 -clone_repo "https://github.com/ppl-ai/pplx-kernels" "pplx-kernels" "setup.py"

111 +clone_repo "https://github.com/ppl-ai/pplx-kernels" "pplx-kernels" "setup.py" "c336faf"

112 cd pplx-kernels

113 # see https://github.com/pypa/pip/issues/9955#issuecomment-838065925

114 # PIP_NO_BUILD_ISOLATION=0 disables build isolation

@@ -106,7 +117,7 @@ popd

117

118 # build and install deepep, require pytorch installed

119 pushd $WORKSPACE

109 -clone_repo "https://github.com/deepseek-ai/DeepEP" "DeepEP" "setup.py"

120 +clone_repo "https://github.com/deepseek-ai/DeepEP" "DeepEP" "setup.py" "e3908bf"

121 cd DeepEP

122 export NVSHMEM_DIR=$WORKSPACE/nvshmem_install

123 PIP_NO_BUILD_ISOLATION=0 pip install -vvv -e .

[CI Fix] Pin deepep and pplx tags in tools/ep_kernels/, gate multigpu tests (#23568)