tpu-inference 0.11.1.dev202511150811__py3-none-any.whl → 0.11.1.dev202511270815__py3-none-any.whl

tpu_inference/runner/tpu_runner.py

@@ -153,6 +153,7 @@ 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))
@@ -190,18 +191,28 @@ def _substitute_placeholder_token(
     return input_ids.at[token_in_tpu_cur_input_indices].set(update_values)
 
 
-def _reorder_logits_indices(logprobs_lists, logits_indices_selector):
+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=[
-            logprobs_lists.logprob_token_ids[i]
-            for i in logits_indices_selector
-        ],
-        logprobs=[logprobs_lists.logprobs[i] for i in logits_indices_selector],
-        sampled_token_ranks=[
-            logprobs_lists.sampled_token_ranks[i]
-            for i in logits_indices_selector
-        ],
-        cu_num_generated_tokens=logprobs_lists.cu_num_generated_tokens,
+        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,
     )
 
 
@@ -211,6 +222,9 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
         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
@@ -423,8 +437,14 @@ 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)
-        self.block_table_cpu = np.zeros(
-            (self.max_num_reqs, self.max_num_blocks_per_req), 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.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)
@@ -458,9 +478,6 @@ 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,
@@ -505,9 +522,14 @@ 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)
+        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")
 
         logger.info(f"Init model | "
                     f"hbm={common_utils.hbm_usage_gb(self.devices)}GiB")
@@ -520,6 +542,7 @@ 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():
@@ -550,16 +573,17 @@ 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) = (
-             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)
+         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)
         self.execute_model_state = None
 
         if grammar_output is not None:
@@ -573,12 +597,10 @@ 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)
+        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)
 
     def _modify_prev_results(self):
         # If copy to host has not been done, we just wait.
@@ -687,13 +709,23 @@ 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.
@@ -701,6 +733,13 @@ 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 = []
 
@@ -733,6 +772,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                      input_ids,
                      attn_metadata,
                      inputs_embeds,
+                     input_positions,
                      tuple(self.layer_name_to_kvcache_index.items()),
                      lora_metadata,
                  )
@@ -754,7 +794,8 @@ 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)
+            logits_indices_selector=logits_indices_selector,
+            padded_num_reqs=padded_num_reqs)
         return attn_metadata, None
 
     def _sample_from_logits(
@@ -768,23 +809,44 @@ 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:
-        padded_num_reqs = runner_utils.get_padded_num_reqs_with_upper_limit(
-            self.input_batch.num_reqs, self.max_num_reqs)
+        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))
+
         tpu_sampling_metadata = TPUSupportedSamplingMetadata.from_input_batch(
-            self.mesh, self.input_batch, padded_num_reqs)
+            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
+
         if spec_decode_metadata is None:
             next_tokens = sample(
-                self.rng_params_for_sampling,
+                step_rng,
                 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(
-                self.rng_params_for_sampling,
+                bonus_rng,
                 self.mesh,
                 bonus_logits,
                 tpu_sampling_metadata,
@@ -798,7 +860,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                 target_logits=target_logits,
                 bonus_token_ids=bonus_token_ids,
                 sampling_metadata=tpu_sampling_metadata,
-                key=self.rng_params_for_sampling,
+                key=rejection_rng,
             )
 
         if tpu_sampling_metadata.logprobs:
@@ -856,10 +918,8 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
             if logprobs is not None:
                 # Map logprobs back to the pre-dp shuffling order
-                logprobs_lists = logprobs.tolists()
-                if logits_indices_selector is not None:
-                    logprobs_lists = _reorder_logits_indices(
-                        logprobs_lists, logits_indices_selector)
+                logprobs_lists = _jax_logprobs_to_lists(
+                    logprobs, logits_indices_selector)
 
             else:
                 logprobs_lists = None
@@ -929,10 +989,8 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
 
         if logprobs is not None:
             # Map logprobs back to the pre-dp shuffling order
-            logprobs_lists = logprobs.tolists()
-            if logits_indices_selector is not None:
-                logprobs_lists = _reorder_logits_indices(
-                    logprobs_lists, logits_indices_selector)
+            logprobs_lists = _jax_logprobs_to_lists(logprobs,
+                                                    logits_indices_selector)
         else:
             logprobs_lists = None
 
@@ -1280,16 +1338,6 @@ 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]
@@ -1331,20 +1379,59 @@ 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, block_tables, query_start_loc, seq_lens,
-         logits_indices, request_distribution) = device_array(
+        (input_ids, positions, query_start_loc, seq_lens, logits_indices,
+         request_distribution) = device_array(
              self.mesh,
-             (input_ids, positions, block_tables, query_start_loc, seq_lens,
-              logits_indices, request_distribution),
+             (input_ids, positions, query_start_loc, seq_lens, logits_indices,
+              request_distribution),
              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
@@ -1373,25 +1460,19 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
                 padded_total_num_scheduled_tokens,
             )
 
-        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
-
+        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,
         )
 
     def _prepare_inputs_non_dp(self, scheduler_output: "VllmSchedulerOutput"):
@@ -1492,9 +1573,6 @@ 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]
@@ -1523,16 +1601,44 @@ 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, block_tables, query_start_loc, seq_lens,
+
+        (input_ids, positions, query_start_loc, seq_lens,
          logits_indices, request_distribution) = device_array(
-             self.mesh, (input_ids, positions, block_tables, query_start_loc,
-                         seq_lens, logits_indices, request_distribution))
+             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
 
         if self.scheduler_config.async_scheduling and len(
                 token_in_tpu_cur_input_indices) > 0:
@@ -1545,20 +1651,15 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
             self.lora_utils.set_active_loras(
                 num_scheduled_tokens_per_req, total_num_scheduled_tokens,
                 padded_total_num_scheduled_tokens)
-
-        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)
+
+        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)
 
     def _get_input_ids_embeds(self, input_ids: jax.Array,
                               mm_embeds: list[jax.Array]):

tpu_inference/spec_decode/jax/eagle3.py

@@ -9,10 +9,13 @@ import numpy as np
 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.
@@ -51,8 +54,22 @@ 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)
-        del self.state.model['embed_tokens']
-        self.state.model.embed_tokens = target_model.model.embed
+
+        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.")
 
     @functools.partial(jax.jit, static_argnums=(0, ))
     def _prepare_input_ids(

tpu_inference/tpu_info.py

@@ -3,6 +3,7 @@ import os
 
 import requests
 
+from tpu_inference import envs
 from tpu_inference.logger import init_logger
 
 logger = init_logger(__name__)
@@ -32,14 +33,14 @@ def get_tpu_metadata(key: str = "") -> str:
 
 
 def get_tpu_type() -> str:
-    tpu_type = os.getenv("TPU_ACCELERATOR_TYPE", None)
+    tpu_type = envs.TPU_ACCELERATOR_TYPE
     if tpu_type is None:
         tpu_type = get_tpu_metadata(key="accelerator-type")
     return tpu_type
 
 
 def get_node_name() -> str:
-    tpu_name = os.getenv("TPU_NAME", None)
+    tpu_name = envs.TPU_NAME
     if not tpu_name:
         tpu_name = get_tpu_metadata(key="instance-id")
     return tpu_name
@@ -47,7 +48,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 = os.getenv("TPU_WORKER_ID", None)
+    worker_id = envs.TPU_WORKER_ID
     if worker_id is None:
         worker_id = get_tpu_metadata(key="agent-worker-number")
     if worker_id is None:

tpu_inference/utils.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-import os
 import time
 from collections import defaultdict
 from collections.abc import Sequence
@@ -14,8 +13,10 @@ 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.sharding import Mesh, NamedSharding, PartitionSpec
-from vllm import envs, utils
+from vllm import envs as vllm_envs
+from vllm import utils
 
+from tpu_inference import envs
 from tpu_inference.logger import init_logger
 
 GBYTES = 1024 * 1024 * 1024
@@ -57,10 +58,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 envs.VLLM_TPU_USING_PATHWAYS:
+    if vllm_envs.VLLM_TPU_USING_PATHWAYS:
         return pathways_hbm_usage_gb(devices)
 
-    multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
+    multihost_backend = envs.TPU_MULTIHOST_BACKEND
     if multihost_backend == "ray":
         # MemoryStats is only supported for addressable PjRt devices.
         # Assume all the devices have similar memory usage for now.