From ef4b76074c60f935f3e3b7d045f6f1920c62e6ce Mon Sep 17 00:00:00 2001
From: shanmugamr1992 <shanmugamr1992@gmail.com>
Date: Mon, 5 Jan 2026 16:14:07 -0800
Subject: [PATCH 1/3] Supporting infernece when called within an asyncio loop

---
 .../core/inference/engines/dynamic_engine.py  | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index c7698b8a4bb..e59c8961c43 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import asyncio
+import concurrent.futures
 import logging
 import multiprocessing
 import os
@@ -1331,11 +1332,29 @@ async def async_step(
         # Keep for compatibility with current test suite.
         return ret
 
+    def _run_coroutine_sync(self, coro):
+        """Run a coroutine synchronously, handling the case when already in an event loop.
+
+        This method safely runs an async coroutine from synchronous code, even when
+        called from within an already running event loop (e.g., when used with async
+        frameworks like pytriton).
+        """
+        try:
+            # Check if there's already a running event loop
+            asyncio.get_running_loop()
+            # We're inside a running loop - run in a separate thread
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(asyncio.run, coro)
+                return future.result()
+        except RuntimeError:
+            # No running loop - safe to use run_until_complete
+            return self._loop.run_until_complete(coro)
+
     def step_modern(
         self,
     ) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest], float]:
         """Synchronous wrapper for `self.async_step`."""
-        return self._loop.run_until_complete(self.async_step())
+        return self._run_coroutine_sync(self.async_step())
 
     def step_legacy(
         self, sampling_params: SamplingParams
@@ -1346,7 +1365,7 @@ def step_legacy(
             "0.16. Please use `step_modern()` going forward, which will eventually "
             "be renamed to `step()`."
         )
-        result = self._loop.run_until_complete(self.async_step())
+        result = self._run_coroutine_sync(self.async_step())
         active_requests = [self.get_request(i) for i in result["active_request_ids"]]
         finished_requests = [r.merge() for r in result["finished_request_records"]]
         return active_requests, finished_requests, result["step_time"]

From 2f31f1cead3d23e6ad4146f3e3b15ad9387dc975 Mon Sep 17 00:00:00 2001
From: shanmugamr1992 <shanmugamr1992@gmail.com>
Date: Tue, 20 Jan 2026 13:48:01 -0800
Subject: [PATCH 2/3] Onurs comments

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

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index d1529498c86..7b5fbcd4cba 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1575,6 +1575,12 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int]
         metadata = req.tracked_metadata
         metadata_types = req.get_metadata_types()
         for m, m_type in zip(metadata, metadata_types):
+            if not isinstance(m, torch.Tensor):
+                m = torch.as_tensor(
+                    m,
+                    device=self.request_metadata[label].device,
+                    dtype=self.request_metadata[label].dtype,
+                )
             label, _, _ = m_type
             self.request_metadata[label][current_id] = m
 

From 83de56ee518948b9533181299ba611472051b9fe Mon Sep 17 00:00:00 2001
From: shanmugamr1992 <shanmugamr1992@gmail.com>
Date: Tue, 20 Jan 2026 13:59:30 -0800
Subject: [PATCH 3/3] Onurs comments

---
 megatron/core/inference/contexts/dynamic_context.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 20cc1c369f3..78f412ea5ab 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -1608,13 +1608,14 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int]
         metadata = req.tracked_metadata
         metadata_types = req.get_metadata_types()
         for m, m_type in zip(metadata, metadata_types):
+            label, _, _ = m_type
             if not isinstance(m, torch.Tensor):
                 m = torch.as_tensor(
                     m,
                     device=self.request_metadata[label].device,
                     dtype=self.request_metadata[label].dtype,
                 )
-            label, _, _ = m_type
+
             self.request_metadata[label][current_id] = m
 
         # Handle length and block assignments.