diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index fcee2c1daef..2c9919bedfb 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1721,7 +1721,7 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]: if will_log_this_step: self.step_start_event.record() - result = await self.controller.async_generate_output_tokens_dynamic_batch() + result = await self.controller.async_generate_output_tokens_dynamic_batch(loop=self._loop) if will_log_this_step: self.step_end_event.record() self.step_end_event.synchronize() diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 87edddea566..add71b7ec8c 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -27,6 +27,7 @@ ) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.utils import ( + GPUFuture, get_attention_mask, set_decode_expert_padding, set_moe_metadata_sync, @@ -1693,12 +1694,15 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: } async def async_generate_output_tokens_dynamic_batch( - self, skip_bookkeeping: Optional[bool] = False + self, + skip_bookkeeping: Optional[bool] = False, + loop: Optional[asyncio.AbstractEventLoop] = None, ) -> Optional[Dict]: """Forward step the model and update the inference context. Args: skip_bookkeeping (Optional[bool]): If true, skip the context bookkeeping step. + loop (Optional[asyncio.AbstractEventLoop]): Event loop to use for GPU synchronization. Return: (Optional[Dict]): A dictionary containing: @@ -1716,8 +1720,12 @@ async def async_generate_output_tokens_dynamic_batch( if context.active_token_count == 0 and active_request_count == 0: return None + loop = get_asyncio_loop(loop) + gpu_done = GPUFuture(loop, interrupt_event_loop=True) + with torch.inference_mode(): input_ids, position_ids = self._dynamic_step_context_init() + return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping() cuda_graph_request_count = ( context.padded_active_request_count @@ -1748,24 +1756,25 @@ async def async_generate_output_tokens_dynamic_batch( context.kv_block_allocator.store_routing_per_block(self._router_record_bookkeeping()) range_pop() - # This is the best place to yield control back to event loop. - # At this point we have enqueued FW pass GPU kernels asynchronously. - # While they are running, we can do other useful CPU work. - # Note: This can be moved further ahead if sampling can be made - # asynchronous. - # Todo [Siddharth]: Can we condition the sleep on a cuda event? - # NOTE [TDE]: This will be moved once CPU and GPU methods are separated. - await asyncio.sleep(0) + # Trigger a faux-interrupt on the event loop after the forward pass GPU work completes. + # The actual await happens after the CPU has also enqueued sampling work. + # This allows the GPU to continue working (on sampling) while the faux-interrupt is handled. + gpu_done.record() with torch.inference_mode(): - range_push("sampling") - return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping() - if self.num_speculative_tokens > 0: # Phase 1: Verify speculative tokens using base logits only. nvtx_range_push("mtp-spec-decoding/verify") self._dynamic_step_sample_logits_and_verify_tokens(input_ids) nvtx_range_pop("mtp-spec-decoding/verify") + else: + nvtx_range_push("sampling") + self._dynamic_step_sample_logits() + nvtx_range_pop("sampling") + + await gpu_done + + if self.num_speculative_tokens > 0: # Phase 2: Rewind KV cache for rejected tokens. nvtx_range_push("mtp-spec-decoding/rewind-kv-cache") blocks_to_release, remove_mask = self._rewind_kv_cache() @@ -1785,8 +1794,6 @@ async def async_generate_output_tokens_dynamic_batch( # Phase 4: Release freed blocks. Deferred from Phase 2 so the # data-dependent boolean-mask sync overlaps with MTP GPU work. context.kv_block_allocator.release_memory_blocks(blocks_to_release[remove_mask]) - else: - self._dynamic_step_sample_logits() log_probs = None top_n_logprobs = None @@ -1805,7 +1812,6 @@ async def async_generate_output_tokens_dynamic_batch( top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs( log_probs_tensor ) - range_pop() if skip_bookkeeping: # _transfer_samples_to_cpu wasn't invoked on this path, so do diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py index 11973551aa2..33808893a62 100644 --- a/megatron/core/inference/utils.py +++ b/megatron/core/inference/utils.py @@ -1,10 +1,12 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. import asyncio +import ctypes import logging import multiprocessing import sys from importlib.metadata import PackageNotFoundError, version +from typing import Optional import torch @@ -302,3 +304,164 @@ def shutdown(self): else: asyncio_QueueShutDown = asyncio.QueueShutDown asyncio_Queue = asyncio.Queue + + +_libcudart: Optional[ctypes.CDLL] = None +_CUDA_HOST_FN_T = ctypes.CFUNCTYPE(None, ctypes.c_void_p) + + +def _get_cudart() -> ctypes.CDLL: + """Lazily load and configure the CUDA runtime library.""" + global _libcudart + if _libcudart is None: + cuda_major = torch.version.cuda.split('.')[0] + _libcudart = ctypes.CDLL(f"libcudart.so.{cuda_major}") + _libcudart.cudaLaunchHostFunc.restype = ctypes.c_int + _libcudart.cudaLaunchHostFunc.argtypes = [ + ctypes.c_void_p, # cudaStream_t + _CUDA_HOST_FN_T, # cudaHostFn_t + ctypes.c_void_p, # void* userData + ] + return _libcudart + + +class PriorityEventLoop(asyncio.SelectorEventLoop): + """Event loop with front-of-queue scheduling for GPU completion callbacks. + + Adds `call_soon_front` and `call_soon_threadsafe_front` which prepend callbacks to the ready + queue, as opposed to appending them. This acts as a pseudo-interrupt. + """ + + def call_soon_front(self, callback, *args, context=None): + """Schedule `callback` at the front of the ready queue.""" + self._check_closed() + handle = asyncio.events.Handle(callback, args, self, context) + self._ready.appendleft(handle) + return handle + + def call_soon_threadsafe_front(self, callback, *args, context=None): + """Thread-safe variant of `call_soon_front`.""" + self._check_closed() + handle = asyncio.events.Handle(callback, args, self, context) + self._ready.appendleft(handle) + self._write_to_self() # wake select() + return handle + + +class GPUFuture: + """Awaitable that resolves when all preceding work on a CUDA stream completes. + + Instead of blocking on the CPU thread, this future fires a callback from CUDA's internal thread. + + When `interrupt_event_loop` is True, the callback is prepended to the event loop queue. + When False (default), the callback is treated like a normal asyncio callback. + Note that `interrupt_event_loop=True` requires a `PriorityEventLoop`. + + Usage: + + gpu_done = GPUFuture(loop) + # Enqueue the main GPU work. + launch_long_gpu_work(...) + # Record the callback. + gpu_done.record() + # Recommendation: enqueue a short GPU workload to cover the latency of the callback. + launch_short_gpu_work(...) + await gpu_done # resumes as soon as GPU finishes. + """ + + _prevent_gc: dict = {} + _asyncio_future_blocking = False + + def __init__(self, loop: asyncio.AbstractEventLoop, interrupt_event_loop: bool = False): + self._loop = loop + self._interrupt_event_loop = interrupt_event_loop + self._done = False + self._callbacks: list = [] + + def get_loop(self): + """Return the event loop this future is bound to.""" + return self._loop + + def done(self): + """Return True once the GPU stream has drained past the recorded callback.""" + return self._done + + def cancelled(self): + """A `GPUFuture` cannot be cancelled; always returns False.""" + return False + + def cancel(self, msg=None): + """A `GPUFuture` cannot be cancelled; always returns False.""" + return False + + def result(self): + """Return None once done; raises `InvalidStateError` if awaited too early.""" + if not self._done: + raise asyncio.InvalidStateError('not ready') + return None + + def _schedule(self, fn, *, context=None): + """Schedule `fn` on the loop, at the front of the queue if interrupting.""" + if self._interrupt_event_loop: + self._loop.call_soon_front(fn, self, context=context) + else: + self._loop.call_soon(fn, self, context=context) + + def _schedule_threadsafe(self, fn): + """Thread-safe variant of `_schedule`, without the future arg / context.""" + if self._interrupt_event_loop: + self._loop.call_soon_threadsafe_front(fn) + else: + self._loop.call_soon_threadsafe(fn) + + def add_done_callback(self, fn, *, context=None): + """Register `fn` to run when the future resolves; fires immediately if already done.""" + if self._done: + self._schedule(fn, context=context) + else: + self._callbacks.append((fn, context)) + + def __await__(self): + if not self._done: + self._asyncio_future_blocking = True + yield self + if not self._done: + raise RuntimeError("await wasn't used with future") + return self.result() + + def record(self, stream: Optional[torch.cuda.Stream] = None) -> None: + """Enqueue a host callback that resolves this future when stream drains. + + Args: + stream: CUDA stream to attach to. Defaults to the current stream. + """ + if stream is None: + stream = torch.cuda.current_stream() + + prevent_gc_ref: Optional[object] = None + + def _host_fn(_user_data: ctypes.c_void_p) -> None: + # Runs on CUDA's internal callback thread; MUST NOT call CUDA API. + try: + self._schedule_threadsafe(self._resolve) + except RuntimeError: + pass # event loop closed + GPUFuture._prevent_gc.pop(id(prevent_gc_ref), None) + + c_fn = _CUDA_HOST_FN_T(_host_fn) + prevent_gc_ref = c_fn + GPUFuture._prevent_gc[id(c_fn)] = c_fn + + err = _get_cudart().cudaLaunchHostFunc( + ctypes.c_void_p(stream.cuda_stream), c_fn, ctypes.c_void_p(0) + ) + if err != 0: + GPUFuture._prevent_gc.pop(id(c_fn), None) + raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}") + + def _resolve(self): + """Mark done and fire pending callbacks.""" + self._done = True + cbs, self._callbacks = self._callbacks, [] + for fn, ctx in cbs: + self._schedule(fn, context=ctx) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index d04fd180bb2..fd82f010c86 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2305,7 +2305,9 @@ def get_asyncio_loop(loop: asyncio.AbstractEventLoop | None = None) -> asyncio.A if _ASYNC_IO_LOOP is not None: return _ASYNC_IO_LOOP else: - _ASYNC_IO_LOOP = loop = asyncio.new_event_loop() + from megatron.core.inference.utils import PriorityEventLoop + + _ASYNC_IO_LOOP = loop = PriorityEventLoop() asyncio.set_event_loop(loop) return loop