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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
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
    async_eplb = parallel_config.eplb_config.use_async
    is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency"
    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.
    #
    # DeepEP low-latency:
    # The hang happens when two ranks interleave kernel launches differently
    # between NCCL collectives (used by async EPLB weight exchange) and DeepEP
    # low-latency (LL) kernels. DeepEP LL uses a cooperative launch and tries
    # to reserve a large fraction of the GPU's SMs; if those SMs are currently
    # occupied by NCCL, the DeepEP LL launch blocks until enough SMs are
    # freed.
    #
    # If rank A enters DeepEP LL in main thread while rank B is still executing
    # NCCL in async thread, rank A can block waiting for SMs, while rank B can
    # block inside NCCL waiting for rank A to participate in the collective.
    # This circular wait causes a deadlock.
    # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP
    # cooperative kernel to launch and complete, breaking the deadlock.
    # See: https://github.com/deepseek-ai/DeepEP/issues/496
    #
    # DeepGEMM Mega MoE also uses cooperative launch and will cause hang even
    # with sync EPLB.
    if (
        is_data_parallel
        and is_eplb_enabled
        and is_nccl_based_eplb_communicator
        and ((is_deepep_ll and async_eplb) or 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)
        backend = "deepep_low_latency" if is_deepep_ll else "deep_gemm_mega_moe"
        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 ({backend})",
            scope="global",
        )