Skip to content

vllm.distributed.eplb.eplb_utils

Utility functions for EPLB (Expert Parallel Load Balancing).

Classes:

  • CpuGpuEvent –

    Combines a CUDA event with a CPU threading event to enforce record->wait

Functions:

CpuGpuEvent

Combines a CUDA event with a CPU threading event to enforce record->wait ordering across two threads.

This class is designed for exactly two threads: one producer that calls record() and one consumer that calls wait(). Using it with more than two threads is not supported and will produce undefined behavior.

CUDA events alone are insufficient for cross-thread synchronization because waiting on an unrecorded CUDA event is a no-op. The wait will return immediately instead of blocking. This class adds a threading.Event so that the waiting thread blocks on the CPU side until record() is called, at which point the CUDA event is guaranteed to be in-flight and event.wait() will correctly synchronize the GPU stream.

Methods:

  • record –

    Unblocks the waiting thread after calling event.record().

  • wait –

    Blocks the calling thread until record finishes. Used to guarantee that the

Source code in vllm/distributed/eplb/eplb_utils.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class CpuGpuEvent:
    """
    Combines a CUDA event with a CPU threading event to enforce record->wait
    ordering across two threads.

    This class is designed for exactly two threads: one producer that calls
    record() and one consumer that calls wait(). Using it with more than two
    threads is not supported and will produce undefined behavior.

    CUDA events alone are insufficient for cross-thread synchronization because
    waiting on an unrecorded CUDA event is a no-op. The wait will return
    immediately instead of blocking. This class adds a threading.Event so
    that the waiting thread blocks on the CPU side until record() is called, at
    which point the CUDA event is guaranteed to be in-flight and event.wait() will
    correctly synchronize the GPU stream.
    """

    def __init__(self):
        self._event = torch.cuda.Event()
        self._recorded = threading.Event()

    def wait(self, stream: torch.cuda.Stream | None = None):
        """
        Blocks the calling thread until record finishes. Used to guarantee that the
        record kernel is called before wait.

        Should only be called by the Async Eplb thread.
        """
        self._recorded.wait()
        self._event.wait(stream)
        self._recorded.clear()

    def record(self, stream: torch.cuda.Stream | None = None):
        """
        Unblocks the waiting thread after calling event.record().

        Should only be called by the main thread.
        """
        if self._recorded.is_set():
            raise RuntimeError(
                "CpuGpuEvent.record() called before the previous event was "
                "consumed by wait()"
            )
        self._event = torch.cuda.Event()
        self._event.record(stream)
        self._recorded.set()

record(stream=None)

Unblocks the waiting thread after calling event.record().

Should only be called by the main thread.

Source code in vllm/distributed/eplb/eplb_utils.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def record(self, stream: torch.cuda.Stream | None = None):
    """
    Unblocks the waiting thread after calling event.record().

    Should only be called by the main thread.
    """
    if self._recorded.is_set():
        raise RuntimeError(
            "CpuGpuEvent.record() called before the previous event was "
            "consumed by wait()"
        )
    self._event = torch.cuda.Event()
    self._event.record(stream)
    self._recorded.set()

wait(stream=None)

Blocks the calling thread until record finishes. Used to guarantee that the record kernel is called before wait.

Should only be called by the Async Eplb thread.

Source code in vllm/distributed/eplb/eplb_utils.py 
37
38
39
40
41
42
43
44
45
46
def wait(self, stream: torch.cuda.Stream | None = None):
    """
    Blocks the calling thread until record finishes. Used to guarantee that the
    record kernel is called before wait.

    Should only be called by the Async Eplb thread.
    """
    self._recorded.wait()
    self._event.wait(stream)
    self._recorded.clear()

override_envs_for_eplb(parallel_config, moe_backend=None)

Override environment variables for EPLB when specific conditions are met.

Parameters:

  • parallel_config

    (ParallelConfig) –

    The parallel configuration object.

  • moe_backend

    (str | None, default: None ) –

    The configured MoE backend (e.g. deep_gemm_mega_moe).

Source code in vllm/distributed/eplb/eplb_utils.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def override_envs_for_eplb(
    parallel_config: ParallelConfig,
    moe_backend: str | None = None,
) -> None:
    """
    Override environment variables for EPLB when specific conditions are met.

    Args:
        parallel_config: The parallel configuration object.
        moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``).
    """
    is_data_parallel = parallel_config.data_parallel_size > 1
    is_eplb_enabled = parallel_config.enable_eplb
    is_mega_moe = moe_backend == "deep_gemm_mega_moe"
    is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in (
        "torch_nccl",
        "pynccl",
    )

    # Override NCCL_MAX_CTAS to avoid hangs when EPLB's NCCL weight exchange
    # contends with MoE backend's cooperative-launch on GPU SMs.
    #
    # DeepGEMM Mega MoE uses cooperative launch, which tries to reserve a
    # large fraction of the GPU's SMs. If those SMs are occupied by NCCL,
    # the cooperative launch blocks until enough SMs are freed, causing a
    # deadlock. Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for
    # the cooperative kernel to launch and complete.
    if (
        is_data_parallel
        and is_eplb_enabled
        and is_nccl_based_eplb_communicator
        and is_mega_moe
    ):
        current_value_str = os.getenv("NCCL_MAX_CTAS")

        if current_value_str and current_value_str.isdigit():
            return

        override_value = 8
        os.environ["NCCL_MAX_CTAS"] = str(override_value)
        logger.info_once(
            f"EPLB: Setting NCCL_MAX_CTAS={override_value} "
            f"for expert parallel with NCCL-based EPLB communicator and "
            f"cooperative MoE backend (deep_gemm_mega_moe)",
            scope="global",
        )