JetStream changes for Jax based implementation of unified_lora_params for decoding batch of multiple different lora adapters.

jetstream/core/orchestrator.py

@@ -1030,6 +1030,10 @@ def _insert_if_possible(
     # Check if there are any free my_slots. We don't want to block here since
     # 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:
+      adapter_tensorstore = self._generate_adapterstore[idx]
+
     while True:
       my_slots_size = my_slots.qsize()
 
@@ -1100,6 +1104,11 @@ def _insert_if_possible(
           slot=slot,
           # request_id=new_request.request_id,
       )
+
+      if adapter_tensorstore:
+        adapter_tensorstore.insert_adapter_in_cache(
+            new_request.adapter_id, slot)
+
       ThreadDebugLog(
           thread_name,
           f"Generate slice {idx} filled slot {slot} at step "
@@ -1239,6 +1248,10 @@ def _generate_thread(self, idx: int):
     my_generate_backlog = self._generate_backlogs[idx]
     my_detokenize_backlog = self._detokenize_backlogs[idx]
 
+    adapter_tensorstore = None
+    if self._generate_adapterstore and idx < len(self._generate_adapterstore):
+      adapter_tensorstore = self._generate_adapterstore[idx]
+
     # Keep track of what step tokens were generated at.
     generate_timestep = 0
     # State to store things like running kv cache in.
@@ -1304,9 +1317,14 @@ def _generate_thread(self, idx: int):
           my_slots.qsize() < max_concurrent_decodes
       ), "At this point we must have some requests inserted into the slots."
 
+      decoding_adapters_cache = None
+
+      if adapter_tensorstore:
+        decoding_adapters_cache = adapter_tensorstore.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
+          generate_params, decode_state, decoding_adapters_cache
       )
       sampled_tokens.copy_to_host_async()
       # Respond to detokenization backpressure.

jetstream/engine/engine_api.py

@@ -211,6 +211,7 @@ 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.