experimental multi-host device_put

Lumosis · Lumosis · commit db36644910cc · 2025-11-14T01:10:51.000Z
diff --git a/tpu_inference/runner/tpu_runner.py b/tpu_inference/runner/tpu_runner.py
@@ -1557,7 +1557,7 @@ def _prepare_inputs_non_dp(self, scheduler_output: "VllmSchedulerOutput"):
         input_tuple_single_device = jax.device_put(
             (input_ids, positions, block_tables, query_start_loc, seq_lens,
              logits_indices, request_distribution),
-            device=self.devices[0],
+            device=jax.local_devices()[0],
         )
         (input_ids, positions, block_tables, query_start_loc, seq_lens,
          logits_indices, request_distribution) = device_array(

Original file line number	Diff line number	Diff line change
`@@ -1557,7 +1557,7 @@ def _prepare_inputs_non_dp(self, scheduler_output: "VllmSchedulerOutput"):`
`1557`	`1557`	`input_tuple_single_device = jax.device_put(`
`1558`	`1558`	`(input_ids, positions, block_tables, query_start_loc, seq_lens,`
`1559`	`1559`	`logits_indices, request_distribution),`
`1560`		`- device=self.devices[0],`
	`1560`	`+ device=jax.local_devices()[0],`
`1561`	`1561`	`)`
`1562`	`1562`	`(input_ids, positions, block_tables, query_start_loc, seq_lens,`
`1563`	`1563`	`logits_indices, request_distribution) = device_array(`