| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import threading |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm import forward_context |
| 8 | from vllm.forward_context import ForwardContext |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.utils.torch_utils import current_stream |
| 11 | |
| 12 | logger = init_logger(__name__) |
| 13 | |
| 14 | _THREAD_ID_TO_CONTEXT: dict = {} |
| 15 | # Here we hardcode the number of microbatches to 2 for default. |
| 16 | _NUM_UBATCHES: int = 2 |
| 17 | _CURRENT_CONTEXTS: list["UBatchContext | None"] = [] |
| 18 | |
| 19 | |
| 20 | class UBatchContext: |
| 21 | """Context manager for micro-batching synchronization using threading events.""" |
| 22 | |
| 23 | def __init__( |
| 24 | self, |
| 25 | id: int, |
| 26 | comm_stream: torch.cuda.Stream, |
| 27 | compute_stream: torch.cuda.Stream, |
| 28 | forward_context: ForwardContext, |
| 29 | ready_barrier: threading.Barrier, |
| 30 | cpu_wait_event: threading.Event, |
| 31 | cpu_signal_event: threading.Event, |
| 32 | gpu_comm_done_event: torch.Event, |
| 33 | gpu_compute_done_event: torch.Event, |
| 34 | schedule: str = "default", |
| 35 | ): |
| 36 | self.id = id |
| 37 | self.comm_stream = comm_stream |
| 38 | self.compute_stream = compute_stream |
| 39 | self.forward_context = forward_context |
| 40 | self.ready_barrier = ready_barrier |
| 41 | self.cpu_wait_event = cpu_wait_event |
| 42 | self.cpu_signal_event = cpu_signal_event |
| 43 | self.current_stream = compute_stream |
| 44 | self.gpu_comm_done_event = gpu_comm_done_event |
| 45 | self.gpu_compute_done_event = gpu_compute_done_event |
| 46 | self.schedule = schedule |
| 47 | self.recv_hook = None |
| 48 | |
| 49 | def __enter__(self): |
| 50 | global _CURRENT_CONTEXTS, _THREAD_ID_TO_CONTEXT |
| 51 | _THREAD_ID_TO_CONTEXT[threading.get_ident()] = self.id |
| 52 | _CURRENT_CONTEXTS[self.id] = self |
| 53 | # _NUM_UBATCHES is set in make_ubatch_contexts |
| 54 | self.ready_barrier.wait() |
| 55 | |
| 56 | self.cpu_wait_event.wait() |
| 57 | self.cpu_wait_event.clear() |
| 58 | self._restore_context() |
| 59 | # Assume we want to start on the compute stream |
| 60 | self.update_stream(self.compute_stream) |
| 61 | return self |
| 62 | |
| 63 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 64 | global _CURRENT_CONTEXTS, _THREAD_ID_TO_CONTEXT |
| 65 | _CURRENT_CONTEXTS[self.id] = None |
| 66 | del _THREAD_ID_TO_CONTEXT[threading.get_ident()] |
| 67 | self.maybe_run_recv_hook() |
| 68 | self.cpu_signal_event.set() |
| 69 | self.cpu_wait_event.clear() |
| 70 | return False |
| 71 | |
| 72 | def _restore_context(self): |
| 73 | forward_context._forward_context = self.forward_context |
| 74 | |
| 75 | def update_stream(self, stream): |
| 76 | self.current_stream = stream |
| 77 | if current_stream() != self.current_stream: |
| 78 | torch.cuda.set_stream(self.current_stream) |
| 79 | |
| 80 | def _signal_comm_done(self): |
| 81 | self.gpu_comm_done_event.record(self.comm_stream) |
| 82 | |
| 83 | def _signal_compute_done(self): |
| 84 | self.gpu_compute_done_event.record(self.compute_stream) |
| 85 | |
| 86 | def _wait_compute_done(self): |
| 87 | self.comm_stream.wait_event(self.gpu_compute_done_event) |
| 88 | |
| 89 | def _wait_comm_done(self): |
| 90 | self.compute_stream.wait_event(self.gpu_comm_done_event) |
| 91 | |
| 92 | def _cpu_yield(self): |
| 93 | # It is critical for correctness that only one thread is running |
| 94 | # at a time. These asserts just make sure that this is the only |
| 95 | # thread running before waking the other one up and going to sleep |
| 96 | assert forward_context._forward_context == self.forward_context |
| 97 | assert current_stream() == self.current_stream |
| 98 | assert not self.cpu_wait_event.is_set() |
| 99 | |
| 100 | self.cpu_signal_event.set() |
| 101 | self.cpu_wait_event.wait() |
| 102 | self.cpu_wait_event.clear() |
| 103 | self._restore_context() |
| 104 | |
| 105 | def switch_to_comm(self): |
| 106 | self.update_stream(self.comm_stream) |
| 107 | |
| 108 | def switch_to_compute(self): |
| 109 | self.update_stream(self.compute_stream) |
| 110 | |
| 111 | def switch_to_comm_sync(self): |
| 112 | self._signal_compute_done() |
| 113 | self.update_stream(self.comm_stream) |
| 114 | self._wait_compute_done() |
| 115 | |
| 116 | def switch_to_compute_sync(self): |
| 117 | self._signal_comm_done() |
| 118 | self.update_stream(self.compute_stream) |
| 119 | self._wait_comm_done() |
| 120 | |
| 121 | def maybe_run_recv_hook(self): |
| 122 | if self.recv_hook is not None: |
| 123 | self.recv_hook() |
| 124 | self.recv_hook = None |
| 125 | |
| 126 | def yield_(self): |
| 127 | self.current_stream = current_stream() |
| 128 | self._cpu_yield() |
| 129 | self.update_stream(self.current_stream) |
| 130 | |
| 131 | def yield_and_switch_from_compute_to_comm(self): |
| 132 | assert current_stream() == self.compute_stream |
| 133 | self._signal_compute_done() |
| 134 | self._cpu_yield() |
| 135 | assert self.current_stream == self.compute_stream |
| 136 | self.update_stream(self.comm_stream) |
| 137 | self._wait_compute_done() |
| 138 | |
| 139 | def yield_and_switch_from_comm_to_compute(self): |
| 140 | assert current_stream() == self.comm_stream |
| 141 | self._signal_comm_done() |
| 142 | self._cpu_yield() |
| 143 | assert self.current_stream == self.comm_stream |
| 144 | self.update_stream(self.compute_stream) |
| 145 | self._wait_comm_done() |
| 146 | |
| 147 | |
| 148 | def dbo_enabled() -> bool: |
| 149 | return len(_THREAD_ID_TO_CONTEXT) > 0 |
| 150 | |
| 151 | |
| 152 | def dbo_current_ubatch_id() -> int: |
| 153 | if len(_THREAD_ID_TO_CONTEXT) == 0: |
| 154 | return 0 |
| 155 | return _THREAD_ID_TO_CONTEXT[threading.get_ident()] |
| 156 | |
| 157 | |
| 158 | def _register_ubatch_function(func): |
| 159 | def wrapper(*args, **kwargs): |
| 160 | if len(_THREAD_ID_TO_CONTEXT) > 0: |
| 161 | ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()] |
| 162 | ctx = _CURRENT_CONTEXTS[ctx_idx] |
| 163 | func(ctx, *args, **kwargs) |
| 164 | |
| 165 | return wrapper |
| 166 | |
| 167 | |
| 168 | dbo_maybe_run_recv_hook = _register_ubatch_function(UBatchContext.maybe_run_recv_hook) |
| 169 | dbo_yield = _register_ubatch_function(UBatchContext.yield_) |
| 170 | dbo_yield_and_switch_from_compute_to_comm = _register_ubatch_function( |
| 171 | UBatchContext.yield_and_switch_from_compute_to_comm |
| 172 | ) |
| 173 | dbo_yield_and_switch_from_comm_to_compute = _register_ubatch_function( |
| 174 | UBatchContext.yield_and_switch_from_comm_to_compute |
| 175 | ) |
| 176 | dbo_switch_to_comm = _register_ubatch_function(UBatchContext.switch_to_comm) |
| 177 | dbo_switch_to_compute = _register_ubatch_function(UBatchContext.switch_to_compute) |
| 178 | dbo_switch_to_comm_sync = _register_ubatch_function(UBatchContext.switch_to_comm_sync) |
| 179 | dbo_switch_to_compute_sync = _register_ubatch_function( |
| 180 | UBatchContext.switch_to_compute_sync |
| 181 | ) |
| 182 | |
| 183 | |
| 184 | def dbo_register_recv_hook(recv_hook): |
| 185 | if len(_THREAD_ID_TO_CONTEXT) > 0: |
| 186 | ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()] |
| 187 | next_ctx = _CURRENT_CONTEXTS[(ctx_idx + 1) % _NUM_UBATCHES] |
| 188 | next_ctx.recv_hook = recv_hook |
| 189 | |
| 190 | |
| 191 | def dbo_get_previous_event(func, *args, **kwargs): |
| 192 | if len(_THREAD_ID_TO_CONTEXT) > 0: |
| 193 | ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()] |
| 194 | ctx = _CURRENT_CONTEXTS[ctx_idx] |
| 195 | # execute callable on the ubatch compute stream to record/wait events there |
| 196 | with torch.cuda.stream(ctx.compute_stream): |
| 197 | return func(*args, **kwargs) |
| 198 | |
| 199 | |
| 200 | def make_ubatch_contexts( |
| 201 | num_micro_batches: int, |
| 202 | compute_stream: torch.cuda.Stream, |
| 203 | comm_stream: torch.cuda.Stream, |
| 204 | forward_contexts: list[ForwardContext], |
| 205 | ready_barrier: threading.Barrier, |
| 206 | schedule: str = "default", |
| 207 | ) -> list[UBatchContext]: |
| 208 | global _NUM_UBATCHES, _CURRENT_CONTEXTS |
| 209 | assert num_micro_batches > 1, "num_micro_batches must be greater than 1" |
| 210 | |
| 211 | _NUM_UBATCHES = num_micro_batches |
| 212 | # Ensure the global context list is large enough |
| 213 | if len(_CURRENT_CONTEXTS) < num_micro_batches: |
| 214 | _CURRENT_CONTEXTS.extend([None] * (num_micro_batches - len(_CURRENT_CONTEXTS))) |
| 215 | |
| 216 | """ |
| 217 | Create a context manager for micro-batching synchronization. |
| 218 | """ |
| 219 | cpu_events = [threading.Event() for _ in range(num_micro_batches)] |
| 220 | gpu_comm_done_events = [torch.Event() for _ in range(num_micro_batches)] |
| 221 | gpu_compute_done_events = [torch.Event() for _ in range(num_micro_batches)] |
| 222 | |
| 223 | ctxs = [] |
| 224 | for i in range(num_micro_batches): |
| 225 | ctx = UBatchContext( |
| 226 | id=i, |
| 227 | compute_stream=compute_stream, |
| 228 | comm_stream=comm_stream, |
| 229 | forward_context=forward_contexts[i], |
| 230 | ready_barrier=ready_barrier, |
| 231 | cpu_wait_event=cpu_events[i], |
| 232 | cpu_signal_event=cpu_events[(i + 1) % num_micro_batches], |
| 233 | gpu_comm_done_event=gpu_comm_done_events[i], |
| 234 | gpu_compute_done_event=gpu_compute_done_events[i], |
| 235 | schedule=schedule, |
| 236 | ) |
| 237 | ctxs.append(ctx) |
| 238 | |
| 239 | return ctxs |