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

Author: sixiang-google

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]:
+    ) -> tuple[Prefix, 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 None, result
 
   def _convert_to_np_caches(
       self, caches: List[Tuple[jax.Array, jax.Array]]