Assert single physical GPU; Cleanup

vinayakdsci · vinayakdsci · commit ad3b24ae766d · 2025-06-05T05:36:33.000Z
diff --git a/shortfin/python/shortfin_apps/llm/components/batcher.py b/shortfin/python/shortfin_apps/llm/components/batcher.py
@@ -54,7 +54,7 @@ def __init__(
         functions: dict[int, sf.ProgramFunction],
         ideal_batch_size: int,
         program_isolation: str,
-        exec_fiber: Fiber,
+        exec_fiber: Fiber | None = None,
     ):
         super().__init__(fiber=fiber)
         self.name = name
@@ -121,7 +121,11 @@ async def board_flights(self):
         scheduled = []
         for job in to_schedule:
             scheduled = scheduled + job
-            self.board(cache, self.exec_fiber, job)
+            self.board(
+                cache,
+                self.exec_fiber if self.exec_fiber is not None else self.fiber,
+                job,
+            )
             logger.debug("Post boarding cache state: %r", cache)
 
         pending = set(pending) - set(scheduled)
@@ -169,7 +173,7 @@ def __init__(
         model_params: ModelParams,
         prefill_functions: dict[int, sf.ProgramFunction],
         program_isolation: str,
-        exec_fiber: Fiber,
+        exec_fiber: Fiber | None = None,
     ):
         super().__init__(
             name="prefill",
@@ -227,7 +231,7 @@ def __init__(
         model_params: ModelParams,
         decode_functions: dict[int, sf.ProgramFunction],
         program_isolation: str,
-        exec_fiber: Fiber,
+        exec_fiber: Fiber | None = None,
     ):
         super().__init__(
             name="decode",
diff --git a/shortfin/python/shortfin_apps/llm/components/fiber_pool.py b/shortfin/python/shortfin_apps/llm/components/fiber_pool.py
@@ -9,6 +9,8 @@
 from .manager import LlmSystemManager
 import asyncio
 
+from typing import override
+
 
 class FiberPool:
     """
@@ -32,65 +34,105 @@ def __init__(
         self.sysman: LlmSystemManager = sysman
         self.name: str = name
 
-        # Name mangle to make outside access harder.
-        self.__fiber_pool: list[sf.Fiber] = []
-        self.__workers: list[sf.Worker] = []
+        self._fiber_pool: list[sf.Fiber] = []
+        self._workers: list[sf.Worker] = []
         # Keep track of how many extra fibers were created
         # during runtime if `resizable` is set to True.
-        self.__extra_fibers: int = 0
-        self.__index_queue = asyncio.Queue()
+        self._extra_fibers: int = 0
+        self._index_queue = asyncio.Queue()
+
+        self._initialize_pool()
 
-        self.__initialize_pool()
+    def resize(self):
+        new_worker = self.sysman.ls.create_worker(
+            f"{self.name}-new-worker-{self._extra_fibers}"
+        )
+        self._workers.append(new_worker)
+        fiber = self.sysman.ls.create_fiber(new_worker)
+        self._fiber_pool.append(fiber)
+        self._extra_fibers += 1
+
+        return [self.size() - 1, fiber]
 
     async def get(self) -> tuple[int, sf.Fiber]:
         try:
-            idx = self.__index_queue.get_nowait()
+            idx = self._index_queue.get_nowait()
             return (
                 idx,
-                self.__fiber_pool[idx],
+                self._fiber_pool[idx],
             )
         except asyncio.QueueEmpty:
             if self.resizable:
                 # Resize the fiber pool by adding a new fiber.
-                devices = self.sysman.ls.devices
-                num_devices = len(devices)
-                new_worker = self.sysman.ls.create_worker(
-                    f"{self.name}-new-worker-{self.__extra_fibers}"
-                )
-                self.__workers.append(new_worker)
-
-                fiber = self.sysman.ls.create_fiber(
-                    new_worker, devices=[devices[self.size() % num_devices]]
-                )
-                self.__fiber_pool.append(fiber)
-                self.__extra_fibers += 1
-                return [self.size() - 1, fiber]
-
-            available_index = await self.__index_queue.get()
-            return (available_index, self.__fiber_pool[available_index])
+                return self.resize()
+
+            available_index = await self._index_queue.get()
+            return (available_index, self._fiber_pool[available_index])
 
     def pool(self) -> list[sf.Fiber]:
-        return self.__fiber_pool
+        return self._fiber_pool
 
-    def __initialize_pool(self):
-        devices = self.sysman.ls.devices
-        num_devices = len(devices)
+    def _initialize_pool(self):
         for idx in range(self.init_size):
             worker = self.sysman.ls.create_worker(f"{self.name}-init-worker-{idx}")
-            self.__workers.append(worker)
-
-            fiber = self.sysman.ls.create_fiber(
-                worker, devices=[devices[idx % num_devices]]
-            )
-            self.__fiber_pool.append(fiber)
+            self._workers.append(worker)
+            fiber = self.sysman.ls.create_fiber(worker)
+            self._fiber_pool.append(fiber)
             assert idx < self.size()
-            self.__index_queue.put_nowait(idx)
+            self._index_queue.put_nowait(idx)
 
     def return_fiber(self, indices: int | list[int]):
         if not isinstance(indices, list):
             indices = [indices]
         for idx in indices:
-            self.__index_queue.put_nowait(idx)
+            self._index_queue.put_nowait(idx)
 
     def size(self) -> int:
-        return len(self.__fiber_pool)
+        return len(self._fiber_pool)
+
+
+class DisaggregatedFiberPool(FiberPool):
+    def __init__(
+        self,
+        sysman: LlmSystemManager,
+        init_size: int,
+        resizable: bool = True,
+        name: str = "default-disagg-fiber-pool",
+    ):
+        super().__init__(
+            sysman=sysman,
+            init_size=init_size,
+            resizable=resizable,
+            name=name,
+        )
+
+    @override
+    def resize(self):
+        devices = self.sysman.ls.devices
+        num_devices = len(devices)
+        new_worker = self.sysman.ls.create_worker(
+            f"{self.name}-new-worker-{self._extra_fibers}"
+        )
+        self._workers.append(new_worker)
+
+        fiber = self.sysman.ls.create_fiber(
+            new_worker, devices=[devices[self.size() % num_devices]]
+        )
+        self._fiber_pool.append(fiber)
+        self._extra_fibers += 1
+        return [self.size() - 1, fiber]
+
+    @override
+    def _initialize_pool(self):
+        devices = self.sysman.ls.devices
+        num_devices = len(devices)
+        for idx in range(self.init_size):
+            worker = self.sysman.ls.create_worker(f"{self.name}-init-worker-{idx}")
+            self._workers.append(worker)
+
+            fiber = self.sysman.ls.create_fiber(
+                worker, devices=[devices[idx % num_devices]]
+            )
+            self._fiber_pool.append(fiber)
+            assert idx < self.size()
+            self._index_queue.put_nowait(idx)
diff --git a/shortfin/python/shortfin_apps/llm/components/lifecycle.py b/shortfin/python/shortfin_apps/llm/components/lifecycle.py
@@ -17,7 +17,7 @@ def lifecycle(app: FastApi):
 from .config_struct import ModelParams, ServerParams
 from .token_selection_strategy import DecodeConfig
 from .manager import LlmSystemManager
-from .service import LlmGenerateService
+from .service import LlmGenerateService, LlmGenerateDisaggregatedService
 from .tokenizer import Tokenizer
 from typing import TYPE_CHECKING
 from fastapi import FastAPI
@@ -64,10 +64,18 @@ def __init__(self, args):
             )
             server_params.decode_config = decode_config
 
+        service_cls = LlmGenerateService
         if args.disaggregate:
+            # To not run into complications with sharded models, assert that the server is
+            # being run only on one physical device.
+            rocr_visible_devices = os.environ.get("ROCR_VISIBLE_DEVICES")
+            assert (
+                rocr_visible_devices is not None and len(rocr_visible_devices) <= 2
+            ), "Running disaggregated prefill on HIP streams is supported only when running on one physical device. Set `ROCR_VISIBLE_DEVICES`=<device_id>."
             # Setup two logical devices on one physical device to disaggregate
             # prefill and decode invocations to distinct streams.
             os.environ["SHORTFIN_AMDGPU_LOGICAL_DEVICES_PER_PHYSICAL_DEVICE"] = "2"
+            service_cls = LlmGenerateDisaggregatedService
 
         # Setup system (configure devices, etc).
         sysman = LlmSystemManager(
@@ -84,7 +92,7 @@ def __init__(self, args):
         tokenizer = Tokenizer.from_tokenizer_json_file(
             args.tokenizer_json, eos_token=eos_token
         )
-        service = LlmGenerateService(
+        service = service_cls(
             name="default",
             sysman=sysman,
             tokenizer=tokenizer,
diff --git a/shortfin/python/shortfin_apps/llm/components/service.py b/shortfin/python/shortfin_apps/llm/components/service.py
