Allow inference to asynchronously await GPU work completion

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()

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

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)

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