[Shortfin][LLM] Add initial support for disaggregated invocations

Author: vinayakdsci

shortfin/python/shortfin_apps/llm/cli.py

@@ -134,13 +134,18 @@ def add_service_args(parser: argparse.ArgumentParser):
     parser.add_argument(
         "--benchmark",
         action="store_true",
-        help="Perform a benchmarking run for throughput",
+        help="Perform a benchmarking run for throughput.",
     )
     parser.add_argument(
         "--benchmark_tasks",
         type=int,
         default=None,
-        help="Workload size to benchmark with",
+        help="Workload size to benchmark with.",
+    )
+    parser.add_argument(
+        "--disaggregate",
+        action="store_true",
+        help="Disaggregate the prefill and decode invocations to separate HIP streams.",
     )
 
 

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

@@ -53,6 +53,7 @@ def __init__(
         functions: dict[int, sf.ProgramFunction],
         ideal_batch_size: int,
         program_isolation: str,
+        exec_fiber: Fiber,
     ):
         super().__init__(fiber=fiber)
         self.name = name
@@ -65,7 +66,7 @@ def __init__(
         self.ideal_batch_size: int = ideal_batch_size
         self.page_seq_stride = self.model_params.paged_kv_cache.block_seq_stride
         self.scheduler = Scheduler(ideal_batch_size=self.ideal_batch_size)
-
+        self.exec_fiber = exec_fiber
         self.program_isolation = program_isolation
 
     def handle_inference_request(self, request):
@@ -161,6 +162,7 @@ def __init__(
         model_params: ModelParams,
         prefill_functions: dict[int, sf.ProgramFunction],
         program_isolation: str,
+        exec_fiber: Fiber,
     ):
         super().__init__(
             name="prefill",
@@ -170,11 +172,12 @@ def __init__(
             functions=prefill_functions,
             ideal_batch_size=max(model_params.prefill_batch_sizes),
             program_isolation=program_isolation,
+            exec_fiber=exec_fiber,
         )
 
     def make_process(self, cache: BasePagedAttentionCache, fiber: Fiber):
         return PrefillExecutorProcess(
-            fiber,
+            self.exec_fiber,
             self.functions,
             self.page_seq_stride,
             cache.page_pool.page_tables,
@@ -216,6 +219,7 @@ def __init__(
         model_params: ModelParams,
         decode_functions: dict[int, sf.ProgramFunction],
         program_isolation: str,
+        exec_fiber: Fiber,
     ):
         super().__init__(
             name="decode",
@@ -225,11 +229,12 @@ def __init__(
             functions=decode_functions,
             ideal_batch_size=max(model_params.decode_batch_sizes),
             program_isolation=program_isolation,
+            exec_fiber=exec_fiber,
         )
 
     def make_process(self, cache: BasePagedAttentionCache, fiber: Fiber):
         return DecodeExecutorProcess(
-            fiber,
+            self.exec_fiber,
             self.functions,
             self.page_seq_stride,
             cache.page_pool.page_tables,

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

@@ -238,6 +238,7 @@ class ServerParams:
     amdgpu_async_caching: bool = False
     amdgpu_allocators: Optional[str] = None
     amdgpu_allow_device_reuse: bool = False
+    disaggregate: bool = False
 
     @staticmethod
     def load(config_path: Optional[Path] = None) -> "ServerParams":

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

@@ -52,12 +52,16 @@ async def get(self) -> tuple[int, sf.Fiber]:
         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)
+                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]
@@ -69,11 +73,15 @@ def pool(self) -> list[sf.Fiber]:
         return self.__fiber_pool
 
     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)
+            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

@@ -25,6 +25,7 @@ def lifecycle(app: FastApi):
 
 from contextlib import asynccontextmanager
 import logging
+import os
 
 
 def get_eos_from_tokenizer_config(json_path):
@@ -63,6 +64,11 @@ def __init__(self, args):
             )
             server_params.decode_config = decode_config
 
+        if args.disaggregate:
+            # 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"
+
         # Setup system (configure devices, etc).
         sysman = LlmSystemManager(
             device=args.device,

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

@@ -33,7 +33,7 @@
 class LlmGenerateService(GenerateService):
     """Top level service interface for generating text against a model."""
 
-    inference_program: sf.Program
+    inference_program: list[sf.Program]
     prefill_functions: dict[int, sf.ProgramFunction]
     decode_functions: dict[int, sf.ProgramFunction]
 
@@ -53,6 +53,7 @@ def __init__(
         self.tokenizer = tokenizer
         self.model_params = model_params
         self.server_params = server_params
+        self.disaggregate = server_params.disaggregate
         self.max_queue_size = max_queue_size
         self.current_queue_size = 0
         self.main_fiber_pool = FiberPool(
@@ -92,23 +93,30 @@ 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)
+        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)
+        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)
+        self.decode_fiber = self.sysman.ls.create_fiber(
+            self.decode_worker, devices=[devices[1 % len(devices)]]
+        )
 
         self.devices = self.prefill_fiber.devices_dict.values()
 
@@ -141,17 +149,36 @@ def _initialize_page_cache(self):
 
     def start(self):
         component_modules = self.initialize_program_modules("main")
-        self.inference_program = self.create_program(
-            modules=component_modules, devices=self.sysman.ls.devices
-        )
+        print(f"{self.disaggregate=}")
+        self.inference_program = [
+            self.create_program(
+                modules=component_modules, devices=[self.sysman.ls.devices[idx]]
+            )
+            for idx in range(len(self.sysman.ls.devices))
+        ]
         self.initialize_function_references()
 
+        task_list = [
+            "prefill-exec",
+            "decode-exec",
+        ]
+
+        devices = self.sysman.ls.devices
+        workers = [self.sysman.ls.create_worker(f"{task}-worker") for task in task_list]
+        fibers = [
+            self.sysman.ls.create_fiber(
+                workers[idx], devices=[devices[idx % len(devices)]]
+            )
+            for idx in range(len(workers))
+        ]
+
         self.prefill_batcher = PrefillBatcherProcess(
             self.prefill_fiber,
             self.page_cache,
             self.model_params,
             self.prefill_functions,
             self.prog_isolation,
+            fibers[0],
         )
 
         self.decode_batcher = DecodeBatcherProcess(
@@ -160,21 +187,24 @@ def start(self):
             self.model_params,
             self.decode_functions,
             self.prog_isolation,
+            fibers[1],
         )
 
         self.prefill_batcher.launch()
         self.decode_batcher.launch()
 
     def initialize_function_references(self):
+        devices = self.sysman.ls.devices
+        num_devices = len(devices)
         self.prefill_functions = {}
         for bs in self.model_params.prefill_batch_sizes:
-            self.prefill_functions[bs] = self.inference_program[
+            self.prefill_functions[bs] = self.inference_program[0][
                 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[
+            self.decode_functions[bs] = self.inference_program[1 % num_devices][
                 f"{self.model_params.module_name}.decode_bs{bs}"
             ]