After merging with main, some of the code gets overridden. So adding back the code to support multi-LoRA adapters into same batch.

Author: aman2930

jetstream/core/lora/adapter_tensorstore.py

@@ -99,6 +99,7 @@ def __init__(
       adapters_dir_path: str,
       hbm_memory_budget: int,
       cpu_memory_budget: int,
+      total_slots: int,
   ):
     """Initializes the AdapterTensorStore."""
     self.engine = engine  # Possibly MaxEngine object
@@ -119,6 +120,8 @@ def __init__(
     self.running_requests: int = (
         0  # Number of async tasks which are in "loading" state
     )
+    self.decoding_adapters_cache: Dict[str, Any] = {}
+    self.total_slots = total_slots
     self.lock = asyncio.Lock()  # Use an asyncio Lock for thread safety
 
   # --- Unsafe Internal methods which assumes that lock is held ---
@@ -207,6 +210,70 @@ def _unsafe_unload_adapter(self, adapter_id: str):
     metadata.size_hbm = 0
     metadata.size_cpu = 0
 
+  def _initialize_decoding_adapters_cache(self, adapter_weights):
+    """
+    Create a new PyTree with zero tensors at the paths corresponding to non-None leaves
+    in the input PyTree. The zero tensors have an added dimension of size `self.totol_slots`.
+    Args:
+      adatper_weights: The input PyTree, whose structure will be mirrored.
+    Returns:
+      A new PyTree with zero Tensors or None values, mirroring the structure of the input PyTree.
+    """
+    def create_zero_leaf(leaf):
+      if leaf is not None:
+        original_shape = leaf.shape
+        if not original_shape: # handle scalar case
+          zero_tensor_shape = (self.total_slots,)
+        else:
+          zero_tensor_shape = (self.total_slots,) + original_shape  # Prepend a new dimension
+
+        return jnp.zeros(zero_tensor_shape, dtype=leaf.dtype)
+      else:
+        return None   # Maintain None structure for None leaves
+
+    return jax.tree_util.tree_map(create_zero_leaf, adapter_weights)
+
+
+  def insert_adapter_in_cache(self, adapter_id: str, slot_id: int):
+    """
+    Insert the specific adapter tensors into a slot in the serving_adapters_cache.
+    Args:
+      adapter_id: The id of the adapter, whose tensors will be inserted
+      slot_id: The id of slot, which represents the index in the serving_adapter_cache
+        where the adapter tensors will be inserted.
+    """
+
+    def insert_leaf(dest_leaf, source_leaf):
+      if dest_leaf is not None and source_leaf is not None:
+        return dest_leaf.at[slot_id].set(source_leaf)   # Insert at the specific index
+      elif dest_leaf is not None:
+        return dest_leaf      # If source_leaf is None, keep the zero_leaf as is
+      elif source_leaf is not None:       # In this case the adapters have different target modules 
+        original_shape = source_leaf.shape
+        if not original_shape:        # Handle scalar case
+          zero_tensor_shape = (self.total_slots,)
+        else:
+          zero_tensor_shape = (self.total_slots,) + original_shape
+        new_dest_leaf = jnp.zeros(zero_tensor_shape, dtype=source_leaf.dtype)
+        return new_dest_leaf.at[slot_id].set(source_leaf)
+      else:
+        return None           # If both are None, return None
+
+    if adapter_id == "":
+      logging.info("Empty adapter id. So no LoRA tensors inserted into the cache in adapter_tensorStore.")
+      return
+
+    asyncio.run(self.load_adapter(adapter_id, None, True))
+
+    adapter_weights = self.loaded_adapters_hbm[adapter_id]
+
+    if not self.decoding_adapters_cache:
+      self.decoding_adapters_cache = self._initialize_decoding_adapters_cache(adapter_weights)
+
+    self.decoding_adapters_cache = jax.tree_util.tree_map(insert_leaf,
+                                                         self.decoding_adapters_cache,
+                                                         adapter_weights)
+
   # --- Public Methods (Acquire lock, then call unsafe methods) ---
 
   async def register_adapter(

jetstream/core/orchestrator.py

@@ -1031,7 +1031,7 @@ def _insert_if_possible(
     # we can still generate if we can't insert. We do this in a while loop to
     # insert as many sequences as possible.
     adapter_tensorstore = None
-    if self._generate_adapterstore:
+    if self._generate_adapterstore and idx < len(self._generate_adapterstore):
       adapter_tensorstore = self._generate_adapterstore[idx]
 
     while True:
@@ -1102,7 +1102,6 @@ def _insert_if_possible(
           new_request.prefill_result,
           decode_state,
           slot=slot,
-          # request_id=new_request.request_id,
       )
 
       if adapter_tensorstore:
@@ -1321,10 +1320,11 @@ def _generate_thread(self, idx: int):
 
       if adapter_tensorstore:
         decoding_adapters_cache = adapter_tensorstore.decoding_adapters_cache
+        #decode_state["lora_adapter_cache"] = decoding_adapters_cache
 
       # Now we actually take a generate step on requests in the slots.
       decode_state, sampled_tokens = generate_engine.generate(
-          generate_params, decode_state, decoding_adapters_cache
+          generate_params, decode_state
       )
       sampled_tokens.copy_to_host_async()
       # Respond to detokenization backpressure.

jetstream/core/server_lib.py

@@ -174,23 +174,26 @@ def create_driver(
   shared_adapterstore = []
 
   if lora_input_adapters_path:
+    # TODO: Make hbm_memory_budget and cpu_memory_budget configurable
     for pe in engines.prefill_engines:
       prefill_adapterstore.append(
           adapterstore.AdapterTensorStore(
               engine=pe,
               adapters_dir_path=lora_input_adapters_path,
               hbm_memory_budget=20 * (1024**3),  # 20 GB HBM
               cpu_memory_budget=100 * (1024**3),  # 100 GB RAM
+              total_slots=pe.max_concurrent_decodes,
           )
       )
-    # TODO: Make hbm_memory_budget and cpu_memory_budget configurable
+
     for ge in engines.generate_engines:
       generate_adapterstore.append(
           adapterstore.AdapterTensorStore(
               engine=ge,
               adapters_dir_path=lora_input_adapters_path,
               hbm_memory_budget=20 * (1024**3),  # 20 GB HBM
               cpu_memory_budget=100 * (1024**3),  # 100 GB RAM
+              total_slots=ge.max_concurrent_decodes,
           )
       )
 
@@ -201,6 +204,7 @@ def create_driver(
               adapters_dir_path=lora_input_adapters_path,
               hbm_memory_budget=20 * (1024**3),  # 20 GB HBM
               cpu_memory_budget=100 * (1024**3),  # 100 GB RAM
+              total_slots=ie.max_concurrent_decodes,
           )
       )
 
@@ -315,6 +319,9 @@ def run(
         "Not starting Prometheus server: --prometheus_port flag not set"
     )
 
+  if multi_sampling and lora_input_adapters_path:
+    raise ValueError("LoRA adapters is not enabled for multi_sampling mode.")
+
   driver = create_driver(
       config,
       devices,

jetstream/engine/engine_api.py

@@ -211,7 +211,6 @@ def generate(
       params: Params,
       decode_state: DecodeState,
       sampler: Optional[Callable[[Any], Any]] = None,
-      lora_params: Params = None,
   ) -> Tuple[DecodeState, ResultTokens]:
     """Generates tokens for each sequence being decoded in parallel.