Make prefilling return first token for loadgen integration (#143)

* Make prefilling return first token for loadgen integration

* minor fix and lint

* enable passing of max_decode_length as a flag

benchmarks/prefill_offline.py

@@ -82,7 +82,7 @@ def prefill_benchmark(tokens_list, engine, params, warmup):
     # pylint: disable-next=all
     warmup_text = "warmup" if warmup else "execute"
     it = time.time()
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params,
         padded_tokens=prefill_tokens,
         true_length=len(prefill_tokens),

benchmarks/run_offline.py

@@ -44,7 +44,7 @@ def run_prefill_time(engine, params, decode_state, seqlen):
   )
 
   for _ in range(3):
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=tokens, true_length=true_length
     )
     decode_state = engine.insert(
@@ -58,7 +58,7 @@ def run_prefill_time(engine, params, decode_state, seqlen):
         jax.profiler.start_trace(FLAGS.profiling_output)
         profiler_started = True
 
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=tokens, true_length=true_length
     )
     decode_state = engine.insert(

jetstream_pt/config.py

@@ -31,6 +31,7 @@
 flags.DEFINE_string("size", "tiny", "size of model")
 flags.DEFINE_bool("quantize_kv_cache", False, "kv_cache_quantize")
 flags.DEFINE_integer("max_cache_length", 1024, "kv_cache_quantize")
+flags.DEFINE_integer("max_decode_length", 1024, "max length of generated text")
 flags.DEFINE_string("sharding_config", "", "config file for sharding")
 flags.DEFINE_bool(
     "shard_on_batch",
@@ -197,6 +198,7 @@ def create_engine_from_config_flags():
       batch_size=FLAGS.batch_size,
       quant_config=quant_config,
       max_cache_length=FLAGS.max_cache_length,
+      max_decode_length=FLAGS.max_decode_length,
       sharding_config=sharding_file_name,
       shard_on_batch=FLAGS.shard_on_batch,
       ragged_mha=FLAGS.ragged_mha,

jetstream_pt/engine.py

@@ -97,7 +97,8 @@ def __init__(
     jax.config.update("jax_enable_x64", False)
 
     self.prefill = jax.jit(
-        self.prefill, out_shardings=self.get_prefix_destination_sharding()
+        self.prefill,
+        out_shardings=(self.get_prefix_destination_sharding(), None),
     )
     self.insert = jax.jit(
         self.insert,
@@ -247,7 +248,7 @@ def prefill(
       existing_prefix: Optional[Prefix] = None,
       padded_tokens: PrefillInputs,  # PrefillInputs[jax.Array],
       true_length: int,
-  ) -> Prefix:
+  ) -> Tuple[Prefix, engine_api.ResultTokens]:
     if isinstance(padded_tokens, jax.Array):
       batched_token = padded_tokens.reshape(1, -1)
     else:
@@ -264,7 +265,6 @@ def prefill(
     )
     if len(logits.shape) == 3:  # b, seqlen, num words
       logits = logits[0]  # seqlen, num words
-
     token = sampling_utils.sampling(
         logits[true_length - 1],
         self.rng,
@@ -273,7 +273,23 @@ def prefill(
         self.env.nucleus_topp,
         self.env.temperature,
     )
-
+    token_out = jnp.reshape(token, (1, 1))
+    data = jnp.concatenate(
+        [
+            token_out,  # First token
+            jnp.ones_like(token_out),  # validity of first token
+            jnp.zeros((1, 1), dtype=jnp.int32),  # length = 0
+        ],
+        axis=-1,
+    )
+    length = token_out.shape[1]
+    result = engine_api.ResultTokens(
+        data=data,
+        tokens_idx=(0, length),
+        valid_idx=(length, 2 * length),
+        length_idx=(2 * length, 2 * length + 1),
+        samples_per_slot=1,
+    )
     # truncate to true_length didnt work need to be out side of jit
     # caches = [
     #   (jax.lax.dynamic_slice_in_dim(
@@ -282,7 +298,7 @@ def prefill(
     #       v, seq_len - true_length, true_length, axis=2))
     #   for k, v in updated_caches
     # ]
-    return Prefix(token, updated_caches, true_length)
+    return Prefix(token, updated_caches, true_length), result
 
   def shrink_prefix(
       self,

run_interactive.py

@@ -66,7 +66,7 @@ def main(argv):
     # pylint: disable-next=all
     if profiling_prefill:
       jax.profiler.start_trace(profiling_output)
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=tokens, true_length=true_length
     )
       # pylint: disable-next=all

run_interactive_disaggregated.py

@@ -161,7 +161,7 @@ def main(argv):
     print(
         f"---- Do prefill in prefill engine pod_slice_name: {prefill_engine.pod_slice_name}"
     )
-    prefill_result = prefill_engine.prefill(
+    prefill_result, _ = prefill_engine.prefill(
         params=None, padded_tokens=tokens, true_length=true_length
     )
     print(

run_interactive_multiple_host.py

@@ -88,7 +88,7 @@ def main(argv):
     print(f"---- Encoded tokens are: {tokens}")
 
     # pylint: disable-next=all
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=None, padded_tokens=tokens, true_length=true_length
     )
     # pylint: disable-next=all

tests/test_llama_e2e.py

@@ -128,7 +128,7 @@ def test_jetstream_llama2_seed(self):
       decode_state = engine.init_decode_state()
       slot = 0
       # pylint: disable-next=all
-      prefill_result = engine.prefill(
+      prefill_result, _ = engine.prefill(
           params=params, padded_tokens=padded_tokens, true_length=true_length
       )
 
@@ -197,7 +197,7 @@ def _llama_e2e(self, env, model_arg):
     decode_state = engine.init_decode_state()
     slot = 0
     # pylint: disable-next=all
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=padded_tokens, true_length=true_length
     )
 
@@ -334,7 +334,7 @@ def test_llama_e2e_two_addtional_tokens(self):
     slot = 0
 
     # pylint: disable-next=all
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=padded_tokens, true_length=true_length
     )
 
@@ -406,7 +406,7 @@ def test_llama_e2e_four_addtional_tokens(self):
     slot = 0
 
     # pylint: disable-next=all
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=padded_tokens, true_length=true_length
     )
 
@@ -472,7 +472,7 @@ def test_llama_with_original_prefill_decode_32(self):
     # pylint: disable-next=all
     decode_state = engine.init_decode_state()
     # pylint: disable-next=all
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=padded_tokens, true_length=true_length
     )
     out_tokens = prefill_result.token
@@ -547,7 +547,7 @@ def test_llama_with_original_prefill_decode(self):
     # pylint: disable-next=all
     decode_state = engine.init_decode_state()
     # pylint: disable-next=all
-    prefill_result = engine.prefill(
+    prefill_result, _ = engine.prefill(
         params=params, padded_tokens=padded_tokens, true_length=true_length
     )
     out_tokens = prefill_result.token