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

tpu_inference/runner/tpu_runner.py

@@ -10,23 +10,24 @@ import jax
 import jax.numpy as jnp
 import jaxtyping
 import numpy as np
-import vllm.envs as vllm_envs
+import torch
+import vllm.envs as envs
 from flax import nnx
 from jax.experimental import mesh_utils
 from jax.sharding import NamedSharding, PartitionSpec
-from torchax.ops.mappings import t2j_dtype
+from torchax.ops.mappings import j2t_dtype
 from vllm.config import VllmConfig
-from vllm.distributed import get_pp_group
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                           has_kv_transfer_group)
 from vllm.forward_context import set_forward_context
+from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput
 from vllm.v1.core.sched.output import SchedulerOutput as VllmSchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,
-                             DraftTokenIds, KVConnectorOutput, LogprobsLists,
+                             DraftTokenIds, KVConnectorOutput,
                              ModelRunnerOutput)
 from vllm.v1.request import Request
 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
@@ -34,22 +35,19 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import \
     KVConnectorModelRunnerMixin
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 
-import tpu_inference.envs as envs
 from tpu_inference import utils as common_utils
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
-from tpu_inference.layers.common.sharding import (MESH_AXIS_NAMES,
-                                                  MESH_AXIS_NAMES_2D,
-                                                  ShardingAxisName,
-                                                  ShardingConfigManager)
 from tpu_inference.layers.jax.sample.rejection_sampler import RejectionSampler
 from tpu_inference.layers.jax.sample.sampling import (compute_logprobs,
                                                       gather_logprobs, sample)
 from tpu_inference.layers.jax.sample.sampling_metadata import \
     TPUSupportedSamplingMetadata
+from tpu_inference.layers.jax.sharding import (MESH_AXIS_NAMES,
+                                               MESH_AXIS_NAMES_2D,
+                                               ShardingAxisName,
+                                               ShardingConfigManager)
 from tpu_inference.logger import init_logger
 from tpu_inference.models.common.model_loader import get_model
-from tpu_inference.models.jax.jax_intermediate_tensor import \
-    JaxIntermediateTensors
 from tpu_inference.models.jax.utils.weight_utils import (
     shard_put, transfer_state_with_mappings)
 from tpu_inference.runner import utils as runner_utils
@@ -66,7 +64,7 @@ from tpu_inference.runner.structured_decoding_manager import \
     StructuredDecodingManager
 from tpu_inference.spec_decode.jax.eagle3 import Eagle3Proposer
 from tpu_inference.utils import (device_array, make_optimized_mesh,
-                                 time_function, to_torch_dtype)
+                                 time_function)
 
 logger = init_logger(__name__)
 
@@ -80,6 +78,17 @@ DUMMY_METADATA = AttentionMetadata(
     request_distribution=[0, 0, 0],
 )
 
+TPU_STR_DTYPE_TO_TORCH_DTYPE = {
+    "half": torch.half,
+    "bfloat16": torch.bfloat16,
+    "float": torch.float,
+    "fp8": torch.float8_e4m3fn,
+    "fp8_e4m3": torch.float8_e4m3fn,
+    "fp8_e5m2": torch.float8_e5m2,
+    "int8": torch.int8,
+    "uint8": torch.uint8,
+}
+
 
 class AsyncTPUModelRunnerOutput(AsyncModelRunnerOutput):
     """Holds asynchronous model output specifically from a TPU runner.
@@ -144,7 +153,6 @@ class ExecuteModelState:
     spec_decode_metadata: Optional[SpecDecodeMetadata]
     kv_connector_output: Optional[KVConnectorOutput]
     logits_indices_selector: Optional[List[int]] = None
-    padded_num_reqs: Optional[int] = None
 
 
 @functools.partial(jax.jit, donate_argnums=(0, 1, 2))
@@ -182,40 +190,12 @@ def _substitute_placeholder_token(
     return input_ids.at[token_in_tpu_cur_input_indices].set(update_values)
 
 
-def _jax_logprobs_to_lists(logprobs_tensors,
-                           logits_indices_selector=None,
-                           cu_num_generated_tokens=None):
-    """Convert JAX LogprobsTensors to LogprobsLists by converting JAX arrays to numpy."""
-    log_token_ids_list = logprobs_tensors.logprob_token_ids.tolist()
-    logprobs_list = logprobs_tensors.logprobs.tolist()
-    selected_token_ranks_list = logprobs_tensors.selected_token_ranks.tolist()
-
-    if logits_indices_selector is not None:
-        log_token_ids_list = [
-            log_token_ids_list[i] for i in logits_indices_selector
-        ]
-        logprobs_list = [logprobs_list[i] for i in logits_indices_selector]
-        selected_token_ranks_list = [
-            selected_token_ranks_list[i] for i in logits_indices_selector
-        ]
-
-    return LogprobsLists(
-        logprob_token_ids=np.asarray(log_token_ids_list),
-        logprobs=np.asarray(logprobs_list),
-        sampled_token_ranks=np.asarray(selected_token_ranks_list),
-        cu_num_generated_tokens=cu_num_generated_tokens,
-    )
-
-
 class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
     def __init__(
         self,
         vllm_config: VllmConfig,
         devices: List[Any],
-        rank: int = 0,
-        is_first_rank: bool = True,
-        is_last_rank: bool = True,
     ):
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
@@ -234,9 +214,6 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         self.maybe_forbid_compile = runner_utils.ForbidCompile(
         ) if envs.VLLM_XLA_CHECK_RECOMPILATION else nullcontext()
         self.dp_size = self.vllm_config.sharding_config.total_dp_size
-        self.rank = rank
-        self.is_first_rank = is_first_rank
-        self.is_last_rank = is_last_rank
 
         self._init_random()
         self._init_mesh()
@@ -247,21 +224,31 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
         # Delegate functions to specific manager classes.
         self.compilation_manager = CompilationManager(self)
-        if self.is_last_rank:
-            self.speculative_decoding_manager = SpeculativeDecodingManager(
-                self)
-            self.structured_decoding_manager = StructuredDecodingManager(self)
+        self.speculative_decoding_manager = SpeculativeDecodingManager(self)
+        self.structured_decoding_manager = StructuredDecodingManager(self)
         self.kv_cache_manager = KVCacheManager(self)
         self.mm_manager = MultiModalManager(self)
         self.persistent_batch_manager = PersistentBatchManager(
             self.requests, self.input_batch, self.encoder_cache,
-            self.uses_mrope, self.model_config, self.is_last_rank)
+            self.uses_mrope, self.model_config)
         self.lora_utils = LoraUtils(self)
 
-        cache_dtype = self.cache_config.cache_dtype
-        if cache_dtype == "auto":
-            cache_dtype = self.dtype
-        self.kv_cache_dtype = to_torch_dtype(cache_dtype)
+        cache_config = self.cache_config
+        if cache_config.cache_dtype == "auto":
+            model_dtype = self.dtype
+            if isinstance(model_dtype, str):
+                self.kv_cache_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
+            elif isinstance(getattr(model_dtype, 'dtype', None), jnp.dtype):
+                self.kv_cache_dtype = j2t_dtype(model_dtype.dtype)
+            elif isinstance(model_dtype, torch.dtype):
+                self.kv_cache_dtype = model_dtype
+            else:
+                raise ValueError(
+                    "KV cache is unsupported for model_dtype of %s",
+                    model_dtype)
+        else:
+            self.kv_cache_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE[
+                cache_config.cache_dtype]
 
         self._pre_async_results: AsyncPreResults | None = None
         self._substitute_placeholder_token_fn = _substitute_placeholder_token
@@ -275,7 +262,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         self.rng_key = jax.random.key(self.model_config.seed)
 
     def _init_mesh(self) -> None:
-        if envs.NEW_MODEL_DESIGN:
+        if os.getenv("NEW_MODEL_DESIGN", False):
             self.mesh = self._create_new_model_mesh()
         else:
             # NOTE(wenxindongwork): The new MoE kernel expects a 2D mesh, so we need
@@ -286,7 +273,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         logger.info(f"Init mesh | mesh={self.mesh}")
 
     def _create_new_model_mesh(self) -> jax.sharding.Mesh:
-        num_slices = envs.NUM_SLICES
+        num_slices = int(os.environ.get('NUM_SLICES', 1))
 
         logger.info(f"Creating new model mesh | devices={len(self.devices)}, "
                     f"num_slices={num_slices}")
@@ -355,7 +342,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                                        devices=self.devices)
 
     def _init_phased_profiling(self) -> None:
-        self.phased_profiling_dir = envs.PHASED_PROFILING_DIR
+        self.phased_profiling_dir = os.getenv("PHASED_PROFILING_DIR", "")
         self.phase_based_profiler = None
         if self.phased_profiling_dir:
             self.phase_based_profiler = runner_utils.PhasedBasedProfiler(
@@ -397,7 +384,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             min_token_size=max(16, self.dp_size),
             max_token_size=scheduler_config.max_num_batched_tokens *
             self.dp_size,
-            padding_gap=vllm_envs.VLLM_TPU_BUCKET_PADDING_GAP)
+            padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP)
         self.num_tokens_paddings_per_dp = [
             padding // self.dp_size for padding in self.num_tokens_paddings
         ]
@@ -421,14 +408,8 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
         self.input_ids_cpu = np.zeros(self.max_num_tokens, dtype=np.int32)
         self.positions_cpu = np.zeros(self.max_num_tokens, dtype=np.int32)
-        # Note: self.input_batch and self.block_tables_cpu are both initialized
-        # with only 1 block_size. For hybrid kv cache, it will be re-init
-        # in kv_cache_manager's maybe_reinitialize_input_batch.
-        self.block_tables_cpu = [
-            np.zeros((self.max_num_reqs, self.max_num_blocks_per_req),
-                     dtype=np.int32)
-        ]
-
+        self.block_table_cpu = np.zeros(
+            (self.max_num_reqs, self.max_num_blocks_per_req), dtype=np.int32)
         self.query_start_loc_cpu = np.zeros(self.max_num_reqs + self.dp_size,
                                             dtype=np.int32)
         self.seq_lens_cpu = np.zeros(self.max_num_reqs, dtype=np.int32)
@@ -462,6 +443,9 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
         # tensors for structured decoding
         self.vocab_size = self.model_config.get_vocab_size()
+        if self.lora_config is not None:
+            # lora_config.lora_extra_vocab_size is the "Maximum size of extra vocabulary that can be present in a LoRA adapter" per https://github.com/vanbasten23/vllm/blob/7f4a8b6705622fde952a2e633e86716f902d6e1b/vllm/config.py#L3040
+            self.vocab_size += self.lora_config.lora_extra_vocab_size
         self.grammar_bitmask_cpu = np.zeros(
             (self.max_num_reqs, cdiv(self.vocab_size, 32)),
             dtype=np.int32,
@@ -506,14 +490,9 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
         self.rng_params_for_sampling = nnx.Rngs(
             jax.random.key(self.model_config.seed)).params()
-        self.is_multimodal_model = (
-            self.model_config.is_multimodal_model
-            and self.get_multimodal_embeddings_fn is not None and hasattr(
-                self.model_config.hf_config, "architectures"
-            )  #TODO: Remove Llama Guard 4 specific condition once the LG4 Vision portion is implemented
-            and len(self.model_config.hf_config.architectures) >= 1
-            and self.model_config.hf_config.architectures[0]
-            != "Llama4ForConditionalGeneration")
+        self.is_multimodal_model = (self.model_config.is_multimodal_model
+                                    and self.get_multimodal_embeddings_fn
+                                    is not None)
 
         logger.info(f"Init model | "
                     f"hbm={common_utils.hbm_usage_gb(self.devices)}GiB")
@@ -526,7 +505,6 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
     def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         self.kv_cache_config = kv_cache_config
-        self.use_hybrid_kvcache = len(kv_cache_config.kv_cache_groups) > 1
         self.kv_caches = []
         self.kv_cache_manager.initialize_kv_cache(kv_cache_config)
         if has_kv_transfer_group():
@@ -539,12 +517,12 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
     def execute_model(
         self,
         scheduler_output: "VllmSchedulerOutput",
-        intermediate_tensors: Optional[JaxIntermediateTensors] = None,
-    ) -> ModelRunnerOutput | JaxIntermediateTensors | None:
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> ModelRunnerOutput | None:
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called "
                                "after execute_model() returns None.")
-        _, output = self._execute_model(scheduler_output, intermediate_tensors)
+        _, output = self._execute_model(scheduler_output)
         return output
 
     def sample_tokens(
@@ -557,17 +535,16 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
         (scheduler_output, attn_metadata, input_ids, hidden_states, logits,
          aux_hidden_states, spec_decode_metadata, kv_connector_output,
-         logits_indices_selector,
-         padded_num_reqs) = (self.execute_model_state.scheduler_output,
-                             self.execute_model_state.attn_metadata,
-                             self.execute_model_state.input_ids,
-                             self.execute_model_state.hidden_states,
-                             self.execute_model_state.logits,
-                             self.execute_model_state.aux_hidden_states,
-                             self.execute_model_state.spec_decode_metadata,
-                             self.execute_model_state.kv_connector_output,
-                             self.execute_model_state.logits_indices_selector,
-                             self.execute_model_state.padded_num_reqs)
+         logits_indices_selector) = (
+             self.execute_model_state.scheduler_output,
+             self.execute_model_state.attn_metadata,
+             self.execute_model_state.input_ids,
+             self.execute_model_state.hidden_states,
+             self.execute_model_state.logits,
+             self.execute_model_state.aux_hidden_states,
+             self.execute_model_state.spec_decode_metadata,
+             self.execute_model_state.kv_connector_output,
+             self.execute_model_state.logits_indices_selector)
         self.execute_model_state = None
 
         if grammar_output is not None:
@@ -581,10 +558,12 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                 logits,
                 arange,
             )
-        return self._sample_from_logits(
-            scheduler_output, attn_metadata, input_ids, hidden_states, logits,
-            aux_hidden_states, spec_decode_metadata, kv_connector_output,
-            logits_indices_selector, padded_num_reqs)
+        return self._sample_from_logits(scheduler_output, attn_metadata,
+                                        input_ids, hidden_states, logits,
+                                        aux_hidden_states,
+                                        spec_decode_metadata,
+                                        kv_connector_output,
+                                        logits_indices_selector)
 
     def _modify_prev_results(self):
         # If copy to host has not been done, we just wait.
@@ -670,9 +649,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
     def _execute_model(
         self,
         scheduler_output: "VllmSchedulerOutput",
-        intermediate_tensors: Optional[JaxIntermediateTensors] = None,
-    ) -> tuple[AttentionMetadata, JaxIntermediateTensors | ModelRunnerOutput
-               | None]:
+    ) -> tuple[AttentionMetadata, ModelRunnerOutput | None]:
         self.persistent_batch_manager.update_states(
             scheduler_output, self.get_mrope_input_positions_fn)
         if not scheduler_output.total_num_scheduled_tokens:
@@ -695,23 +672,13 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         # TODO(pooyam): I guess we can remove returning sampling_metadata in `_prepare_inputs` after https://github.com/njhill/vllm/commit/b7433ca1a47732394b1bdea4099d98389515954b
         (
             input_ids,
-            input_positions,
             attn_metadata,
             _,
             logits_indices,
             spec_decode_metadata,
             logits_indices_selector,
-            padded_num_reqs,
         ) = self._prepare_inputs(scheduler_output)
 
-        is_llama_guard_4 = (
-            hasattr(
-                self.model_config.hf_config, "architectures"
-            )  #TODO: Remove Llama Guard 4 specific condition once the LG4 Vision portion is implemented
-            and len(self.model_config.hf_config.architectures) >= 1
-            and self.model_config.hf_config.architectures[0]
-            == "Llama4ForConditionalGeneration")
-
         # multi-modal support
         if self.is_multimodal_model:
             # Run the multimodal encoder if any.
@@ -719,13 +686,6 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             self.mm_manager.execute_mm_encoder(scheduler_output)
             mm_embeds = self.mm_manager.gather_mm_embeddings(
                 scheduler_output, input_ids.shape[0])
-        #TODO: Remove the follow elif statement once Llama Guard 4 Vision portion has been implemented
-        elif is_llama_guard_4 and any(
-                self.mm_manager.runner.requests[req_id].mm_features
-                for req_id in self.mm_manager.runner.input_batch.req_ids):
-            raise NotImplementedError(
-                "Llama Guard 4 (JAX) currently supports only text inputs. "
-                "Multimodal processing not yet implemented.")
         else:
             mm_embeds = []
 
@@ -750,6 +710,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                     scheduler_output) as kv_connector_output:
                 # NOTE(Wenlong): It takes both `input_ids` and `inputs_embeds`,
                 # but one of them would be `None`
+
                 (self.kv_caches, hidden_states,
                  aux_hidden_states) = self.model_fn(
                      self.state,
@@ -757,17 +718,10 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                      input_ids,
                      attn_metadata,
                      inputs_embeds,
-                     input_positions,
                      tuple(self.layer_name_to_kvcache_index.items()),
                      lora_metadata,
-                     intermediate_tensors,
-                     self.is_first_rank,
-                     self.is_last_rank,
                  )
-            if not get_pp_group().is_last_rank:
-                assert isinstance(hidden_states, JaxIntermediateTensors)
-                hidden_states.kv_connector_output = kv_connector_output
-                return attn_metadata, hidden_states
+
             hidden_states = self._select_from_array_fn(hidden_states,
                                                        logits_indices)
             logits = self.compute_logits_fn(
@@ -785,8 +739,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             aux_hidden_states=aux_hidden_states,
             spec_decode_metadata=spec_decode_metadata,
             kv_connector_output=kv_connector_output,
-            logits_indices_selector=logits_indices_selector,
-            padded_num_reqs=padded_num_reqs)
+            logits_indices_selector=logits_indices_selector)
         return attn_metadata, None
 
     def _sample_from_logits(
@@ -800,44 +753,23 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         spec_decode_metadata: Optional[SpecDecodeMetadata],
         kv_connector_output: Optional[KVConnectorOutput],
         logits_indices_selector: Optional[List[int]] = None,
-        padded_num_reqs: Optional[int] = None,
     ) -> ModelRunnerOutput | AsyncTPUModelRunnerOutput:
-        if padded_num_reqs is None:
-            padded_num_reqs = runner_utils.get_padded_num_reqs_with_upper_limit(
-                self.input_batch.num_reqs, self.max_num_reqs)
-
-        sharding = None
-        if self.dp_size > 1:
-            sharding = NamedSharding(self.mesh,
-                                     PartitionSpec(ShardingAxisName.ATTN_DATA))
-
+        padded_num_reqs = runner_utils.get_padded_num_reqs_with_upper_limit(
+            self.input_batch.num_reqs, self.max_num_reqs)
         tpu_sampling_metadata = TPUSupportedSamplingMetadata.from_input_batch(
-            self.mesh, self.input_batch, padded_num_reqs, sharding=sharding)
-
-        # TODO(pooyam): Should we move this to `_prepare_inputs`?
-        if tpu_sampling_metadata.do_sampling:
-            self.rng_params_for_sampling, step_rng = jax.random.split(
-                self.rng_params_for_sampling)
-        else:
-            step_rng = self.rng_params_for_sampling
-
+            self.mesh, self.input_batch, padded_num_reqs)
         if spec_decode_metadata is None:
             next_tokens = sample(
-                step_rng,
+                self.rng_params_for_sampling,
                 self.mesh,
                 logits,
                 tpu_sampling_metadata,
             )
         else:
-            if tpu_sampling_metadata.do_sampling:
-                bonus_rng, rejection_rng = jax.random.split(step_rng)
-            else:
-                bonus_rng = step_rng
-                rejection_rng = step_rng
             bonus_logits = self._select_from_array_fn(
                 logits, spec_decode_metadata.bonus_logits_indices)
             bonus_token_ids = sample(
-                bonus_rng,
+                self.rng_params_for_sampling,
                 self.mesh,
                 bonus_logits,
                 tpu_sampling_metadata,
@@ -851,7 +783,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                 target_logits=target_logits,
                 bonus_token_ids=bonus_token_ids,
                 sampling_metadata=tpu_sampling_metadata,
-                key=rejection_rng,
+                key=self.rng_params_for_sampling,
             )
 
         if tpu_sampling_metadata.logprobs:
@@ -908,10 +840,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                     logits_indices_selector)
 
             if logprobs is not None:
-                # Map logprobs back to the pre-dp shuffling order
-                logprobs_lists = _jax_logprobs_to_lists(
-                    logprobs, logits_indices_selector)
-
+                logprobs_lists = logprobs.tolists()
             else:
                 logprobs_lists = None
 
@@ -979,9 +908,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             req_state.output_token_ids.extend(sampled_ids)
 
         if logprobs is not None:
-            # Map logprobs back to the pre-dp shuffling order
-            logprobs_lists = _jax_logprobs_to_lists(logprobs,
-                                                    logits_indices_selector)
+            logprobs_lists = logprobs.tolists()
         else:
             logprobs_lists = None
 
@@ -1329,6 +1256,16 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         mrope_positions = self.mrope_positions_cpu[:, :
                                                    padded_total_num_scheduled_tokens]
 
+        block_tables = self.block_table_cpu[:self.max_num_reqs]
+        for dp_rank in range(dp_size):
+            req_offset = dp_rank * max_num_reqs_per_dp_rank
+            _num_reqs = num_req_per_dp_rank[dp_rank]
+
+            block_tables[
+                req_offset:req_offset + _num_reqs, :self.
+                max_num_blocks_per_req] = self.input_batch.block_table[
+                    0].get_cpu_tensor()[req_indices_dp[dp_rank]]
+
         query_start_loc = self.query_start_loc_cpu[:self.max_num_reqs +
                                                    dp_size]
         seq_lens = self.seq_lens_cpu[:self.max_num_reqs]
@@ -1370,59 +1307,20 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         if self.uses_mrope:
             positions = mrope_positions
 
+        # Convert block_tables to 1D on cpu.
+        block_tables = block_tables.reshape(-1)
+
         query_start_loc_cpu = query_start_loc
         logits_indices_cpu = logits_indices
         seq_lens_cpu = seq_lens
 
-        (input_ids, positions, query_start_loc, seq_lens, logits_indices,
-         request_distribution) = device_array(
+        (input_ids, positions, block_tables, query_start_loc, seq_lens,
+         logits_indices, request_distribution, logits_indices) = device_array(
              self.mesh,
-             (input_ids, positions, query_start_loc, seq_lens, logits_indices,
-              request_distribution),
+             (input_ids, positions, block_tables, query_start_loc, seq_lens,
+              logits_indices, request_distribution, logits_indices),
              sharding=data_parallel_attn_sharding,
          )
-
-        attention_metadata_per_layer: Dict[str, AttentionMetadata] = {}
-        uniform_attention_metadata: AttentionMetadata = None
-        for kv_cache_gid, kv_cache_group in enumerate(
-                self.kv_cache_config.kv_cache_groups):
-            block_tables = self.block_tables_cpu[kv_cache_gid][:self.
-                                                               max_num_reqs]
-            for dp_rank in range(dp_size):
-                req_offset = dp_rank * max_num_reqs_per_dp_rank
-                _num_reqs = num_req_per_dp_rank[dp_rank]
-
-                block_tables[
-                    req_offset:req_offset + _num_reqs, :self.
-                    max_num_blocks_per_req] = self.input_batch.block_table[
-                        0].get_cpu_tensor()[req_indices_dp[dp_rank]]
-            # Convert block_tables to 1D on cpu.
-            block_tables = block_tables.reshape(-1)
-            block_tables = device_array(
-                self.mesh,
-                (block_tables),
-                sharding=data_parallel_attn_sharding,
-            )
-
-            attention_metadata_gid = AttentionMetadata(
-                input_positions=positions,
-                block_tables=block_tables,
-                seq_lens=seq_lens,
-                query_start_loc=query_start_loc,
-                request_distribution=request_distribution,
-            )
-
-            # This is for making these cpu buffers hidden during tracing
-            attention_metadata_gid.query_start_loc_cpu = query_start_loc_cpu
-            attention_metadata_gid.seq_lens_cpu = seq_lens_cpu
-
-            if not self.use_hybrid_kvcache:
-                uniform_attention_metadata = attention_metadata_gid
-            else:
-                for layer_name in kv_cache_group.layer_names:
-                    attention_metadata_per_layer[
-                        layer_name] = attention_metadata_gid
-
         # Async scheduling: substitute placeholder tokens for DP
         if self.scheduler_config.async_scheduling and self._pre_async_results is not None:
             # Collect all token indices that need substitution across all DP ranks
@@ -1451,19 +1349,25 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                 padded_total_num_scheduled_tokens,
             )
 
-        if self.use_hybrid_kvcache:
-            attention_metadata = attention_metadata_per_layer
-        else:
-            attention_metadata = uniform_attention_metadata
+        attention_metadata = AttentionMetadata(
+            input_positions=positions,
+            block_tables=block_tables,
+            seq_lens=seq_lens,
+            query_start_loc=query_start_loc,
+            request_distribution=request_distribution,
+        )
+
+        # This is for making these cpu buffers hidden during tracing
+        attention_metadata.query_start_loc_cpu = query_start_loc_cpu
+        attention_metadata.seq_lens_cpu = seq_lens_cpu
+
         return (
             input_ids,
-            positions,
             attention_metadata,
             sampling_metadata,
             logits_indices,
             spec_decode_metadata,
             logits_indices_selector,
-            padded_num_reqs,
         )
 
     def _prepare_inputs_non_dp(self, scheduler_output: "VllmSchedulerOutput"):
@@ -1564,6 +1468,9 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         positions = self.positions_cpu[:padded_total_num_scheduled_tokens]
         mrope_positions = self.mrope_positions_cpu[:, :
                                                    padded_total_num_scheduled_tokens]
+        block_tables = self.block_table_cpu[:self.max_num_reqs]
+        block_tables[:num_reqs, :self.max_num_blocks_per_req] = (
+            self.input_batch.block_table[0].get_cpu_tensor()[:num_reqs])
 
         # TODO(pooyam): Some paddings are up to `num_reqs_paddings` (spec decoding, select hidden states, etc) and some other are to `max_num_reqs` (block table, seq_lens). We should stick to one of them maybe?
         query_start_loc = self.query_start_loc_cpu[:self.max_num_reqs + 1]
@@ -1592,44 +1499,16 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             self.mesh, self.input_batch, padded_num_reqs)
         if self.uses_mrope:
             positions = mrope_positions
+
+        # Convert block_tables to 1D on cpu.
+        block_tables = block_tables.reshape(-1)
+
         query_start_loc_cpu = query_start_loc
         seq_lens_cpu = seq_lens
-
-        (input_ids, positions, query_start_loc, seq_lens,
+        (input_ids, positions, block_tables, query_start_loc, seq_lens,
          logits_indices, request_distribution) = device_array(
-             self.mesh, (input_ids, positions, query_start_loc, seq_lens,
-                         logits_indices, request_distribution))
-
-        attention_metadata_per_layer: Dict[str, AttentionMetadata] = {}
-        uniform_attention_metadata: AttentionMetadata = None
-        for kv_cache_gid, kv_cache_group in enumerate(
-                self.kv_cache_config.kv_cache_groups):
-            block_tables = self.block_tables_cpu[kv_cache_gid][:self.
-                                                               max_num_reqs]
-            block_tables[:num_reqs] = (
-                self.input_batch.block_table[kv_cache_gid].get_cpu_tensor()
-                [:num_reqs])
-            # Convert block_tables to 1D on cpu.
-            block_tables = block_tables.reshape(-1)
-            block_tables = device_array(self.mesh, (block_tables))
-
-            attention_metadata_gid = AttentionMetadata(
-                input_positions=positions,
-                block_tables=block_tables,
-                seq_lens=seq_lens,
-                query_start_loc=query_start_loc,
-                request_distribution=request_distribution)
-            # This is for making these cpu buffers hidden during tracing
-            attention_metadata_gid.query_start_loc_cpu = query_start_loc_cpu
-            attention_metadata_gid.seq_lens_cpu = seq_lens_cpu
-
-            if not self.use_hybrid_kvcache:
-                # all layers share the same attention metadata
-                uniform_attention_metadata = attention_metadata_gid
-            else:
-                for layer_name in kv_cache_group.layer_names:
-                    attention_metadata_per_layer[
-                        layer_name] = attention_metadata_gid
+             self.mesh, (input_ids, positions, block_tables, query_start_loc,
+                         seq_lens, logits_indices, request_distribution))
 
         if self.scheduler_config.async_scheduling and len(
                 token_in_tpu_cur_input_indices) > 0:
@@ -1642,15 +1521,20 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             self.lora_utils.set_active_loras(
                 num_scheduled_tokens_per_req, total_num_scheduled_tokens,
                 padded_total_num_scheduled_tokens)
-        logits_indices_selector = None
 
-        if self.use_hybrid_kvcache:
-            attention_metadata = attention_metadata_per_layer
-        else:
-            attention_metadata = uniform_attention_metadata
-        return (input_ids, positions, attention_metadata, sampling_metadata,
-                logits_indices, spec_decode_metadata, logits_indices_selector,
-                padded_num_reqs)
+        attention_metadata = AttentionMetadata(
+            input_positions=positions,
+            block_tables=block_tables,
+            seq_lens=seq_lens,
+            query_start_loc=query_start_loc,
+            request_distribution=request_distribution)
+
+        # This is for making these cpu buffers hidden during tracing
+        attention_metadata.query_start_loc_cpu = query_start_loc_cpu
+        attention_metadata.seq_lens_cpu = seq_lens_cpu
+        logits_indices_selector = None
+        return (input_ids, attention_metadata, sampling_metadata,
+                logits_indices, spec_decode_metadata, logits_indices_selector)
 
     def _get_input_ids_embeds(self, input_ids: jax.Array,
                               mm_embeds: list[jax.Array]):
@@ -1710,35 +1594,3 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             mappings=mappings,
             transpose_keys=transpose_keys,
             shard=shard)
-
-    def get_intermediate_tensor_spec(self, num_tokens: int):
-        impl = os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower()
-        jax_dtype = t2j_dtype(self.dtype) if impl == "vllm" else self.dtype
-        num_padded_tokens = runner_utils.get_padded_token_len(
-            self.num_tokens_paddings, num_tokens)
-        sharding = NamedSharding(self.mesh, PartitionSpec())
-        hidden_size = self.model_config.get_hidden_size()
-        spec = jax.ShapeDtypeStruct(shape=(num_padded_tokens, hidden_size),
-                                    dtype=jax_dtype,
-                                    sharding=sharding)
-        tensor_spec = {"hidden_states": spec, "residual": spec}
-        return tensor_spec
-
-    def get_uuid_for_jax_transfer(self,
-                                  scheduler_output: "VllmSchedulerOutput",
-                                  rank: int, step: int) -> int:
-        '''
-        Get a uuid for jax.transfer, here we use the hash of
-        scheduler_output + counter_step + sender's rank
-        '''
-        scheduler_output_str = ""
-        if not scheduler_output.num_scheduled_tokens:
-            scheduler_output_str = "empty_batch"
-        else:
-            scheduler_output_str = str(
-                sorted(scheduler_output.num_scheduled_tokens.items()))
-        unique_str = f'{scheduler_output_str} {step} {rank}'
-        import hashlib
-        hasher = hashlib.sha1()
-        hasher.update(unique_str.encode('utf-8'))
-        return int.from_bytes(hasher.digest()[:8], 'big')