tpu-inference 0.0.1rc1__py3-none-any.whl → 0.11.1.dev202511130813__py3-none-any.whl

tpu_inference/runner/utils.py

@@ -15,7 +15,6 @@ import jax
 from jax._src.interpreters import pxla
 from vllm.v1.core.sched.output import SchedulerOutput as VllmSchedulerOutput
 
-from tpu_inference import envs
 from tpu_inference.logger import init_logger
 from tpu_inference.runner.input_batch import InputBatch
 
@@ -307,7 +306,8 @@ class PhasedBasedProfiler:
             InferencePhase.BALANCED: False
         }
         self.default_profiling_options = jax.profiler.ProfileOptions()
-        self.default_profiling_options.python_tracer_level = envs.PYTHON_TRACER_LEVEL
+        self.default_profiling_options.python_tracer_level = os.getenv(
+            "PYTHON_TRACER_LEVEL", 0)
 
         self.current_phase: str = ""
 

tpu_inference/spec_decode/jax/eagle3.py

@@ -6,19 +6,13 @@ from typing import Any, Optional
 import jax
 import jax.numpy as jnp
 import numpy as np
-from flax import nnx
-from jax import lax
-from jax.sharding import NamedSharding, PartitionSpec
 from vllm.config import VllmConfig
 
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
-from tpu_inference.logger import init_logger
 from tpu_inference.models.common.model_loader import get_model
 from tpu_inference.runner import utils as runner_utils
 from tpu_inference.utils import device_array
 
-logger = init_logger(__name__)
-
 
 class Eagle3Proposer:
     """A proposer for speculative decoding using the Eagle3 method.
@@ -57,22 +51,8 @@ class Eagle3Proposer:
         """Loads the draft model."""
         self.model_fn, self.compute_logits_fn, self.combine_hidden_states_fn, _, self.state, _, _ = get_model(
             self.vllm_config, self.rng_key, self.mesh, is_draft_model=True)
-
-        draft_embed_tokens = getattr(self.state.model, 'embed_tokens', None)
-        if draft_embed_tokens is None or ~jnp.any(
-                draft_embed_tokens.embedding):
-            logger.info(
-                "Draft model does not have embedding. Setting draft model's embed_tokens to target model's embed"
-            )
-            self.state.model.embed_tokens = target_model.model.embed
-        elif jnp.array_equal(draft_embed_tokens.embedding,
-                             target_model.model.embed.embedding):
-            logger.info(
-                "Draft model's embed_tokens is identical to target model's embed. Sharing the embedding."
-            )
-            self.state.model.embed_tokens = target_model.model.embed
-        else:
-            logger.info("Draft model has its own embed_tokens.")
+        del self.state.model['embed_tokens']
+        self.state.model.embed_tokens = target_model.model.embed
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _prepare_input_ids(
@@ -130,17 +110,6 @@ class Eagle3Proposer:
                                            max_num_blocks_per_req)
         new_block_tables = jnp.where(expanded_exceeds_mask, -1, block_tables)
 
-        positions = lax.with_sharding_constraint(
-            positions, NamedSharding(self.mesh, PartitionSpec(None, )))
-        clamped_positions = lax.with_sharding_constraint(
-            clamped_positions, NamedSharding(self.mesh, PartitionSpec(None, )))
-        new_seq_lens = lax.with_sharding_constraint(
-            new_seq_lens, NamedSharding(self.mesh, PartitionSpec(None, )))
-        query_start_loc = lax.with_sharding_constraint(
-            query_start_loc, NamedSharding(self.mesh, PartitionSpec()))
-        new_block_tables = lax.with_sharding_constraint(
-            new_block_tables, NamedSharding(self.mesh, PartitionSpec(None, )))
-
         return positions, clamped_positions, new_seq_lens, query_start_loc, new_block_tables
 
     @functools.partial(jax.jit, static_argnums=(0, ))
@@ -152,7 +121,6 @@ class Eagle3Proposer:
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _prepare_hidden_states_and_input_ids(
         self,
-        state: nnx.State,
         aux_hidden_states: tuple[jax.Array, ...],
         query_start_loc: jax.Array,
         target_token_ids: jax.Array,
@@ -161,7 +129,7 @@ class Eagle3Proposer:
     ) -> tuple[jax.Array, jax.Array, jax.Array]:
         target_hidden_states = jnp.concatenate(aux_hidden_states, axis=-1)
         target_hidden_states = self.combine_hidden_states_fn(
-            state, target_hidden_states)
+            self.state, target_hidden_states)
 
         input_ids, last_token_indices = self._prepare_input_ids(
             query_start_loc, target_token_ids, next_token_ids, num_reqs)
@@ -208,8 +176,8 @@ class Eagle3Proposer:
                                     block_tables=device_array(
                                         self.mesh, block_tables))
             target_hidden_states, input_ids, last_token_indices = self._prepare_hidden_states_and_input_ids(
-                self.state, aux_hidden_states, attn_metadata.query_start_loc,
-                input_ids, next_token_ids, num_reqs)
+                aux_hidden_states, attn_metadata.query_start_loc, input_ids,
+                next_token_ids, num_reqs)
             return target_hidden_states, input_ids, last_token_indices, attn_metadata
 
         # Host copies from the metadata prepared by the runner.
@@ -273,13 +241,12 @@ class Eagle3Proposer:
 
         attn_metadata = replace(attn_metadata, block_tables=block_tables)
         return self._filter_token_and_prepare_initial_inputs(
-            self.state, token_indices, query_start_loc, seq_lens, input_ids,
+            token_indices, query_start_loc, seq_lens, input_ids,
             aux_hidden_states, attn_metadata, next_token_ids, num_reqs)
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _filter_token_and_prepare_initial_inputs(
         self,
-        state: nnx.State,
         token_indices: jax.Array,
         query_start_loc: jax.Array,
         seq_lens: jax.Array,
@@ -307,51 +274,35 @@ class Eagle3Proposer:
         )
 
         target_hidden_states, input_ids, last_token_indices = self._prepare_hidden_states_and_input_ids(
-            state, [h[token_indices] for h in aux_hidden_states],
-            query_start_loc, target_token_ids, next_token_ids, num_reqs)
+            [h[token_indices] for h in aux_hidden_states], query_start_loc,
+            target_token_ids, next_token_ids, num_reqs)
 
         return target_hidden_states, input_ids, last_token_indices, attn_metadata
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _select_draft_token_ids(
         self,
-        state: nnx.State,
         hidden_states: jax.Array,
         last_token_indices: jax.Array,
     ) -> jax.Array:
         sample_hidden_states = hidden_states[last_token_indices]
-        sample_hidden_states = lax.with_sharding_constraint(
-            sample_hidden_states,
-            NamedSharding(self.mesh, PartitionSpec(None, None)))
-        return self._get_draft_token_ids(state, sample_hidden_states)
+        return self._get_draft_token_ids(sample_hidden_states)
 
     @functools.partial(jax.jit, static_argnums=(0, ))
-    def _get_draft_token_ids(self, state: nnx.State,
-                             hidden_states: jax.Array) -> jax.Array:
+    def _get_draft_token_ids(self, hidden_states: jax.Array) -> jax.Array:
         lora_metadata = None
-        logits = self.compute_logits_fn(state, hidden_states, lora_metadata)
-        draft_token_ids = jnp.argmax(logits, axis=-1)
-        return lax.with_sharding_constraint(
-            draft_token_ids, NamedSharding(self.mesh, PartitionSpec()))
+        logits = self.compute_logits_fn(self.state, hidden_states,
+                                        lora_metadata)
+        return jnp.argmax(logits, axis=-1)
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _select_inputs_for_loop_speculation(
-            self, state: nnx.State, positions: jax.Array, residual: jax.Array,
+            self, positions: jax.Array, residual: jax.Array,
             hidden_states: jax.Array,
             last_token_indices: jax.Array) -> tuple[jax.Array, jax.Array]:
-        positions = positions[last_token_indices]
-        residual = residual[last_token_indices]
-        draft_token_ids = self._select_draft_token_ids(state, hidden_states,
-                                                       last_token_indices)
-
-        positions = lax.with_sharding_constraint(
-            positions, NamedSharding(self.mesh, PartitionSpec(None, )))
-        residual = lax.with_sharding_constraint(
-            residual, NamedSharding(self.mesh, PartitionSpec(None, None)))
-        draft_token_ids = lax.with_sharding_constraint(
-            draft_token_ids, NamedSharding(self.mesh, PartitionSpec()))
-
-        return positions, residual, draft_token_ids
+        return positions[last_token_indices], residual[
+            last_token_indices], self._select_draft_token_ids(
+                hidden_states, last_token_indices)
 
     def propose(
         self,
@@ -378,11 +329,11 @@ class Eagle3Proposer:
 
         if self.num_speculative_tokens == 1:
             return kv_caches, self._select_draft_token_ids(
-                self.state, hidden_states, last_token_indices)
+                hidden_states, last_token_indices)
 
         positions, hidden_states, draft_token_ids = self._select_inputs_for_loop_speculation(
-            self.state, attn_metadata.input_positions, residual[0],
-            hidden_states, last_token_indices)
+            attn_metadata.input_positions, residual[0], hidden_states,
+            last_token_indices)
 
         draft_token_ids_list = [draft_token_ids]
 
@@ -407,8 +358,7 @@ class Eagle3Proposer:
                 attn_metadata,
             )
             hidden_states = residual[0]
-            draft_token_ids = self._get_draft_token_ids(
-                self.state, new_hidden_states)
+            draft_token_ids = self._get_draft_token_ids(new_hidden_states)
             draft_token_ids_list.append(draft_token_ids)
 
         # [batch_size, num_speculative_tokens]

tpu_inference/tpu_info.py

@@ -3,7 +3,6 @@ import os
 
 import requests
 
-from tpu_inference import envs
 from tpu_inference.logger import init_logger
 
 logger = init_logger(__name__)
@@ -33,14 +32,14 @@ def get_tpu_metadata(key: str = "") -> str:
 
 
 def get_tpu_type() -> str:
-    tpu_type = envs.TPU_ACCELERATOR_TYPE
+    tpu_type = os.getenv("TPU_ACCELERATOR_TYPE", None)
     if tpu_type is None:
         tpu_type = get_tpu_metadata(key="accelerator-type")
     return tpu_type
 
 
 def get_node_name() -> str:
-    tpu_name = envs.TPU_NAME
+    tpu_name = os.getenv("TPU_NAME", None)
     if not tpu_name:
         tpu_name = get_tpu_metadata(key="instance-id")
     return tpu_name
@@ -48,7 +47,7 @@ def get_node_name() -> str:
 
 def get_node_worker_id() -> int:
     """For multi-host TPU VM, this returns the worker id for the current node."""
-    worker_id = envs.TPU_WORKER_ID
+    worker_id = os.getenv("TPU_WORKER_ID", None)
     if worker_id is None:
         worker_id = get_tpu_metadata(key="agent-worker-number")
     if worker_id is None:

tpu_inference/utils.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import os
 import time
 from collections import defaultdict
 from collections.abc import Sequence
@@ -8,54 +9,30 @@ from typing import Any, Callable, List, Tuple
 import jax
 import jax.numpy as jnp
 import numpy as np
-import torch
 from jax._src import dtypes
 from jax._src import mesh as mesh_lib
 from jax._src import xla_bridge as xb
 from jax._src.lib import xla_client as xc
-from jax._src.numpy.scalar_types import _ScalarMeta
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
-from torchax.ops.mappings import j2t_dtype, t2j_dtype
-from vllm import envs as vllm_envs
-from vllm import utils
+from vllm import envs, utils
 
-from tpu_inference import envs
 from tpu_inference.logger import init_logger
 
 GBYTES = 1024 * 1024 * 1024
 TPU_HEAD_SIZE_ALIGNMENT = 128
 TPU_SECOND_LAST_MINOR = 8
 
-# Map vllm dtype string that doesn't exactly match jax dtype string name.
-_VLLM_DTYPE_STR_TO_JAX_DTYPE = {
+# This is used to translate from a string name for a dtype
+# to formal jax.numpy DType.  One use case for this is
+# converting the `--kv_cache_dtype` flag to a dtype.
+TPU_STR_DTYPE_TO_JAX_DTYPE = {
+    "bfloat16": jnp.bfloat16,
     "fp8": jnp.float8_e4m3fn,
-    "fp8_e4m3": jnp.float8_e4m3fn,
+    "fp8_e4m3": jnp.float8_e4m3,
     "fp8_e5m2": jnp.float8_e5m2,
+    "int8": jnp.int8,
 }
 
-
-def to_jax_dtype(dtype: str | jnp.dtype | torch.dtype) -> jnp.dtype:
-    if isinstance(dtype, str):
-        if dict_dtype := _VLLM_DTYPE_STR_TO_JAX_DTYPE.get(dtype, None):
-            return dict_dtype
-        return jnp.dtype(dtype)
-    elif isinstance(dtype, torch.dtype):
-        return t2j_dtype(dtype)
-    elif isinstance(dtype, jnp.dtype):
-        return dtype
-    elif isinstance(dtype, _ScalarMeta):
-        return dtype.dtype
-    else:
-        raise ValueError(f"Argument is unsupported data type {type(dtype)}")
-
-
-def to_torch_dtype(dtype: str | jnp.dtype | torch.dtype) -> torch.dtype:
-    # Use jax dtype as an intermediate dtype which we'll be used to convert it
-    # into torch dtype.
-    dtype = to_jax_dtype(dtype)
-    return j2t_dtype(dtype)
-
-
 _megacore = False
 logger = init_logger(__name__)
 
@@ -80,10 +57,10 @@ def get_num_kv_heads_by_tp(num_kv_heads: int, tp_size: int) -> int:
 
 def hbm_usage_bytes(devices: Any) -> List[Tuple[int, int]]:
     usage = []
-    if vllm_envs.VLLM_TPU_USING_PATHWAYS:
+    if envs.VLLM_TPU_USING_PATHWAYS:
         return pathways_hbm_usage_gb(devices)
 
-    multihost_backend = envs.TPU_MULTIHOST_BACKEND
+    multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
     if multihost_backend == "ray":
         # MemoryStats is only supported for addressable PjRt devices.
         # Assume all the devices have similar memory usage for now.
@@ -155,8 +132,8 @@ def pathways_hbm_usage_gb(devices: Any) -> List[Tuple[float, float]]:
     hbm_used = defaultdict(int)
     hbm_limit = get_device_hbm_limit()
     for array in live_arrays:
-        for buffer in array.addressable_shards:
-            hbm_used[buffer.data.device] += buffer.data.nbytes
+        for buffer in array.device_buffers:
+            hbm_used[buffer.device] += buffer.nbytes
     return [(hbm_used[device], hbm_limit) for device in devices]
 
 
@@ -317,8 +294,8 @@ def get_jax_dtype_from_str_dtype(str_dtype: str) -> jnp.dtype:
     Returns:
         jnp.dtype: The JAX dtype.
     """
-    # TODO(kyuyeunk): Replace all reference of this function into TpuDtype.
-    return to_jax_dtype(str_dtype)
+    str_dtype = str_dtype.lower().strip()
+    return TPU_STR_DTYPE_TO_JAX_DTYPE.get(str_dtype)
 
 
 def time_function(func):

tpu_inference/worker/tpu_worker.py

@@ -2,7 +2,6 @@
 
 import os
 import tempfile
-from dataclasses import dataclass, field
 from typing import Callable, Dict, Optional, Tuple
 
 import jax
@@ -11,7 +10,6 @@ import jaxlib
 import jaxtyping
 import vllm.envs as vllm_envs
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed import get_pp_group
 from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,
                                           has_kv_transfer_group)
 from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
@@ -25,13 +23,10 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
 
 from tpu_inference import envs, utils
-from tpu_inference.distributed import jax_parallel_state
 from tpu_inference.distributed.utils import (get_host_ip, get_kv_transfer_port,
                                              get_node_id)
-from tpu_inference.layers.common.sharding import ShardingConfigManager
+from tpu_inference.layers.jax.sharding import ShardingConfigManager
 from tpu_inference.logger import init_logger
-from tpu_inference.models.jax.jax_intermediate_tensor import \
-    JaxIntermediateTensors
 from tpu_inference.runner.kv_cache import get_rpa_page_size_bytes
 from tpu_inference.runner.tpu_runner import TPUModelRunner
 
@@ -44,39 +39,15 @@ _DTYPE: dict[str, jnp.dtype] = {
 }
 
 
-@dataclass
-class PPConfig:
-    rank: int
-    ip: str
-    prev_worker_ip: str
-    pp_world_size: int
-
-    # default env vars for
-    # TPU_PROCESS_BOUNDS, TPU_CHIPS_PER_PROCESS_BOUNDS, TPU_VISIBLE_CHIPS
-    # if PP is used in single host.
-    default_tpu_process_bounds: str = field(init=False)
-    default_tpu_chips_per_process_bounds: str = field(init=False)
-    default_tpu_visible_chips: str = field(init=False)
-
-    def __post_init__(self):
-        self.default_tpu_process_bounds = f"1,{self.pp_world_size},1"
-        self.default_tpu_chips_per_process_bounds = "1,1,1"
-        self.default_tpu_visible_chips = f"{self.rank}"
-
-
 class TPUWorker:
 
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        local_rank: int,
-        rank: int,
-        distributed_init_method: str,
-        is_driver_worker: bool = False,
-        devices=None,
-        ip: str = "localhost",
-        prev_worker_ip: str = "localhost",
-    ):
+    def __init__(self,
+                 vllm_config: VllmConfig,
+                 local_rank: int,
+                 rank: int,
+                 distributed_init_method: str,
+                 is_driver_worker: bool = False,
+                 devices=None):
         # If we use vLLM's model implementation in PyTorch, we should set it
         # with torch version of the dtype.
         impl = envs.MODEL_IMPL_TYPE
@@ -103,12 +74,10 @@ class TPUWorker:
         self.devices = devices if devices is not None else []
         self.device_ranks = set(device.id for device in self.devices
                                 if isinstance(device, jaxlib._jax.Device))
-        self.pp_config = PPConfig(rank, ip, prev_worker_ip,
-                                  self.parallel_config.pipeline_parallel_size)
 
         if self.model_config.trust_remote_code:
             # note: lazy import to avoid importing torch before initializing
-            from vllm.utils.import_utils import init_cached_hf_modules
+            from vllm.utils import init_cached_hf_modules
 
             init_cached_hf_modules()
 
@@ -117,7 +86,7 @@ class TPUWorker:
         # TPU Worker is initialized. The profiler server needs to start after
         # MP runtime is initialized.
         self.profile_dir = None
-        if vllm_envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1 and self.pp_config.pp_world_size == 1:
+        if vllm_envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
             if not self.devices or 0 in self.device_ranks:
                 # For TPU, we can only have 1 active profiler session for 1 profiler
                 # server. So we only profile on rank0.
@@ -125,14 +94,6 @@ class TPUWorker:
                 logger.info("Profiling enabled. Traces will be saved to: %s",
                             self.profile_dir)
 
-        # For PP, we use MPMD so we want to profile every worker.
-        if self.pp_config.pp_world_size > 1 and vllm_envs.VLLM_TORCH_PROFILER_DIR:
-            self.profile_dir = os.path.join(
-                vllm_envs.VLLM_TORCH_PROFILER_DIR,
-                f"pprank_{self.rank}_ppworldsize_{self.pp_config.pp_world_size}"
-            )
-            os.makedirs(self.profile_dir, exist_ok=True)
-
         use_jax_profiler_server = os.getenv("USE_JAX_PROFILER_SERVER", False)
         # Only one instance of profiler is allowed
         if use_jax_profiler_server and self.rank < 1:
@@ -144,87 +105,31 @@ class TPUWorker:
                 )
                 jax.profiler.start_server(jax_profiler_server_port)
 
-        # step_counter is used to calculate uuid to transfer intermediate tensors.
-        self.step_counter = 0
-
     def initialize_cache(self, num_gpu_blocks: int,
                          num_cpu_blocks: int) -> None:
         self.cache_config.num_gpu_blocks = num_gpu_blocks
         self.cache_config.num_cpu_blocks = num_cpu_blocks
 
-    def init_device(self,
-                    tpu_process_bounds="",
-                    tpu_chips_per_process_bounds="",
-                    tpu_visible_chips=""):
-        # set tpu visible devices for Jax runtime in single host PP.
-        multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
-        if multihost_backend != "ray" and self.parallel_config.pipeline_parallel_size > 1:
-            tpu_ports = [
-                jax_parallel_state.BASE_JAX_PORT + i
-                for i in range(self.pp_config.pp_world_size)
-            ]
-            os.environ["TPU_PROCESS_ADDRESSES"] = ",".join(
-                [f"localhost:{port}" for port in tpu_ports])
-            os.environ["TPU_PROCESS_PORT"] = f"{tpu_ports[self.rank]}"
-            os.environ["CLOUD_TPU_TASK_ID"] = f"{self.rank}"
-
-            # Note: Below is the setting for v6e8 host (8 chips of v6e)
-            # Replace with your own topology.
-            # There are 2 ways of subslicing a v6e
-            # 1) 2 slices with 4 TPU chips each, we can do PP=2, TP=1/2/3/4
-            #   TPU_PROCESS_BOUNDS = "1,1,1"
-            #   TPU_CHIPS_PER_PROCESS_BOUNDS = "1,4,1"
-            #   TPU_VISIBLE_CHIPS = "0,1,2,3" or "4,5,6,7"
-            # 2) 1 chip for each subslice, with at most 8 subslices,
-            #    we can do TP=1, PP=1/2/3/4/5/6/7/8
-            os.environ[
-                "TPU_PROCESS_BOUNDS"] = tpu_process_bounds \
-                    if tpu_process_bounds \
-                        else self.pp_config.default_tpu_process_bounds
-            os.environ[
-                "TPU_CHIPS_PER_PROCESS_BOUNDS"] = tpu_chips_per_process_bounds \
-                    if tpu_chips_per_process_bounds \
-                        else self.pp_config.default_tpu_chips_per_process_bounds
-            os.environ[
-                "TPU_VISIBLE_CHIPS"] = tpu_visible_chips \
-                    if tpu_visible_chips \
-                        else self.pp_config.default_tpu_visible_chips
-
+    def init_device(self):
         if not self.devices:
             sharding_config: ShardingConfigManager = self.vllm_config.sharding_config
             device_indexes = sharding_config.device_indexes
             if device_indexes is not None and len(device_indexes) > 0:
                 # Enforcing the devices sequence to be consistent with the specified device indexes
-                all_local_devices = jax.local_devices()
-                device_dict = {
-                    device.id: device
-                    for device in all_local_devices
-                }
+                all_devices = jax.devices()
+                device_dict = {device.id: device for device in all_devices}
                 self.devices = []
                 for device_index in device_indexes:
                     device = device_dict[device_index]
                     if device is None:
                         raise KeyError(
                             f"Device index {device_index} not found in "
-                            f"jax.local_devices() with IDs {list(device_dict.keys())}!"
+                            f"jax.devices() with IDs {list(device_dict.keys())}!"
                         )
                     self.devices.append(device)
-                assert len(self.devices) >= sharding_config.total_devices
                 self.devices = self.devices[:sharding_config.total_devices]
             else:
-                if self.pp_config.pp_world_size > 1:
-                    # We only support a mixed tp + pp scenario that tp size is
-                    #  smaller or equals the total TPUs in one node
-                    # say: we have 4 nodes with 4 TPUs each, we can only do pp:4, tp:4, but not pp:2, tp:8
-                    assert jax.local_device_count(
-                    ) >= sharding_config.total_devices
-                    self.devices = jax.local_devices()[:sharding_config.
-                                                       total_devices]
-                else:
-                    # In a multi-host distributed env, say: Ray, local_device count may smaller
-                    # than the total devices, we just choose the smaller set here.
-                    self.devices = jax.devices()[:sharding_config.
-                                                 total_devices]
+                self.devices = jax.devices()[:sharding_config.total_devices]
 
         # Initialize the vLLM distribution layer as a single chip environment,
         # we'll swap the model's parallel modules with TPU SPMD equivalents.
@@ -241,18 +146,8 @@ class TPUWorker:
                 tensor_model_parallel_size=1,
                 pipeline_model_parallel_size=1,
             )
-
-        jax_parallel_state.init_pp_distributed_environment(
-            self.pp_config.ip,
-            self.rank,
-            self.parallel_config.pipeline_parallel_size,
-            self.devices[0],
-            need_pp=self.parallel_config.pipeline_parallel_size > 1)
-
         ensure_kv_transfer_initialized(self.vllm_config)
-        self.model_runner = TPUModelRunner(
-            self.vllm_config, self.devices, self.rank, self.rank == 0,
-            self.rank == self.pp_config.pp_world_size - 1)
+        self.model_runner = TPUModelRunner(self.vllm_config, self.devices)
         logger.info(f"Init worker | "
                     f"rank={self.rank} | "
                     f"node_id={get_node_id()} | "
@@ -260,12 +155,6 @@ class TPUWorker:
                     f"hbm={utils.hbm_usage_gb(self.devices)}GiB")
         vllm_utils.report_usage_stats(self.vllm_config)
 
-    def initialize_pp_transfer_connect(self):
-        if self.rank == 0:
-            return
-        jax_parallel_state.connect(self.pp_config.prev_worker_ip,
-                                   self.rank - 1)
-
     def determine_available_memory(self) -> int:
         gpu_memory_utilization = self.cache_config.gpu_memory_utilization
         hbm_usage = utils.hbm_usage_bytes(self.devices)
@@ -305,39 +194,14 @@ class TPUWorker:
         # deliberate, temporary compromise for the same reasons outlined in
         # the `get_kv_cache_spec` method.
 
-        if self.parallel_config.pipeline_parallel_size == 1 or self.rank == 0:
-            intermediate_tensors = None
-        else:
-            # receive intermediate tensors
-            uuid = self.model_runner.get_uuid_for_jax_transfer(
-                scheduler_output, self.rank - 1, self.step_counter)
-            # TODO: this method might only works for vllm model, not sure about jax models.
-            tensor_spec = self.model_runner.get_intermediate_tensor_spec(
-                scheduler_output.total_num_scheduled_tokens)
-            intermediate_tensors_dict = get_pp_group().recv_tensor_dict(
-                uuid, tensor_spec)
-            intermediate_tensors = JaxIntermediateTensors(
-                intermediate_tensors_dict)
-
-        output = self.model_runner.execute_model(scheduler_output,
-                                                 intermediate_tensors)
-
-        if isinstance(output, JaxIntermediateTensors):
-            assert self.parallel_config.pipeline_parallel_size > 1
-            assert not get_pp_group().is_last_rank
-            # send intermediate tensors
-            uuid = self.model_runner.get_uuid_for_jax_transfer(
-                scheduler_output, self.rank, self.step_counter)
-            get_pp_group().send_tensor_dict(uuid, output.tensors)
-            self.step_counter += 1
-            return None
-        else:
-            self.step_counter += 1
-            # With a connector, the scheduler expects output from all workers
-            # TODO(mrjunwan): Figure out if this is ok after https://github.com/vllm-project/vllm/pull/26866
-            if has_kv_transfer_group():
-                return output
-            return output if self.is_driver_worker else None
+        output = self.model_runner.execute_model(scheduler_output)
+
+        # With a connector, the scheduler expects output from all workers
+        # TODO(mrjunwan): Figure out if this is ok after https://github.com/vllm-project/vllm/pull/26866
+        if has_kv_transfer_group():
+            return output
+
+        return output if self.is_driver_worker else None
 
     def sample_tokens(self,
                       grammar_output: GrammarOutput) -> ModelRunnerOutput:
@@ -357,7 +221,7 @@ class TPUWorker:
         if is_start:
             options = jax.profiler.ProfileOptions()
             # default: https://docs.jax.dev/en/latest/profiling.html#general-options
-            options.python_tracer_level = envs.PYTHON_TRACER_LEVEL
+            options.python_tracer_level = os.getenv("PYTHON_TRACER_LEVEL", 0)
             options.host_tracer_level = os.getenv("HOST_TRACER_LEVEL", 1)
             jax.profiler.start_trace(self.profile_dir,
                                      profiler_options=options)
@@ -402,8 +266,7 @@ class TPUWorker:
 
         # TODO(kyuyeunk): Instead of checking page_size_bytes here, introduce
         # feature that allows overriding page_size_bytes of KVCacheSpec.
-        vllm_page_size_bytes = get_uniform_page_size(
-            list(kv_cache_specs.values()))
+        vllm_page_size_bytes = get_uniform_page_size(kv_cache_specs)
         rpa_page_size_bytes = get_rpa_page_size_bytes(self.model_runner.mesh,
                                                       kv_cache_specs)