fix Ray engine crashes on multihost when fetching Jax.array from prefill_ray

Author: sixiang-google

jetstream_pt/ray_engine.py

@@ -113,7 +113,7 @@ def prefill_impl(
     results = ray.get(all_outputs)
     # The prefill function does not return any values;
     # the worker itself manages and maintains the prefill states.
-    return results[0]
+    return None, results[0]
 
   def transfer(self, np_prefix: NpPrefix) -> Any:
     """Store prefill result into object store, then transfer to decode engine workers."""

jetstream_pt/ray_worker.py

@@ -454,7 +454,7 @@ def prefill_ray(
       existing_prefix: Optional[Prefix] = None,
       padded_tokens: PrefillInputs,  # PrefillInputs[np.ndarray],
       true_length: int,
-  ) -> tuple[Prefix, engine_api.ResultTokens]:
+  ) -> engine_api.ResultTokens:
     """Do prefill in ray worker"""
     logits, updated_caches = self.prefill(
         params=params,
@@ -466,9 +466,6 @@ def prefill_ray(
       logits = logits[0]
 
     token = np.argmax(logits[true_length - 1])
-    updated_caches = multihost_utils.process_allgather(
-        updated_caches, tiled=True
-    )
     prefix = Prefix(token, updated_caches, true_length)
     self.prefix_queue.put(prefix, block=False)
 
@@ -490,7 +487,7 @@ def prefill_ray(
         samples_per_slot=1,
     )
 
-    return prefix, result
+    return result
 
   def _convert_to_np_caches(
       self, caches: List[Tuple[jax.Array, jax.Array]]