Syntactic sugar for CG and awaiting

Author: tdene

megatron/core/inference/utils.py

@@ -236,6 +236,97 @@ def tensor_swap(x, src_idxs, dst_idxs):
     x[dst_idxs], x[src_idxs] = x[src_idxs], x[dst_idxs]
 
 
+class CUDAGraphCache:
+    """Syntactic sugar for capturing and replaying graphs over in-line code blocks.
+
+    In Python, neither context managers nor function calls allow the flexibility of this sugar.
+    This class allows for the following type of usage:
+
+    ```
+        cache_for_my_op = CUDAGraphCache()
+        token_count = some_token_count_value
+        request_count = some_request_count_value
+
+        for _ in cache_for_my_op(token_count, request_count):
+            logits = extract_logits(n).float()
+            probs = torch.softmax(logits, dim=-1)
+            output.copy_(sample(probs))
+    ```
+
+    The above example handles both capture and replay.
+    During capture, the for-loop yields twice; during replay, the for-loop body is skipped.
+    """
+
+    def __init__(self):
+        self._graphs: dict = {}
+
+    def __call__(self, *keys, pool=None, skip_warmup=False, eager=False):
+        """Default use: captures if the keys are new, replays if the keys are cached.
+
+        Args:
+            *keys: Arguments to key the graph on.
+            pool (Optional): CUDA memory pool to build the graphs within.
+            skip_warmup (Optional): If True, skip the warmup run before capture.
+            eager (Optional): If True, run the body once without capturing or replaying.
+        """
+        return self.capture(
+            *keys, pool=pool, replay_on_hit=True, skip_warmup=skip_warmup, eager=eager
+        )
+
+    def __contains__(self, key) -> bool:
+        return key in self._graphs
+
+    def __delitem__(self, key) -> None:
+        del self._graphs[key]
+
+    def clear(self) -> None:
+        """Remove all cached graphs."""
+        self._graphs.clear()
+
+    def capture(self, *keys, pool=None, replay_on_hit=False, skip_warmup=False, eager=False):
+        """Inline capture: captures if the keys are new, does nothing if the keys are cached.
+
+        Args:
+            *keys: Cache key components (stored as a tuple).
+            pool (Optional): CUDA memory pool to build the graphs within.
+            replay_on_hit (Optional): If True, replay the graph on cache hit.
+            skip_warmup (Optional): If True, skip the warmup run before capture.
+            eager (Optional): If True, run the body once without capturing or replaying.
+        """
+        if eager:
+            yield
+            return
+        key = keys if len(keys) != 1 else keys[0]
+        if key in self._graphs:
+            if replay_on_hit:
+                self._graphs[key].replay()
+            return
+        if not skip_warmup:
+            # Suggested best practice: run the body once before recording to ensure
+            # the kernels are JIT-compiled and the allocator state is stable.
+            yield
+        # Now perform the actual graph capture.
+        g = torch.cuda.CUDAGraph()
+        kwargs = {"pool": pool} if pool is not None else {}
+        with torch.cuda.graph(g, **kwargs):
+            yield
+        self._graphs[key] = g
+
+    def replay(self, *keys) -> None:
+        """Replay a previously captured graph."""
+        key = keys if len(keys) != 1 else keys[0]
+        self._graphs[key].replay()
+
+
+async def torch_awaitable(stream: torch.cuda.Stream | None = None):
+    """Syntactic sugar for returning an awaitable handle for non-distributed torch."""
+    if stream is None:
+        stream = torch.cuda.current_stream()
+    event = stream.record_event()
+    while not event.query():
+        await asyncio.sleep(0)
+
+
 async def await_process_call(call, process: multiprocessing.Process, timeout: float = 1.0):
     """Repeatedly wait for a multiprocessing callable to resolve, aborting upon process failure.