From cffc0cfa311f4997be58cfe8f5f6473ab4898269 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 16 Oct 2026 12:00:12 -0400 Subject: [PATCH] [CI] Fix B200 FusedMoE Layer Kernels hang on NVSHMEM IB setup (#61695) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- .buildkite/test_areas/kernels.yaml | 2 +- .../moe/modular_kernel_tools/parallel_utils.py | 18 ++++++++++++++++-- tests/kernels/moe/test_moe_layer.py | 12 +++++++++++- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 12b8516926..a1612bc3e8 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -939,7 +939,7 @@ steps: - label: ":nvidia: (B200) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-b200s - timeout_in_minutes: 90 + timeout_in_minutes: 45 device: b200-k8s num_devices: 2 source_file_dependencies: diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index d7394d72cb..235d1b81a4 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import os +import time import traceback from collections.abc import Callable from typing import Any, Concatenate @@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port P = ParamSpec("P") +LAUNCH_TIMEOUT_S = 900 + @dataclasses.dataclass class ProcessGroupInfo: @@ -148,7 +151,7 @@ def parallel_launch_with_config( *args: P.args, **kwargs: P.kwargs, ) -> None: - spawn( + context = spawn( _worker_parallel_launch, args=( world_size, @@ -162,5 +165,16 @@ def parallel_launch_with_config( ) + args, nprocs=world_size, - join=True, + join=False, ) + # A worker stuck in communicator setup never exits, so bound the wait. + deadline = time.monotonic() + LAUNCH_TIMEOUT_S + while not context.join(timeout=max(0.0, deadline - time.monotonic())): + if time.monotonic() >= deadline: + for process in context.processes: + process.kill() + for process in context.processes: + process.join() + raise TimeoutError( + f"Parallel workers did not finish within {LAUNCH_TIMEOUT_S}s" + ) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index ba79cc5c34..91a6bc3f70 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -127,6 +127,14 @@ if has_nixl_ep(): BACKENDS += ["nixl_ep"] DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"} + +# These tests run on one node, where NVSHMEM reaches peers over NVLink, so +# keep it off the NICs to avoid depending on IB setup on shared CI hosts. +DEEPEP_SINGLE_NODE_NVSHMEM_ENV = { + "NVSHMEM_REMOTE_TRANSPORT": "none", + # DeepEP forces IBGDA on, so also hide the NICs from it. + "NVSHMEM_HCA_LIST": "none", +} MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"} QUANT_METHODS = [ @@ -2187,6 +2195,8 @@ def test_moe_layer( rocm_deepep = current_platform.is_rocm() and backend in DEEPEP_BACKENDS parallel_worker = _parallel_worker_rocm_deepep if rocm_deepep else _parallel_worker + cuda_deepep = current_platform.is_cuda() and backend in DEEPEP_BACKENDS + env_dict = DEEPEP_SINGLE_NODE_NVSHMEM_ENV if cuda_deepep else None launch_failures: list[str] = [] try: @@ -2197,7 +2207,7 @@ def test_moe_layer( world_size, parallel_worker, vllm_config, - None, + env_dict, test_config_batch, verbosity, failure_report_path=failure_report_path,