From 27aca2621d3d6fb54fb15ae3dd990984b2b747d6 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 14 Apr 2026 03:45:15 -0500
Subject: [PATCH 01/10] Allow GPU work exclusive ownership of asyncio loop

---
 .../gpu_event_loop_synchronization.py         | 285 ++++++++++++++++++
 megatron/core/utils.py                        |   6 +-
 2 files changed, 290 insertions(+), 1 deletion(-)
 create mode 100644 megatron/core/inference/gpu_event_loop_synchronization.py

diff --git a/megatron/core/inference/gpu_event_loop_synchronization.py b/megatron/core/inference/gpu_event_loop_synchronization.py
new file mode 100644
index 00000000000..f03f3a1abe0
--- /dev/null
+++ b/megatron/core/inference/gpu_event_loop_synchronization.py
@@ -0,0 +1,285 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import asyncio
+import collections
+import contextlib
+import ctypes
+import sys
+from typing import Optional
+
+import torch
+
+_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 GpuFuture:
+    """Awaitable that resolves when all preceding work on a CUDA stream completes.
+
+    If the event loop is an `ExclusiveTaskEventLoop` and exclusive mode is active,
+    awaiting this future temporarily releases exclusivity so other tasks can run while the GPU is
+    busy, then re-acquires it on completion. Pass `yield_exclusive=False` to suppress this behavior.
+
+    Usage:
+        gpu_done = GpuFuture(loop)
+        launch_gpu_work(...)
+        gpu_done.record()
+        launch_more_gpu_work(...)
+        await gpu_done          # releases exclusivity while waiting, if active
+    """
+
+    # Prevent garbage-collection of live ctypes callbacks.
+    # Keyed by id() because ctypes function pointers are not hashable.
+    _prevent_gc: dict = {}
+
+    def __init__(self, loop: asyncio.AbstractEventLoop, yield_exclusive: bool = True):
+        self._loop = loop
+        self._future: asyncio.Future = loop.create_future()
+        self._yield_exclusive = yield_exclusive
+
+    def record(self, stream: Optional[torch.cuda.Stream] = None) -> None:
+        """Enqueue a host callback that resolves this future.
+
+        Args:
+            stream: CUDA stream to attach to. Defaults to the current stream.
+        """
+        if stream is None:
+            stream = torch.cuda.current_stream()
+
+        # This closure prevents the ctypes wrapper from being collected
+        # while the callback is in the 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 any CUDA API.
+            try:
+                self._loop.call_soon_threadsafe(self._future.set_result, None)
+            except RuntimeError:
+                # Event loop closed; nothing to do.
+                pass
+            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:
+            # Allow the callback to be garbage-collected on failure.
+            GpuFuture._prevent_gc.pop(id(c_fn), None)
+            raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}")
+
+    def __await__(self):
+        ready = getattr(self._loop, '_ready', None)
+        exclusive_task = getattr(ready, '_exclusive_task', None) if ready else None
+        should_yield = self._yield_exclusive and exclusive_task is not None
+
+        if should_yield:
+            ready.clear_exclusive()
+        result = yield from self._future.__await__()
+        if should_yield:
+            ready.set_exclusive(exclusive_task)
+        return result
+
+
+def _verify_task_callback_valid(loop: asyncio.AbstractEventLoop) -> None:
+    """Verify that Task step callbacks expose __self__ on this Python runtime.
+
+    This currently works on all Python versions, but is not guaranteed to be stable.
+    Internal API changes can break the mechanism at any point.
+
+    This function creates a throwaway task, inspects its callback in `_ready`,
+    and raises `RuntimeError` if the internal API pattern no longer holds.
+    """
+
+    async def _canary():
+        pass
+
+    task = loop.create_task(_canary())
+
+    # The task's __step should now be in _ready.  Inspect without running it.
+    found = False
+    for handle in loop._ready:
+        cb = handle._callback
+        owner = getattr(cb, '__self__', None)
+        if owner is task:
+            found = True
+            break
+
+    # Cancel the canary — it will be cleaned up on the next loop iteration.
+    task.cancel()
+
+    if not found:
+        raise RuntimeError(
+            f"ExclusiveTaskEventLoop cannot identify Task ownership of callbacks "
+            f"on this Python runtime (Python {sys.version}).  "
+            f"Task step callbacks do not expose __self__.  "
+            f"Exclusive-task scheduling requires this for correctness."
+        )
+
+
+class _PriorityReadyQueue:  # pylint: disable=missing-function-docstring
+    """Drop-in replacement for the event loop's `_ready` deque.
+
+    When an exclusive task is set, callbacks are routed into two categories at `append` time:
+
+    - Allowed: the exclusive task's own callbacks. These enter the live `_queue`.
+    - Deferred: every other task's callbacks.
+      These accumulate in `_deferred` and are drained back into `_queue` when exclusive mode ends.
+
+    The event loop's `run_once` only ever sees `_queue` via the standard deque interface,
+    so no event-loop internals need to be patched.
+    """
+
+    __slots__ = ('_queue', '_deferred', '_exclusive_task')
+
+    def __init__(self) -> None:
+        self._queue: collections.deque = collections.deque()
+        self._deferred: collections.deque = collections.deque()
+        self._exclusive_task: Optional[asyncio.Task] = None
+
+    def set_exclusive(self, task: asyncio.Task) -> None:
+        """Activate exclusive mode for `task`."""
+        assert self._exclusive_task is None, "Exclusive task already set"
+        self._exclusive_task = task
+
+    def clear_exclusive(self) -> None:
+        """Deactivate exclusive mode and drain deferred callbacks."""
+        self._exclusive_task = None
+        if self._deferred:
+            self._queue.extend(self._deferred)
+            self._deferred.clear()
+
+    def _is_allowed(self, handle) -> bool:
+        if self._exclusive_task is None:
+            return True
+
+        cb = handle._callback
+        owner = getattr(cb, '__self__', None)
+
+        # The exclusive task's own __step callback.
+        if owner is self._exclusive_task:
+            return True
+
+        # Non-Task callbacks: Future resolution, I/O dispatch, signal handlers, call_soon_threadsafe
+        # Infrastructure that the exclusive task depends on to function.
+        if not isinstance(owner, asyncio.Task):
+            return True
+
+        # A different task's __step; defer it.
+        return False
+
+    def append(self, handle) -> None:
+        if self._is_allowed(handle):
+            self._queue.append(handle)
+        else:
+            self._deferred.append(handle)
+
+    def appendleft(self, handle) -> None:
+        if self._is_allowed(handle):
+            self._queue.appendleft(handle)
+        else:
+            self._deferred.appendleft(handle)
+
+    def popleft(self):
+        return self._queue.popleft()
+
+    def extend(self, iterable) -> None:
+        for item in iterable:
+            self.append(item)
+
+    def remove(self, item) -> None:
+        try:
+            self._queue.remove(item)
+        except ValueError:
+            self._deferred.remove(item)
+
+    def clear(self) -> None:
+        self._queue.clear()
+        self._deferred.clear()
+
+    def __len__(self) -> int:
+        return len(self._queue)
+
+    def __bool__(self) -> bool:
+        return bool(self._queue)
+
+    def __iter__(self):
+        return iter(self._queue)
+
+    def __contains__(self, item) -> bool:
+        return item in self._queue or item in self._deferred
+
+
+class ExclusiveTaskEventLoop(asyncio.SelectorEventLoop):
+    """Event loop with exclusive-task support.
+
+    When `set_exclusive_task` is called, only the designated tasks' callbacks remain visible.
+    This gives the exclusive task immediate, contention-free access to the event loop.
+
+    Usage:
+        loop = ExclusiveTaskEventLoop()
+        asyncio.set_event_loop(loop)
+        [...]
+        # Inside the critical coroutine (i.e. GPU forward pass loop):
+        async with loop.exclusive():
+            while True:
+                gpu_done = GpuFuture(loop)
+                launch_gpu_work_1()         # exclusive — no other task can interleave
+                gpu_done.record()
+                launch_gpu_work_2()         # still exclusive
+                await gpu_done              # yields exclusivity while GPU is busy
+                                            # re-acquires when GPU work 1 completes
+        # Exclusive mode is always released on exit, even on exception.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._ready = _PriorityReadyQueue()
+
+        # Verify that Task callbacks expose __self__ before we rely on it in the ready queue.
+        _verify_task_callback_valid(self)
+
+    @contextlib.asynccontextmanager
+    async def exclusive(self, task: Optional[asyncio.Task] = None):
+        """Async context manager for exclusive-task mode."""
+        if task is None:
+            task = asyncio.current_task()
+        self.set_exclusive_task(task)
+        try:
+            yield
+        finally:
+            self.set_exclusive_task(None)
+
+    def set_exclusive_task(self, task: Optional[asyncio.Task]) -> None:
+        """Set or clear the exclusive task.
+
+        Prefer the :meth:`exclusive` context manager over calling this directly,
+        to ensure exclusive mode is always released.
+
+        Args:
+            task: The `asyncio.Task` that should have exclusive access,
+                or `None` to end the exclusive section and drain deferred callbacks.
+        """
+        if task is None:
+            self._ready.clear_exclusive()
+        else:
+            self._ready.set_exclusive(task)
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index d04fd180bb2..04ab167eb9e 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -2305,7 +2305,11 @@ 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.gpu_event_loop_synchronization import (
+                    ExclusiveTaskEventLoop,
+                )
+
+                _ASYNC_IO_LOOP = loop = ExclusiveTaskEventLoop()
                 asyncio.set_event_loop(loop)
     return loop
 

From 14cbdaf5e3a0ce0f1cccd24bd70f7472d297c9b1 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 14 Apr 2026 05:48:48 -0500
Subject: [PATCH 02/10] Utilize the exclusive GPU work sugar

---
 .../core/inference/engines/dynamic_engine.py  |  5 ++++-
 .../text_generation_controller.py             | 22 +++++++++++--------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index fcee2c1daef..11221300003 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -1721,7 +1721,10 @@ 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()
+        async with self._loop.exclusive():
+            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..4b2b0afa6fb 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -21,6 +21,7 @@
 )
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.contexts.static_context import StaticInferenceContext
+from megatron.core.inference.gpu_event_loop_synchronization import GpuFuture
 from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
@@ -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,6 +1720,9 @@ 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)
+
         with torch.inference_mode():
             input_ids, position_ids = self._dynamic_step_context_init()
 
@@ -1748,14 +1755,11 @@ 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)
+        # Record after forward pass kernels are enqueued.  Awaiting this
+        # yields exclusivity (if active) so other tasks can run while the
+        # GPU is busy, then re-acquires once the forward pass completes.
+        gpu_done.record()
+        await gpu_done
 
         with torch.inference_mode():
             range_push("sampling")

From 22b3bd11c4a1c280fb87f53d2d2308f1036e4793 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 14 Apr 2026 06:54:53 -0500
Subject: [PATCH 03/10] Undo the asyncio exclusive ownership overhead

---
 .../core/inference/engines/dynamic_engine.py  |  16 +-
 .../gpu_event_loop_synchronization.py         | 222 +-----------------
 .../text_generation_controller.py             |   8 +-
 megatron/core/utils.py                        |   6 +-
 4 files changed, 28 insertions(+), 224 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 11221300003..5772564f56c 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -31,6 +31,7 @@
     DataParallelInferenceCoordinator,
 )
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
+from megatron.core.inference.gpu_event_loop_synchronization import GPUFuture
 from megatron.core.inference.headers import Headers, UnknownHeaderError
 from megatron.core.inference.inference_request import (
     DynamicInferenceEvent,
@@ -1721,13 +1722,14 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
 
         if will_log_this_step:
             self.step_start_event.record()
-        async with self._loop.exclusive():
-            result = await self.controller.async_generate_output_tokens_dynamic_batch(
-                loop=self._loop
-            )
+        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()
+            step_done = GPUFuture(self._loop)
+            step_done.record()
+            await step_done
             step_time = self.step_start_event.elapsed_time(self.step_end_event) / 1e3
         else:
             step_time = 0.0
@@ -2379,7 +2381,9 @@ async def run_engine_with_coordinator(
                             self.step_start_event.record()
                             self.controller.dummy_forward()
                             self.step_end_event.record()
-                            self.step_end_event.synchronize()
+                            step_done = GPUFuture(self._loop)
+                            step_done.record()
+                            await step_done
                             self.context.step_count += 1
                             self.context.prefix_cache_lru_clock += 1
                     else:
diff --git a/megatron/core/inference/gpu_event_loop_synchronization.py b/megatron/core/inference/gpu_event_loop_synchronization.py
index f03f3a1abe0..d08b58838f1 100644
--- a/megatron/core/inference/gpu_event_loop_synchronization.py
+++ b/megatron/core/inference/gpu_event_loop_synchronization.py
@@ -1,10 +1,7 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import asyncio
-import collections
-import contextlib
 import ctypes
-import sys
 from typing import Optional
 
 import torch
@@ -28,29 +25,30 @@ def _get_cudart() -> ctypes.CDLL:
     return _libcudart
 
 
-class GpuFuture:
+class GPUFuture:
     """Awaitable that resolves when all preceding work on a CUDA stream completes.
 
-    If the event loop is an `ExclusiveTaskEventLoop` and exclusive mode is active,
-    awaiting this future temporarily releases exclusivity so other tasks can run while the GPU is
-    busy, then re-acquires it on completion. Pass `yield_exclusive=False` to suppress this behavior.
+    Instead of blocking the CPU with ``torch.cuda.synchronize()`` or
+    ``event.synchronize()``, this uses ``cudaLaunchHostFunc`` to enqueue a
+    host-side callback on the CUDA stream.  The callback resolves an asyncio
+    ``Future`` via ``call_soon_threadsafe``, allowing other asyncio tasks to
+    run while the GPU is busy.
 
     Usage:
-        gpu_done = GpuFuture(loop)
+        gpu_done = GPUFuture(loop)
         launch_gpu_work(...)
         gpu_done.record()
         launch_more_gpu_work(...)
-        await gpu_done          # releases exclusivity while waiting, if active
+        await gpu_done          # other tasks run while GPU is busy
     """
 
     # Prevent garbage-collection of live ctypes callbacks.
     # Keyed by id() because ctypes function pointers are not hashable.
     _prevent_gc: dict = {}
 
-    def __init__(self, loop: asyncio.AbstractEventLoop, yield_exclusive: bool = True):
+    def __init__(self, loop: asyncio.AbstractEventLoop):
         self._loop = loop
         self._future: asyncio.Future = loop.create_future()
-        self._yield_exclusive = yield_exclusive
 
     def record(self, stream: Optional[torch.cuda.Stream] = None) -> None:
         """Enqueue a host callback that resolves this future.
@@ -73,213 +71,19 @@ def _host_fn(_user_data: ctypes.c_void_p) -> None:
             except RuntimeError:
                 # Event loop closed; nothing to do.
                 pass
-            GpuFuture._prevent_gc.pop(id(prevent_gc_ref), None)
+            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
+        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:
             # Allow the callback to be garbage-collected on failure.
-            GpuFuture._prevent_gc.pop(id(c_fn), None)
+            GPUFuture._prevent_gc.pop(id(c_fn), None)
             raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}")
 
     def __await__(self):
-        ready = getattr(self._loop, '_ready', None)
-        exclusive_task = getattr(ready, '_exclusive_task', None) if ready else None
-        should_yield = self._yield_exclusive and exclusive_task is not None
-
-        if should_yield:
-            ready.clear_exclusive()
-        result = yield from self._future.__await__()
-        if should_yield:
-            ready.set_exclusive(exclusive_task)
-        return result
-
-
-def _verify_task_callback_valid(loop: asyncio.AbstractEventLoop) -> None:
-    """Verify that Task step callbacks expose __self__ on this Python runtime.
-
-    This currently works on all Python versions, but is not guaranteed to be stable.
-    Internal API changes can break the mechanism at any point.
-
-    This function creates a throwaway task, inspects its callback in `_ready`,
-    and raises `RuntimeError` if the internal API pattern no longer holds.
-    """
-
-    async def _canary():
-        pass
-
-    task = loop.create_task(_canary())
-
-    # The task's __step should now be in _ready.  Inspect without running it.
-    found = False
-    for handle in loop._ready:
-        cb = handle._callback
-        owner = getattr(cb, '__self__', None)
-        if owner is task:
-            found = True
-            break
-
-    # Cancel the canary — it will be cleaned up on the next loop iteration.
-    task.cancel()
-
-    if not found:
-        raise RuntimeError(
-            f"ExclusiveTaskEventLoop cannot identify Task ownership of callbacks "
-            f"on this Python runtime (Python {sys.version}).  "
-            f"Task step callbacks do not expose __self__.  "
-            f"Exclusive-task scheduling requires this for correctness."
-        )
-
-
-class _PriorityReadyQueue:  # pylint: disable=missing-function-docstring
-    """Drop-in replacement for the event loop's `_ready` deque.
-
-    When an exclusive task is set, callbacks are routed into two categories at `append` time:
-
-    - Allowed: the exclusive task's own callbacks. These enter the live `_queue`.
-    - Deferred: every other task's callbacks.
-      These accumulate in `_deferred` and are drained back into `_queue` when exclusive mode ends.
-
-    The event loop's `run_once` only ever sees `_queue` via the standard deque interface,
-    so no event-loop internals need to be patched.
-    """
-
-    __slots__ = ('_queue', '_deferred', '_exclusive_task')
-
-    def __init__(self) -> None:
-        self._queue: collections.deque = collections.deque()
-        self._deferred: collections.deque = collections.deque()
-        self._exclusive_task: Optional[asyncio.Task] = None
-
-    def set_exclusive(self, task: asyncio.Task) -> None:
-        """Activate exclusive mode for `task`."""
-        assert self._exclusive_task is None, "Exclusive task already set"
-        self._exclusive_task = task
-
-    def clear_exclusive(self) -> None:
-        """Deactivate exclusive mode and drain deferred callbacks."""
-        self._exclusive_task = None
-        if self._deferred:
-            self._queue.extend(self._deferred)
-            self._deferred.clear()
-
-    def _is_allowed(self, handle) -> bool:
-        if self._exclusive_task is None:
-            return True
-
-        cb = handle._callback
-        owner = getattr(cb, '__self__', None)
-
-        # The exclusive task's own __step callback.
-        if owner is self._exclusive_task:
-            return True
-
-        # Non-Task callbacks: Future resolution, I/O dispatch, signal handlers, call_soon_threadsafe
-        # Infrastructure that the exclusive task depends on to function.
-        if not isinstance(owner, asyncio.Task):
-            return True
-
-        # A different task's __step; defer it.
-        return False
-
-    def append(self, handle) -> None:
-        if self._is_allowed(handle):
-            self._queue.append(handle)
-        else:
-            self._deferred.append(handle)
-
-    def appendleft(self, handle) -> None:
-        if self._is_allowed(handle):
-            self._queue.appendleft(handle)
-        else:
-            self._deferred.appendleft(handle)
-
-    def popleft(self):
-        return self._queue.popleft()
-
-    def extend(self, iterable) -> None:
-        for item in iterable:
-            self.append(item)
-
-    def remove(self, item) -> None:
-        try:
-            self._queue.remove(item)
-        except ValueError:
-            self._deferred.remove(item)
-
-    def clear(self) -> None:
-        self._queue.clear()
-        self._deferred.clear()
-
-    def __len__(self) -> int:
-        return len(self._queue)
-
-    def __bool__(self) -> bool:
-        return bool(self._queue)
-
-    def __iter__(self):
-        return iter(self._queue)
-
-    def __contains__(self, item) -> bool:
-        return item in self._queue or item in self._deferred
-
-
-class ExclusiveTaskEventLoop(asyncio.SelectorEventLoop):
-    """Event loop with exclusive-task support.
-
-    When `set_exclusive_task` is called, only the designated tasks' callbacks remain visible.
-    This gives the exclusive task immediate, contention-free access to the event loop.
-
-    Usage:
-        loop = ExclusiveTaskEventLoop()
-        asyncio.set_event_loop(loop)
-        [...]
-        # Inside the critical coroutine (i.e. GPU forward pass loop):
-        async with loop.exclusive():
-            while True:
-                gpu_done = GpuFuture(loop)
-                launch_gpu_work_1()         # exclusive — no other task can interleave
-                gpu_done.record()
-                launch_gpu_work_2()         # still exclusive
-                await gpu_done              # yields exclusivity while GPU is busy
-                                            # re-acquires when GPU work 1 completes
-        # Exclusive mode is always released on exit, even on exception.
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._ready = _PriorityReadyQueue()
-
-        # Verify that Task callbacks expose __self__ before we rely on it in the ready queue.
-        _verify_task_callback_valid(self)
-
-    @contextlib.asynccontextmanager
-    async def exclusive(self, task: Optional[asyncio.Task] = None):
-        """Async context manager for exclusive-task mode."""
-        if task is None:
-            task = asyncio.current_task()
-        self.set_exclusive_task(task)
-        try:
-            yield
-        finally:
-            self.set_exclusive_task(None)
-
-    def set_exclusive_task(self, task: Optional[asyncio.Task]) -> None:
-        """Set or clear the exclusive task.
-
-        Prefer the :meth:`exclusive` context manager over calling this directly,
-        to ensure exclusive mode is always released.
-
-        Args:
-            task: The `asyncio.Task` that should have exclusive access,
-                or `None` to end the exclusive section and drain deferred callbacks.
-        """
-        if task is None:
-            self._ready.clear_exclusive()
-        else:
-            self._ready.set_exclusive(task)
+        return self._future.__await__()
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 4b2b0afa6fb..4fda559404f 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -21,7 +21,7 @@
 )
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.contexts.static_context import StaticInferenceContext
-from megatron.core.inference.gpu_event_loop_synchronization import GpuFuture
+from megatron.core.inference.gpu_event_loop_synchronization import GPUFuture
 from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
@@ -1721,7 +1721,7 @@ async def async_generate_output_tokens_dynamic_batch(
             return None
 
         loop = get_asyncio_loop(loop)
-        gpu_done = GpuFuture(loop)
+        gpu_done = GPUFuture(loop)
 
         with torch.inference_mode():
             input_ids, position_ids = self._dynamic_step_context_init()
@@ -1756,8 +1756,8 @@ async def async_generate_output_tokens_dynamic_batch(
             range_pop()
 
         # Record after forward pass kernels are enqueued.  Awaiting this
-        # yields exclusivity (if active) so other tasks can run while the
-        # GPU is busy, then re-acquires once the forward pass completes.
+        # lets other asyncio tasks run while the GPU is busy, and resumes
+        # as soon as the forward pass completes.
         gpu_done.record()
         await gpu_done
 
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index 04ab167eb9e..d04fd180bb2 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -2305,11 +2305,7 @@ def get_asyncio_loop(loop: asyncio.AbstractEventLoop | None = None) -> asyncio.A
             if _ASYNC_IO_LOOP is not None:
                 return _ASYNC_IO_LOOP
             else:
-                from megatron.core.inference.gpu_event_loop_synchronization import (
-                    ExclusiveTaskEventLoop,
-                )
-
-                _ASYNC_IO_LOOP = loop = ExclusiveTaskEventLoop()
+                _ASYNC_IO_LOOP = loop = asyncio.new_event_loop()
                 asyncio.set_event_loop(loop)
     return loop
 

From ea9d427d250ed3ea8b17707b0d4c7463fa118dc8 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 14 Apr 2026 07:17:55 -0500
Subject: [PATCH 04/10] Get rid of new file, now that code is simpler

---
 .../core/inference/engines/dynamic_engine.py  |  7 +-
 .../gpu_event_loop_synchronization.py         | 89 -------------------
 .../text_generation_controller.py             |  2 +-
 megatron/core/inference/utils.py              | 85 ++++++++++++++++++
 4 files changed, 88 insertions(+), 95 deletions(-)
 delete mode 100644 megatron/core/inference/gpu_event_loop_synchronization.py

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 5772564f56c..c8b21e77f4e 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -31,7 +31,6 @@
     DataParallelInferenceCoordinator,
 )
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
-from megatron.core.inference.gpu_event_loop_synchronization import GPUFuture
 from megatron.core.inference.headers import Headers, UnknownHeaderError
 from megatron.core.inference.inference_request import (
     DynamicInferenceEvent,
@@ -44,7 +43,7 @@
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
-from megatron.core.inference.utils import Counter, await_process_call
+from megatron.core.inference.utils import Counter, GPUFuture, await_process_call
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
 from megatron.core.transformer.enums import CudaGraphScope
@@ -1722,9 +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(
-            loop=self._loop,
-        )
+        result = await self.controller.async_generate_output_tokens_dynamic_batch(loop=self._loop)
         if will_log_this_step:
             self.step_end_event.record()
             step_done = GPUFuture(self._loop)
diff --git a/megatron/core/inference/gpu_event_loop_synchronization.py b/megatron/core/inference/gpu_event_loop_synchronization.py
deleted file mode 100644
index d08b58838f1..00000000000
--- a/megatron/core/inference/gpu_event_loop_synchronization.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-
-import asyncio
-import ctypes
-from typing import Optional
-
-import torch
-
-_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 GPUFuture:
-    """Awaitable that resolves when all preceding work on a CUDA stream completes.
-
-    Instead of blocking the CPU with ``torch.cuda.synchronize()`` or
-    ``event.synchronize()``, this uses ``cudaLaunchHostFunc`` to enqueue a
-    host-side callback on the CUDA stream.  The callback resolves an asyncio
-    ``Future`` via ``call_soon_threadsafe``, allowing other asyncio tasks to
-    run while the GPU is busy.
-
-    Usage:
-        gpu_done = GPUFuture(loop)
-        launch_gpu_work(...)
-        gpu_done.record()
-        launch_more_gpu_work(...)
-        await gpu_done          # other tasks run while GPU is busy
-    """
-
-    # Prevent garbage-collection of live ctypes callbacks.
-    # Keyed by id() because ctypes function pointers are not hashable.
-    _prevent_gc: dict = {}
-
-    def __init__(self, loop: asyncio.AbstractEventLoop):
-        self._loop = loop
-        self._future: asyncio.Future = loop.create_future()
-
-    def record(self, stream: Optional[torch.cuda.Stream] = None) -> None:
-        """Enqueue a host callback that resolves this future.
-
-        Args:
-            stream: CUDA stream to attach to. Defaults to the current stream.
-        """
-        if stream is None:
-            stream = torch.cuda.current_stream()
-
-        # This closure prevents the ctypes wrapper from being collected
-        # while the callback is in the 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 any CUDA API.
-            try:
-                self._loop.call_soon_threadsafe(self._future.set_result, None)
-            except RuntimeError:
-                # Event loop closed; nothing to do.
-                pass
-            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:
-            # Allow the callback to be garbage-collected on failure.
-            GPUFuture._prevent_gc.pop(id(c_fn), None)
-            raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}")
-
-    def __await__(self):
-        return self._future.__await__()
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 4fda559404f..9ba3a431639 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -21,13 +21,13 @@
 )
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.contexts.static_context import StaticInferenceContext
-from megatron.core.inference.gpu_event_loop_synchronization import GPUFuture
 from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
 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,
diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index 11973551aa2..d6c6318c98c 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,86 @@ 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 GPUFuture:
+    """Awaitable that resolves when all preceding work on a CUDA stream completes.
+
+    Instead of blocking the CPU with ``torch.cuda.synchronize()`` or
+    ``event.synchronize()``, this uses ``cudaLaunchHostFunc`` to enqueue a
+    host-side callback on the CUDA stream.  The callback resolves an asyncio
+    ``Future`` via ``call_soon_threadsafe``, allowing other asyncio tasks to
+    run while the GPU is busy.
+
+    Usage:
+        gpu_done = GPUFuture(loop)
+        launch_gpu_work(...)
+        gpu_done.record()
+        launch_more_gpu_work(...)
+        await gpu_done          # other tasks run while GPU is busy
+    """
+
+    # Prevent garbage-collection of live ctypes callbacks.
+    # Keyed by id() because ctypes function pointers are not hashable.
+    _prevent_gc: dict = {}
+
+    def __init__(self, loop: asyncio.AbstractEventLoop):
+        self._loop = loop
+        self._future: asyncio.Future = loop.create_future()
+
+    def record(self, stream: Optional[torch.cuda.Stream] = None) -> None:
+        """Enqueue a host callback that resolves this future.
+
+        Args:
+            stream: CUDA stream to attach to. Defaults to the current stream.
+        """
+        if stream is None:
+            stream = torch.cuda.current_stream()
+
+        # This closure prevents the ctypes wrapper from being collected
+        # while the callback is in the 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 any CUDA API.
+            try:
+                self._loop.call_soon_threadsafe(self._future.set_result, None)
+            except RuntimeError:
+                # Event loop closed; nothing to do.
+                pass
+            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:
+            # Allow the callback to be garbage-collected on failure.
+            GPUFuture._prevent_gc.pop(id(c_fn), None)
+            raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}")
+
+    def __await__(self):
+        return self._future.__await__()

From 2ed33f26ab32b934672dbdc721b80bb1a20ad97a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 14 Apr 2026 16:02:08 -0500
Subject: [PATCH 05/10] Undo synchronize changes in engine; out of place

---
 megatron/core/inference/engines/dynamic_engine.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index c8b21e77f4e..2c9919bedfb 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -43,7 +43,7 @@
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
-from megatron.core.inference.utils import Counter, GPUFuture, await_process_call
+from megatron.core.inference.utils import Counter, await_process_call
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
 from megatron.core.transformer.enums import CudaGraphScope
@@ -1724,9 +1724,7 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
         result = await self.controller.async_generate_output_tokens_dynamic_batch(loop=self._loop)
         if will_log_this_step:
             self.step_end_event.record()
-            step_done = GPUFuture(self._loop)
-            step_done.record()
-            await step_done
+            self.step_end_event.synchronize()
             step_time = self.step_start_event.elapsed_time(self.step_end_event) / 1e3
         else:
             step_time = 0.0
@@ -2378,9 +2376,7 @@ async def run_engine_with_coordinator(
                             self.step_start_event.record()
                             self.controller.dummy_forward()
                             self.step_end_event.record()
-                            step_done = GPUFuture(self._loop)
-                            step_done.record()
-                            await step_done
+                            self.step_end_event.synchronize()
                             self.context.step_count += 1
                             self.context.prefix_cache_lru_clock += 1
                     else:

From 1587f74f80228fb395fce6e9e0fbbf9e9626243d Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Thu, 23 Apr 2026 08:32:23 -0500
Subject: [PATCH 06/10] Bring back exclusive ownership

---
 megatron/core/inference/utils.py | 107 ++++++++++++++++++++++++-------
 megatron/core/utils.py           |   4 +-
 2 files changed, 88 insertions(+), 23 deletions(-)

diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index d6c6318c98c..ad9e9d5dfe9 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -325,33 +325,91 @@ def _get_cudart() -> ctypes.CDLL:
     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 the CPU with ``torch.cuda.synchronize()`` or
-    ``event.synchronize()``, this uses ``cudaLaunchHostFunc`` to enqueue a
-    host-side callback on the CUDA stream.  The callback resolves an asyncio
-    ``Future`` via ``call_soon_threadsafe``, allowing other asyncio tasks to
-    run while the GPU is busy.
+    Instead of blocking on the CPU thread, this future fires a callback from CUDA's internal thread.
+    When the callback triggers, this class attempts to prepend the task to the event loop queue.
 
     Usage:
+
         gpu_done = GPUFuture(loop)
-        launch_gpu_work(...)
+        # Enqueue the main GPU work.
+        launch_long_gpu_work(...)
+        # Record the callback.
         gpu_done.record()
-        launch_more_gpu_work(...)
-        await gpu_done          # other tasks run while GPU is busy
+        # Recommendation: enqueue a short GPU workload to cover the latency of the callback.
+        launch_short_gpu_work(...)
+        await gpu_done    # resumes at front of queue as soon as GPU finishes.
     """
 
-    # Prevent garbage-collection of live ctypes callbacks.
-    # Keyed by id() because ctypes function pointers are not hashable.
     _prevent_gc: dict = {}
+    _asyncio_future_blocking = False
 
     def __init__(self, loop: asyncio.AbstractEventLoop):
         self._loop = loop
-        self._future: asyncio.Future = loop.create_future()
+        self._done = False
+        self._callbacks: list = []
+
+    def get_loop(self):
+        return self._loop
+
+    def done(self):
+        return self._done
+
+    def cancelled(self):
+        return False
+
+    def cancel(self, msg=None):
+        return False
+
+    def result(self):
+        if not self._done:
+            raise asyncio.InvalidStateError('not ready')
+        return None
+
+    def add_done_callback(self, fn, *, context=None):
+        if self._done:
+            try:
+                self._loop.call_soon_front(fn, self, context=context)
+            except AttributeError:
+                self._loop.call_soon(fn, self, 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.
+        """Enqueue a host callback that resolves this future when stream drains.
 
         Args:
             stream: CUDA stream to attach to. Defaults to the current stream.
@@ -359,18 +417,17 @@ def record(self, stream: Optional[torch.cuda.Stream] = None) -> None:
         if stream is None:
             stream = torch.cuda.current_stream()
 
-        # This closure prevents the ctypes wrapper from being collected
-        # while the callback is in the 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 any CUDA API.
+            # Runs on CUDA's internal callback thread; MUST NOT call CUDA API.
             try:
-                self._loop.call_soon_threadsafe(self._future.set_result, None)
+                try:
+                    self._loop.call_soon_threadsafe_front(self._resolve)
+                except AttributeError:
+                    self._loop.call_soon_threadsafe(self._resolve)
             except RuntimeError:
-                # Event loop closed; nothing to do.
-                pass
+                pass  # event loop closed
             GPUFuture._prevent_gc.pop(id(prevent_gc_ref), None)
 
         c_fn = _CUDA_HOST_FN_T(_host_fn)
@@ -381,9 +438,15 @@ def _host_fn(_user_data: ctypes.c_void_p) -> None:
             ctypes.c_void_p(stream.cuda_stream), c_fn, ctypes.c_void_p(0)
         )
         if err != 0:
-            # Allow the callback to be garbage-collected on failure.
             GPUFuture._prevent_gc.pop(id(c_fn), None)
             raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}")
 
-    def __await__(self):
-        return self._future.__await__()
+    def _resolve(self):
+        """Mark done and fire callbacks at the front of the ready queue."""
+        self._done = True
+        cbs, self._callbacks = self._callbacks, []
+        for fn, ctx in cbs:
+            try:
+                self._loop.call_soon_front(fn, self, context=ctx)
+            except AttributeError:
+                self._loop.call_soon(fn, self, 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
 

From 696a7d2aa26f0b0032e6e6eba29e3ba1db6a32fe Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 5 May 2026 05:43:12 -0500
Subject: [PATCH 07/10] Move await past sampling enqueue

---
 .../text_generation_controller.py             | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

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 9ba3a431639..9d952863ed1 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -1755,14 +1755,12 @@ async def async_generate_output_tokens_dynamic_batch(
             context.kv_block_allocator.store_routing_per_block(self._router_record_bookkeeping())
             range_pop()
 
-        # Record after forward pass kernels are enqueued.  Awaiting this
-        # lets other asyncio tasks run while the GPU is busy, and resumes
-        # as soon as the forward pass completes.
+        # 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()
-        await gpu_done
 
         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:
@@ -1770,6 +1768,14 @@ async def async_generate_output_tokens_dynamic_batch(
                 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()
@@ -1789,8 +1795,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
@@ -1809,7 +1813,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

From e26743dd1ba666190a3bb5007c50f0104f2526a1 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 5 May 2026 05:52:55 -0500
Subject: [PATCH 08/10] Make faux-interrupt optional

---
 .../text_generation_controller.py             |  2 +-
 megatron/core/inference/utils.py              | 41 +++++++++++--------
 2 files changed, 26 insertions(+), 17 deletions(-)

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 9d952863ed1..aba68ed97e5 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -1721,7 +1721,7 @@ async def async_generate_output_tokens_dynamic_batch(
             return None
 
         loop = get_asyncio_loop(loop)
-        gpu_done = GPUFuture(loop)
+        gpu_done = GPUFuture(loop, interrupt_event_loop=True)
 
         with torch.inference_mode():
             input_ids, position_ids = self._dynamic_step_context_init()
diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index ad9e9d5dfe9..f2aa3aa9783 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -352,7 +352,10 @@ 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 the callback triggers, this class attempts to prepend the task to the event loop queue.
+
+    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:
 
@@ -363,14 +366,15 @@ class GPUFuture:
         gpu_done.record()
         # Recommendation: enqueue a short GPU workload to cover the latency of the callback.
         launch_short_gpu_work(...)
-        await gpu_done    # resumes at front of queue as soon as GPU finishes.
+        await gpu_done    # resumes as soon as GPU finishes.
     """
 
     _prevent_gc: dict = {}
     _asyncio_future_blocking = False
 
-    def __init__(self, loop: asyncio.AbstractEventLoop):
+    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 = []
 
@@ -391,12 +395,23 @@ def result(self):
             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):
         if self._done:
-            try:
-                self._loop.call_soon_front(fn, self, context=context)
-            except AttributeError:
-                self._loop.call_soon(fn, self, context=context)
+            self._schedule(fn, context=context)
         else:
             self._callbacks.append((fn, context))
 
@@ -422,10 +437,7 @@ def record(self, stream: Optional[torch.cuda.Stream] = None) -> None:
         def _host_fn(_user_data: ctypes.c_void_p) -> None:
             # Runs on CUDA's internal callback thread; MUST NOT call CUDA API.
             try:
-                try:
-                    self._loop.call_soon_threadsafe_front(self._resolve)
-                except AttributeError:
-                    self._loop.call_soon_threadsafe(self._resolve)
+                self._schedule_threadsafe(self._resolve)
             except RuntimeError:
                 pass  # event loop closed
             GPUFuture._prevent_gc.pop(id(prevent_gc_ref), None)
@@ -442,11 +454,8 @@ def _host_fn(_user_data: ctypes.c_void_p) -> None:
             raise RuntimeError(f"cudaLaunchHostFunc failed with CUDA error {err}")
 
     def _resolve(self):
-        """Mark done and fire callbacks at the front of the ready queue."""
+        """Mark done and fire pending callbacks."""
         self._done = True
         cbs, self._callbacks = self._callbacks, []
         for fn, ctx in cbs:
-            try:
-                self._loop.call_soon_front(fn, self, context=ctx)
-            except AttributeError:
-                self._loop.call_soon(fn, self, context=ctx)
+            self._schedule(fn, context=ctx)

From 911cab4bb183602daef12fb5ea42215f3f723a9c Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 5 May 2026 06:02:26 -0500
Subject: [PATCH 09/10] Move `_dynamic_step_log_probs_bookkeeping`

---
 .../text_generation_controllers/text_generation_controller.py  | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

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 aba68ed97e5..add71b7ec8c 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -1725,6 +1725,7 @@ async def async_generate_output_tokens_dynamic_batch(
 
         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
@@ -1761,8 +1762,6 @@ async def async_generate_output_tokens_dynamic_batch(
         gpu_done.record()
 
         with torch.inference_mode():
-            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")

From 8d3f778eb3073f2b177e0bd11ad72b7f09eedd7a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 5 May 2026 06:05:52 -0500
Subject: [PATCH 10/10] lint

---
 megatron/core/inference/utils.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index f2aa3aa9783..33808893a62 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -379,18 +379,23 @@ def __init__(self, loop: asyncio.AbstractEventLoop, interrupt_event_loop: bool =
         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
@@ -410,6 +415,7 @@ def _schedule_threadsafe(self, fn):
             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: