1) Fixed linter errors.

2) Removed hard-coded scale factor for batch processing. Also added functionality to support different scale_factor values for different adapters.

Author: aman2930

jetstream/core/lora/adapter_tensorstore.py

@@ -121,9 +121,26 @@ def __init__(
         0  # Number of async tasks which are in "loading" state
     )
     self.decoding_adapters_cache: Dict[str, Any] = {}
+
+    # TODO: Make dtype configurable for the scale factor array
+    self.adapters_scale_factor = jnp.empty(1, dtype=jnp.bfloat16)
+
     self.total_slots = total_slots
     self.lock = asyncio.Lock()  # Use an asyncio Lock for thread safety
 
+  def _get_adapter_scale_factor(self, adapter_id: str):
+    """
+    Internal: Get the LoRA scale_factor using the adapter_id.
+    """
+    adapter_config = self.adapter_registry[adapter_id].config
+    lora_scale_factor = float(1)
+
+    if "r" in adapter_config and "lora_alpha" in adapter_config:
+      lora_rank = int(adapter_config["r"])
+      lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank
+
+    return lora_scale_factor
+
   # --- Unsafe Internal methods which assumes that lock is held ---
   def _unsafe_transfer_to_hbm(self, adapter_id: str):
     """
@@ -212,67 +229,87 @@ def _unsafe_unload_adapter(self, adapter_id: str):
 
   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`.
+    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.
+      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
+        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
+          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 None  # Maintain None structure for None leaves
 
+    self.adapters_scale_factor = jnp.ones(self.total_slots, dtype=jnp.bfloat16)
     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.
+    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.
+      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
+        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 
+        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
+        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
+        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.")
+      logging.info(
+          "Empty adapter id. No LoRA tensors added to adapter_tensorstore cache"
+      )
       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 = self._initialize_decoding_adapters_cache(
+          adapter_weights
+      )
 
-    self.decoding_adapters_cache = jax.tree_util.tree_map(insert_leaf,
-                                                         self.decoding_adapters_cache,
-                                                         adapter_weights)
+    adapter_scale_factor = jnp.bfloat16(
+        self._get_adapter_scale_factor(adapter_id)
+    )
+    self.adapters_scale_factor = self.adapters_scale_factor.at[slot_id].set(
+        adapter_scale_factor
+    )
+    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) ---
 

jetstream/core/orchestrator.py

@@ -1106,7 +1106,8 @@ def _insert_if_possible(
 
       if adapter_tensorstore:
         adapter_tensorstore.insert_adapter_in_cache(
-            new_request.adapter_id, slot)
+            new_request.adapter_id, slot
+        )
 
       ThreadDebugLog(
           thread_name,
@@ -1316,11 +1317,23 @@ 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
-        #decode_state["lora_adapter_cache"] = decoding_adapters_cache
+        decoding_adapters_params = adapter_tensorstore.decoding_adapters_cache
+        adapters_scale_factor = adapter_tensorstore.adapters_scale_factor
+        b = adapters_scale_factor.shape[0]
+
+        # Reshaped the scale_factors array to 4-D to align with shape of
+        # the vectors `(batch, hidden_size, num_heads, head_dim)`.
+        reshaped_scale_factors = adapters_scale_factor.reshape((b, 1, 1, 1))
+
+        lora_state = {}
+        lora_state["scale_factor"] = reshaped_scale_factors
+        lora_state["lora_params"] = decoding_adapters_params
+
+        if isinstance(decode_state, dict): # For flax.struct.dataclass
+          decode_state["lora_state"] = lora_state
+        else: # For standard mutable dataclasses.dataclass
+          decode_state = decode_state.replace(lora_state=lora_state)
 
       # Now we actually take a generate step on requests in the slots.
       decode_state, sampled_tokens = generate_engine.generate(

jetstream/engine/mock_engine.py

@@ -32,7 +32,7 @@
 
 import functools
 from dataclasses import asdict
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Dict, Callable, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -71,6 +71,7 @@ class DecodeState:
   generate_cache_index: int
   generate_lengths: jax.Array
   generate_tokens: jax.Array
+  lora_state: Optional[Dict[str, Any]] = None
 
 
 class TestEngine(engine_api.Engine):
@@ -509,6 +510,7 @@ def init_decode_state(self) -> DecodeState:
         generate_tokens=jnp.zeros(
             (self.generate_cache_batch, 1), dtype=jnp.int32
         ),
+        lora_state={},
     )
 
   @property

jetstream/tests/core/lora/test_adapter_tensorstore.py

@@ -145,6 +145,7 @@ async def asyncSetUp(self):
         adapters_dir_path=self.adapters_dir_path,
         hbm_memory_budget=self.hbm_budget,
         cpu_memory_budget=self.cpu_budget,
+        total_slots=8,
     )
 
     # Pre-register adapters for most tests to simplify setup

jetstream/tests/core/test_orchestrator.py

@@ -123,13 +123,15 @@ async def _setup_driver_with_adapterstore(
         adapters_dir_path="/tmp/",
         hbm_memory_budget=20 * (1024**3),  # 20 GB HBM
         cpu_memory_budget=100 * (1024**3),  # 100 GB RAM
+        total_slots=8,
     )
 
     generate_adapterstore = adapterstore.AdapterTensorStore(
         engine=generate_engine,
         adapters_dir_path="/tmp/",
         hbm_memory_budget=20 * (1024**3),  # 20 GB HBM
         cpu_memory_budget=100 * (1024**3),  # 100 GB RAM
+        total_slots=8,
     )
 
     await prefill_adapterstore.register_adapter(