Separate prefill and decode fiber pools

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

@@ -142,7 +142,7 @@ def __init__(
         responder: FastAPIResponder,
         fiber: sf.Fiber | None = None,
     ):
-        super().__init__(fiber=service.fiber_pool.fibers[0] if fiber is None else fiber)
+        super().__init__(fiber=service.prefill_fiber_pool.fibers[0] if fiber is None else fiber)
         self.service = service
         self.gen_req = gen_req
         self.responder = responder

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

@@ -85,18 +85,26 @@ def initialize_worker_and_fiber(self):
         logger.info(
             f"Creating {num_workers} workers, with {fibers_per_worker} fibers per worker..."
         )
-        fibers = []
+        prefill_fibers = []
+        decode_fibers = []
         for i in range(num_workers):
-            worker = self.sysman.ls.create_worker(f"{self.name}-inference-{i}")
+            prefill_worker = self.sysman.ls.create_worker(f"{self.name}-prefill_inference-{i}")
+            decode_worker = self.sysman.ls.create_worker(f"{self.name}-decode_inference-{i}")
             for _ in range(fibers_per_worker):
-                fiber = self.sysman.ls.create_fiber(worker)
-                fibers.append(fiber)
-
-        self.fiber_pool = FiberPool(
-            fibers,
-            fibers,
+                prefill_fiber = self.sysman.ls.create_fiber(prefill_worker)
+                prefill_fibers.append(prefill_fiber)
+                decode_fiber = self.sysman.ls.create_fiber(decode_worker)
+                decode_fibers.append(decode_fiber)
+
+        self.prefill_fiber_pool = FiberPool(
+            prefill_fibers,
+            prefill_fibers,
+        )
+        self.decode_fiber_pool = FiberPool(
+            decode_fibers,
+            decode_fibers,
         )
-        self.devices = fibers[0].devices_dict.values()
+        self.devices = prefill_fibers[0].devices_dict.values()
 
     def initialize_page_cache(self):
         """Initialize page pool and attention cache."""
@@ -133,15 +141,15 @@ def start(self):
         self.initialize_function_references()
 
         self.prefill_batcher = PrefillBatcherProcess(
-            self.fiber_pool,
+            self.prefill_fiber_pool,
             self.page_cache,
             self.model_params,
             self.prefill_functions,
             self.prog_isolation,
         )
 
         self.decode_batcher = DecodeBatcherProcess(
-            self.fiber_pool,
+            self.decode_fiber_pool,
             self.page_cache,
             self.model_params,
             self.decode_functions,