Assert single physical GPU; Cleanup

Author: vinayakdsci

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",

shortfin/python/shortfin_apps/llm/components/fiber_pool.py

@@ -32,65 +32,103 @@ 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,
+        )
+
+    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]
+
+    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)

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,

shortfin/python/shortfin_apps/llm/components/service.py

@@ -25,7 +25,7 @@
 from .token_selection_strategy import is_multi_response
 
 from ...utils import GenerateService
-from .fiber_pool import FiberPool
+from .fiber_pool import FiberPool, DisaggregatedFiberPool
 
 logger = logging.getLogger(__name__)
 
@@ -56,7 +56,8 @@ def __init__(
         self.disaggregate = server_params.disaggregate
         self.max_queue_size = max_queue_size
         self.current_queue_size = 0
-        self.main_fiber_pool = FiberPool(
+        fiber_pool_cls = DisaggregatedFiberPool if self.disaggregate else FiberPool
+        self.main_fiber_pool = fiber_pool_cls(
             self.sysman, self.max_queue_size, resizable=True
         )
 
@@ -93,30 +94,23 @@ def remove_from_queue(self, num_beams: int):
     def _initialize_worker_and_fiber(self):
         num_workers = self.server_params.workers
         fibers_per_worker = self.server_params.fibers_per_worker
-        devices = self.sysman.ls.devices
 
         logger.info(
             f"Creating {num_workers} workers, with {fibers_per_worker} fibers per worker..."
         )
 
         self.main_worker = self.sysman.ls.create_worker(f"{self.name}-inference-main-0")
-        self.main_fiber = self.sysman.ls.create_fiber(
-            self.main_worker, devices=[devices[0]]
-        )
+        self.main_fiber = self.sysman.ls.create_fiber(self.main_worker)
 
         self.prefill_worker = self.sysman.ls.create_worker(
             f"{self.name}-inference-prefill-0"
         )
-        self.prefill_fiber = self.sysman.ls.create_fiber(
-            self.prefill_worker, devices=[devices[0]]
-        )
+        self.prefill_fiber = self.sysman.ls.create_fiber(self.prefill_worker)
 
         self.decode_worker = self.sysman.ls.create_worker(
             f"{self.name}-inference-decode-0"
         )
-        self.decode_fiber = self.sysman.ls.create_fiber(
-            self.decode_worker, devices=[devices[1 % len(devices)]]
-        )
+        self.decode_fiber = self.sysman.ls.create_fiber(self.decode_worker)
 
         self.devices = self.prefill_fiber.devices_dict.values()
 
@@ -147,6 +141,108 @@ def _initialize_page_cache(self):
                 f"Unknown prefix_sharing_algorithm {self.server_params.prefix_sharing_algorithm}. Currently only supporting 'trie' and 'none'."
             )
 
+    def start(self):
+        component_modules = self.initialize_program_modules("main")
+        self.inference_program = self.create_program(
+            modules=component_modules, devices=self.sysman.ls.devices
+        )
+        self.initialize_function_references()
+
+        self.prefill_batcher = PrefillBatcherProcess(
+            self.prefill_fiber,
+            self.page_cache,
+            self.model_params,
+            self.prefill_functions,
+            self.prog_isolation,
+        )
+
+        self.decode_batcher = DecodeBatcherProcess(
+            self.decode_fiber,
+            self.page_cache,
+            self.model_params,
+            self.decode_functions,
+            self.prog_isolation,
+        )
+
+        self.prefill_batcher.launch()
+        self.decode_batcher.launch()
+
+    def initialize_function_references(self):
+        self.prefill_functions = {}
+        for bs in self.model_params.prefill_batch_sizes:
+            self.prefill_functions[bs] = self.inference_program[
+                f"{self.model_params.module_name}.prefill_bs{bs}"
+            ]
+        # Resolve decode entrypoints.
+        self.decode_functions = {}
+        for bs in self.model_params.decode_batch_sizes:
+            self.decode_functions[bs] = self.inference_program[
+                f"{self.model_params.module_name}.decode_bs{bs}"
+            ]
+
+    def __repr__(self):
+        return (
+            f"ServiceManager(\n"
+            f"  model_params={self.model_params}\n"
+            f"  server_params={self.server_params}\n"
+            f"  inference_modules={self.inference_modules}\n"
+            f"  page_cache={self.page_cache}\n"
+            f")"
+        )
+
+
+class LlmGenerateDisaggregatedService(LlmGenerateService):
+    def __init__(
+        self,
+        *,
+        name: str,
+        sysman: LlmSystemManager,
+        tokenizer: Tokenizer,
+        model_params: ModelParams,
+        server_params: "ServerParams",
+        program_isolation: str = "per_call",
+        max_queue_size: int = 3,  # Maximum number of requests in queue
+    ):
+        super().__init__(
+            name=name,
+            sysman=sysman,
+            tokenizer=tokenizer,
+            model_params=model_params,
+            server_params=server_params,
+            program_isolation=program_isolation,
+            max_queue_size=max_queue_size,
+        )
+
+    def _initialize_worker_and_fiber(self):
+        num_workers = self.server_params.workers
+        fibers_per_worker = self.server_params.fibers_per_worker
+        devices = self.sysman.ls.devices
+
+        logger.info(
+            f"Creating {num_workers} workers, with {fibers_per_worker} fibers per worker..."
+        )
+
+        self.main_worker = self.sysman.ls.create_worker(f"{self.name}-inference-main-0")
+        self.main_fiber = self.sysman.ls.create_fiber(
+            self.main_worker, devices=[devices[0]]
+        )
+
+        self.prefill_worker = self.sysman.ls.create_worker(
+            f"{self.name}-inference-prefill-0"
+        )
+        self.prefill_fiber = self.sysman.ls.create_fiber(
+            self.prefill_worker, devices=[devices[0]]
+        )
+
+        self.decode_worker = self.sysman.ls.create_worker(
+            f"{self.name}-inference-decode-0"
+        )
+        self.decode_fiber = self.sysman.ls.create_fiber(
+            self.decode_worker, devices=[devices[1 % len(devices)]]
+        )
+
+        self.devices = self.prefill_fiber.devices_dict.values()
+
     def start(self):
         component_modules = self.initialize_program_modules("main")
         print(f"{self.disaggregate=}")
@@ -210,7 +306,7 @@ def initialize_function_references(self):
 
     def __repr__(self):
         return (
-            f"ServiceManager(\n"
+            f"DisaggregatedServiceManager(\n"
             f"  model_params={self.model_params}\n"
             f"  server_params={self.server_params}\n"
             f"  inference_modules={self.inference_modules}\n"