Add unit tests to increase coverage for multi-sampling (#220)

* Add unit tests to increase coverage for multi-samplingh

Author: Lumosis

jetstream/engine/mock_engine.py

@@ -192,7 +192,9 @@ def prefill(
     )
     return (prefix, result_tokens)
 
-  @functools.partial(jax.jit, static_argnums=(0,))
+  @functools.partial(
+      jax.jit, static_argnums=(0,), static_argnames=("num_samples",)
+  )
   def prefill_multisampling(
       self,
       *,
@@ -216,26 +218,30 @@ def prefill_multisampling(
     # Generate dummy prefill cache content
     prefill_cache = padded_tokens[None, :] * params
 
-    # Create a dummy first generated token.
-    first_generated_token = (prefill_cache.sum(axis=-1).astype(jnp.int32))[
-        :, jnp.newaxis
-    ]
+    # Create dummy first generated tokens.
+    first_generated_tokens = []
+    for _ in range(num_samples):
+      first_generated_token = (prefill_cache.sum(axis=-1).astype(jnp.int32))[
+          :, jnp.newaxis
+      ]
+      first_generated_tokens.append(first_generated_token)
+    first_generated_tokens = jnp.concatenate(first_generated_tokens, axis=0)
 
     prefix = Prefix(
         logits=jax.random.normal(self._prng_key, (1, self.vocab_size)),
         cache=prefill_cache,
         next_pos=jnp.full((1, 1), true_length, dtype=jnp.int32),
-        num_generated_tokens=jnp.zeros((1, 1), dtype=jnp.int32),
-        first_token=first_generated_token,
+        num_generated_tokens=jnp.zeros((num_samples, 1), dtype=jnp.int32),
+        first_token=first_generated_tokens,
     )
 
     speculations = first_generated_token.shape[1]
     result_tokens = engine_api.ResultTokens(
         data=jnp.concatenate(
             (
-                first_generated_token,
-                jnp.ones_like(first_generated_token),
-                jnp.ones_like(first_generated_token),
+                first_generated_tokens,
+                jnp.ones_like(first_generated_tokens),
+                jnp.ones_like(first_generated_tokens),
             ),
             axis=-1,
         ),
@@ -244,7 +250,7 @@ def prefill_multisampling(
         valid_idx=(speculations, 2 * speculations),
         # And lengths is rank 1.
         length_idx=(2 * speculations, 2 * speculations + 1),
-        samples_per_slot=self.generate_cache_batch // self.prefill_cache_batch,
+        samples_per_slot=num_samples,
     )
     return (prefix, result_tokens)
 
@@ -398,21 +404,21 @@ def bulk_insert(
     """Insert a single computed prefill cache into multiple slots in
     KV cache.
     """
-    prefill_cache = prefix.cache
+    prefill_cache = decode_state.prefill_cache
     generate_cache = decode_state.generate_cache
     generate_lengths = decode_state.generate_lengths
     generate_tokens = decode_state.generate_tokens
     for slot in slots:
       prefill_cache = jax.lax.dynamic_update_slice_in_dim(
-          decode_state.prefill_cache, prefill_cache, slot, axis=0
+          prefill_cache, prefix.cache, slot, axis=0
       )
       generate_cache = jax.lax.dynamic_update_slice_in_dim(
           generate_cache,
           jnp.zeros((1, self.cache_length)),
           slot,
           axis=0,
       )
-      samples_per_slot = self.generate_cache_batch // self.prefill_cache_batch
+      samples_per_slot = 1
       generate_lengths = jax.lax.dynamic_update_slice_in_dim(
           generate_lengths,
           jnp.ones((samples_per_slot), dtype=jnp.int32),

jetstream/engine/warmup_utils.py

@@ -99,7 +99,13 @@ def initialize_prefill_jit_cache(
   def compile_prefill(length):
     padded_tokens, true_length = jnp.ones((length), dtype="int32"), length
 
-    _, _ = prefill_engine._downstream_engine.prefill(  # pylint: disable=protected-access
+    _, _ = prefill_engine.prefill(
+        params=prefill_params,
+        padded_tokens=padded_tokens,
+        true_length=true_length,
+    )
+
+    _, _ = prefill_engine.prefill_multisampling(
         params=prefill_params,
         padded_tokens=padded_tokens,
         true_length=true_length,
@@ -166,6 +172,10 @@ def compile_insert(length):
 
     generate_engine.insert(prefix=prefill, decode_state=decode_state, slot=0)
 
+    generate_engine.bulk_insert(
+        prefix=prefill, decode_state=decode_state, slots=[0]
+    )
+
     logging.info(
         "---------Generate engine %d compiled for insert length %d.---------",
         generate_idx,

jetstream/tests/core/test_orchestrator.py

@@ -51,7 +51,9 @@
 
 class OrchestratorTest(unittest.IsolatedAsyncioTestCase):
 
-  def _setup_driver(self, interleaved_mode: bool = True):
+  def _setup_driver(
+      self, interleaved_mode: bool = True, multi_sampling: bool = False
+  ):
     prefill_engine = mock_engine.TestEngine(
         batch_size=32, cache_length=256, weight=2.0
     )
@@ -66,6 +68,7 @@ def _setup_driver(self, interleaved_mode: bool = True):
         prefill_params=[prefill_engine.load_params()],
         generate_params=[generate_engine.load_params()],
         interleaved_mode=interleaved_mode,
+        multi_sampling=multi_sampling,
     )
     return driver
 
@@ -150,6 +153,38 @@ async def test_orchestrator(self, interleaved_mode: bool):
     driver.stop()
     print("Orchestrator driver stopped.")
 
+  @parameterized.expand([1, 2, 3, 4])
+  async def test_orchestrator_multi_sampling(self, num_samples: int):
+    """Test the multithreaded orchestration."""
+    driver = self._setup_driver(interleaved_mode=True, multi_sampling=True)
+    client = orchestrator.LLMOrchestrator(driver=driver)
+
+    # The string representation of np.array([[65, 66]]), [2] will be prepend
+    # as BOS.
+    text = "AB"
+
+    request = jetstream_pb2.DecodeRequest(
+        text_content=jetstream_pb2.DecodeRequest.TextContent(text=text),
+        max_tokens=3,
+        num_samples=num_samples,
+    )
+    iterator = client.Decode(request)
+    # chr of [266, 332, 415].
+    expected_text = ["Ċ", "Ō", "Ɵ", ""]
+    expected_token_ids = [266, 332, 415, None]
+    counter = 0
+    async for resp in iterator:
+      for sample in resp.stream_content.samples:
+        output_text = sample.text
+        token_ids = sample.token_ids
+        output_token_id = token_ids[0] if len(token_ids) > 0 else None
+        print(f"actual output: {output_text=} {output_token_id=}")
+        assert output_text == expected_text[counter]
+        assert output_token_id == expected_token_ids[counter]
+      counter += 1
+    driver.stop()
+    print("Orchestrator driver stopped.")
+
   @parameterized.expand([True, False])
   async def test_orchestrator_client_tokenization_chunked_prefill(
       self, interleaved_mode: bool