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

tpu_inference/platforms/tpu_platform.py

@@ -1,22 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, cast
+import os
+from typing import TYPE_CHECKING, Optional, Tuple, Union, cast
 
 import jax.numpy as jnp
-import torch
 import vllm.envs as vllm_envs
+from torchax.ops.mappings import j2t_dtype
 from tpu_info import device
 from vllm.inputs import ProcessorInputs, PromptType
 from vllm.platforms.interface import Platform, PlatformEnum
 from vllm.sampling_params import SamplingParams, SamplingType
 
 from tpu_inference import envs
-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.utils import to_jax_dtype, to_torch_dtype
 
 if TYPE_CHECKING:
-    from vllm.attention.backends.registry import AttentionBackendEnum
+    from vllm.attention.backends.registry import _Backend
     from vllm.config import BlockSize, ModelConfig, VllmConfig
     from vllm.pooling_params import PoolingParams
 else:
@@ -24,10 +24,16 @@ else:
     ModelConfig = None
     VllmConfig = None
     PoolingParams = None
-    AttentionBackendEnum = None
+    _Backend = None
 
 logger = init_logger(__name__)
 
+_DTYPE: dict[str, jnp.dtype] = {
+    "bfloat16": jnp.bfloat16,
+    "float": jnp.float32,
+    "float32": jnp.float32,
+}
+
 
 class TpuPlatform(Platform):
     _enum = PlatformEnum.TPU
@@ -48,13 +54,12 @@ class TpuPlatform(Platform):
     ]
 
     @classmethod
-    def get_attn_backend_cls(cls, selected_backend: "AttentionBackendEnum",
-                             head_size: int, dtype: jnp.dtype,
-                             kv_cache_dtype: Optional[str], block_size: int,
-                             use_v1: bool, use_mla: bool, has_sink: bool,
-                             use_sparse: bool, attn_type: Any) -> str:
-        from vllm.attention.backends.registry import AttentionBackendEnum
-        if selected_backend != AttentionBackendEnum.PALLAS:
+    def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
+                             dtype: jnp.dtype, kv_cache_dtype: Optional[str],
+                             block_size: int, use_v1: bool, use_mla: bool,
+                             has_sink: bool, use_sparse: bool) -> str:
+        from vllm.attention.backends.registry import _Backend
+        if selected_backend != _Backend.PALLAS:
             logger.info("Cannot use %s backend on TPU.", selected_backend)
 
         if use_v1:
@@ -77,14 +82,6 @@ class TpuPlatform(Platform):
             logger.warning(f"Error getting device name: {e}")
             return 'TPU'
 
-    @classmethod
-    def fp8_dtype(cls) -> torch.dtype:
-        if cls.get_device_name().lower() == "tpu v6e":
-            logger.info(
-                "Automatically using fp8_e5m2 for FP8 KV cache on TPU v6e.")
-            return torch.float8_e5m2
-        return torch.float8_e4m3fn
-
     @classmethod
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         raise NotImplementedError
@@ -135,7 +132,6 @@ class TpuPlatform(Platform):
         # For v0, the default block size is 16.
         if cache_config and cache_config.block_size is None:
             cache_config.block_size = cast(BlockSize, 16)
-
         compilation_config = vllm_config.compilation_config
 
         # TPU only supports DYNAMO_TRACE_ONCE compilation level
@@ -152,19 +148,20 @@ class TpuPlatform(Platform):
         # NOTE(xiang): convert dtype to jnp.dtype
         # NOTE(wenlong): skip this logic for mm model preprocessing
         # For mm model preprocessors, it may need the output dtype to be torch.
-        # In order to avoid a PR to vLLM, we postpone the dtype checking during
-        # tpu_worker initialization
+        # In order to avoid a PR to vLLM, we postpone the dtype checking during tpu_worker initialization
         if not vllm_config.scheduler_config.is_multimodal_model or impl == "vllm":
-            model_dtype = vllm_config.model_config.dtype
-            try:
-                dtype = to_jax_dtype(model_dtype)
-            except ValueError:
-                logger.warning(f"{model_dtype=} is not supported. "
-                               "Falling back to jnp.bfloat16")
-                dtype = jnp.bfloat16
-            if impl == "vllm":
-                dtype = to_torch_dtype(dtype)
-            vllm_config.model_config.dtype = dtype
+            if not isinstance(vllm_config.model_config.dtype, str):
+                logger.warning(
+                    "The model dtype is not properly set for JAX backend. "
+                    "Overwriting it to jnp.bfloat16")
+                vllm_config.model_config.dtype = jnp.bfloat16
+            else:
+                vllm_config.model_config.dtype = _DTYPE.get(
+                    vllm_config.model_config.dtype, jnp.bfloat16)
+
+        if impl == "vllm":
+            vllm_config.model_config.dtype = j2t_dtype(
+                vllm_config.model_config.dtype.dtype)
 
         # TODO(cuiq): remove this dependency.
         from vllm.v1.attention.backends.pallas import PallasAttentionBackend
@@ -185,16 +182,10 @@ class TpuPlatform(Platform):
         parallel_config.worker_cls = \
                         "tpu_inference.worker.tpu_worker.TPUWorker"
 
-        multihost_backend = envs.TPU_MULTIHOST_BACKEND
+        multihost_backend = os.environ.get("TPU_MULTIHOST_BACKEND", "").lower()
         if not multihost_backend:  # Single host
-            if parallel_config.pipeline_parallel_size == 1:
-                logger.info("Force using UniProcExecutor for JAX on \
-                        single host without pipeline parallelism.")
-                parallel_config.distributed_executor_backend = "uni"
-            else:
-                logger.info("Force using MultiprocExecutor for JAX on \
-                        single host with pipeline parallelism.")
-                parallel_config.distributed_executor_backend = "mp"
+            logger.info("Force using UniProcExecutor for JAX on single host.")
+            parallel_config.distributed_executor_backend = "uni"
         elif multihost_backend == "ray":
             from tpu_inference.executors.ray_distributed_executor import \
                 RayDistributedExecutor
@@ -269,7 +260,3 @@ class TpuPlatform(Platform):
         Returns if the current platform needs to sync weight loader.
         """
         return True
-
-    @classmethod
-    def support_hybrid_kv_cache(cls) -> bool:
-        return True

tpu_inference/runner/compilation_manager.py

@@ -1,22 +1,20 @@
+import os
 import time
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-import vllm.envs as vllm_envs
+import vllm.envs as envs
 from jax.sharding import NamedSharding, PartitionSpec
 
-import tpu_inference.envs as envs
 from tpu_inference.core.disagg_utils import is_disagg_enabled
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
-from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.layers.jax.sample.sampling import sample
 from tpu_inference.layers.jax.sample.sampling_metadata import \
     TPUSupportedSamplingMetadata
+from tpu_inference.layers.jax.sharding import ShardingAxisName
 from tpu_inference.logger import init_logger
-from tpu_inference.models.jax.jax_intermediate_tensor import \
-    JaxIntermediateTensors
 from tpu_inference.utils import device_array
 
 if TYPE_CHECKING:
@@ -32,10 +30,10 @@ class CompilationManager:
 
     def __init__(self, runner: "TPUModelRunner"):
         self.runner = runner
-        if not vllm_envs.VLLM_DISABLE_COMPILE_CACHE:
+        if not envs.VLLM_DISABLE_COMPILE_CACHE:
             logger.info("Enabling JAX compile cache.")
             jax.config.update("jax_compilation_cache_dir",
-                              vllm_envs.VLLM_XLA_CACHE_PATH)
+                              envs.VLLM_XLA_CACHE_PATH)
 
     def _create_dummy_tensor(self,
                              shape: Tuple[int, ...],
@@ -69,7 +67,8 @@ class CompilationManager:
         logger.info("Compilation finished in %.2f [secs].", end - start)
 
     def capture_model(self) -> None:
-        if envs.SKIP_JAX_PRECOMPILE or self.runner.model_config.enforce_eager:
+        if os.getenv("SKIP_JAX_PRECOMPILE",
+                     False) or self.runner.model_config.enforce_eager:
             return
         logger.info("Precompile all the subgraphs with possible input shapes.")
 
@@ -82,8 +81,6 @@ class CompilationManager:
                 self._precompile_backbone_with_inputs_embeds()
             if self.runner.scheduler_config.async_scheduling:
                 self._precompile_substitute_placeholder_token()
-            if not self.runner.is_last_rank:
-                return
             self._precompile_select_from_array()
             self._precompile_compute_logits()
             self._precompile_disagg_utils()
@@ -123,15 +120,8 @@ class CompilationManager:
                 num_tokens=num_tokens,
             )
 
-    def _precompile_backbone_helper(self,
-                                    name,
-                                    *,
-                                    input_ids,
-                                    positions,
-                                    inputs_embeds,
-                                    intermediate_tensors=None,
-                                    is_first_rank=True,
-                                    is_last_rank=True) -> None:
+    def _precompile_backbone_helper(self, name, *, input_ids, positions,
+                                    inputs_embeds) -> None:
         num_tokens = None
         if input_ids is not None:
             num_tokens = input_ids.shape[0]
@@ -145,6 +135,12 @@ class CompilationManager:
                 ShardingAxisName.ATTN_DATA, )) if dp_size > 1 else None
 
         # Keep existing pattern for complex array operations
+        block_tables = self.runner.block_table_cpu[:self.runner.max_num_reqs]
+        block_tables = block_tables.reshape(-1)
+        block_tables = device_array(self.runner.mesh,
+                                    block_tables,
+                                    sharding=dp_sharding)
+
         seq_lens = self._create_dummy_tensor((self.runner.max_num_reqs, ),
                                              jnp.int32, dp_sharding)
         query_start_loc = self._create_dummy_tensor(
@@ -156,49 +152,26 @@ class CompilationManager:
                                             request_distribution,
                                             sharding=dp_sharding)
 
-        attention_metadata_per_layer: Dict[str, AttentionMetadata] = {}
-        uniform_attention_metadata: AttentionMetadata = None
-        for kv_cache_gid, kv_cache_group in enumerate(
-                self.runner.kv_cache_config.kv_cache_groups):
-            block_tables = self.runner.block_tables_cpu[
-                kv_cache_gid][:self.runner.max_num_reqs]
-            block_tables = block_tables.reshape(-1)
-            block_tables = device_array(self.runner.mesh,
-                                        block_tables,
-                                        sharding=dp_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,
-            )
-            if not self.runner.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
+        attention_metadata = AttentionMetadata(
+            input_positions=positions,
+            block_tables=block_tables,
+            seq_lens=seq_lens,
+            query_start_loc=query_start_loc,
+            request_distribution=request_distribution,
+        )
 
         def model_fn_wrapper(
             state,
             kv_caches,
             input_ids,
             attention_metadata,
-            positions,
             inputs_embeds,
             layer_name_to_kvcache_index,
             lora_metadata,
-            intermediate_tensors,
-            is_first_rank,
-            is_last_rank,
         ):
             kv_caches, hidden_states, _ = self.runner.model_fn(
                 state, kv_caches, input_ids, attention_metadata, inputs_embeds,
-                positions, layer_name_to_kvcache_index, lora_metadata,
-                intermediate_tensors, is_first_rank, is_last_rank)
+                layer_name_to_kvcache_index, lora_metadata)
             self.runner.kv_caches = kv_caches
             return hidden_states
 
@@ -206,10 +179,6 @@ class CompilationManager:
                 self.runner.lora_config, np.array([num_tokens],
                                                   dtype=np.int32)):
             lora_metadata = self.runner.lora_utils.extract_lora_metadata()
-            if self.runner.use_hybrid_kvcache:
-                attention_metadata = attention_metadata_per_layer
-            else:
-                attention_metadata = uniform_attention_metadata
             self._run_compilation(
                 name,
                 model_fn_wrapper,
@@ -217,13 +186,9 @@ class CompilationManager:
                 self.runner.kv_caches,
                 input_ids,
                 attention_metadata,
-                positions,
                 inputs_embeds,
                 tuple(self.runner.layer_name_to_kvcache_index.items()),
                 lora_metadata,
-                intermediate_tensors,
-                is_first_rank,
-                is_last_rank,
                 num_tokens=num_tokens,
             )
 
@@ -274,7 +239,6 @@ class CompilationManager:
                 )
 
     def _precompile_backbone_text_only(self) -> None:
-        hidden_size = self.runner.model_config.get_hidden_size()
         for num_tokens in self.runner.num_tokens_paddings:
             dp_sharding = NamedSharding(
                 self.runner.mesh, PartitionSpec(ShardingAxisName.ATTN_DATA, )
@@ -284,28 +248,10 @@ class CompilationManager:
                                                   dp_sharding)
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32,
                                                   dp_sharding)
-            is_first_rank = self.runner.is_first_rank
-            is_last_rank = self.runner.is_last_rank
-            if is_first_rank:
-                intermediate_tensors = None
-            else:
-                hidden_states = self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16)
-                residual = self._create_dummy_tensor((num_tokens, hidden_size),
-                                                     jnp.bfloat16)
-                intermediate_tensors = JaxIntermediateTensors(
-                    tensors={
-                        "hidden_states": hidden_states,
-                        "residual": residual
-                    })
-            self._precompile_backbone_helper(
-                f"worker{self.runner.rank} backbone",
-                input_ids=input_ids,
-                positions=positions,
-                inputs_embeds=None,
-                intermediate_tensors=intermediate_tensors,
-                is_first_rank=is_first_rank,
-                is_last_rank=is_last_rank)
+            self._precompile_backbone_helper("backbone",
+                                             input_ids=input_ids,
+                                             positions=positions,
+                                             inputs_embeds=None)
 
     def _precompile_backbone_with_inputs_embeds(self) -> None:
         hidden_size = self.runner.model_config.get_hidden_size()
@@ -319,28 +265,10 @@ class CompilationManager:
             else:
                 positions = self._create_dummy_tensor((num_tokens, ),
                                                       jnp.int32)
-            is_first_rank = self.runner.is_first_rank
-            is_last_rank = self.runner.is_last_rank
-            if not is_first_rank:
-                hidden_states = self._create_dummy_tensor(
-                    (num_tokens, hidden_size), jnp.bfloat16)
-                residual = self._create_dummy_tensor((num_tokens, hidden_size),
-                                                     jnp.bfloat16)
-                intermediate_tensors = JaxIntermediateTensors(
-                    tensors={
-                        "hidden_states": hidden_states,
-                        "residual": residual
-                    })
-            else:
-                intermediate_tensors = None
-            self._precompile_backbone_helper(
-                f"worker{self.runner.rank} backbone with embeds",
-                input_ids=None,
-                positions=positions,
-                inputs_embeds=inputs_embeds,
-                intermediate_tensors=intermediate_tensors,
-                is_first_rank=is_first_rank,
-                is_last_rank=is_last_rank)
+            self._precompile_backbone_helper("backbone with embeds",
+                                             input_ids=None,
+                                             positions=positions,
+                                             inputs_embeds=inputs_embeds)
 
     def _precompile_select_from_array_helper(
         self,
@@ -404,23 +332,20 @@ class CompilationManager:
             index_paddings = self.runner.num_reqs_paddings
         dp_sharding = NamedSharding(self.runner.mesh,
                                     PartitionSpec(ShardingAxisName.ATTN_DATA))
-        hidden_states_sharding = NamedSharding(
-            self.runner.mesh, PartitionSpec(ShardingAxisName.ATTN_DATA, None))
         dp_size = self.runner.vllm_config.sharding_config.total_dp_size
         self._precompile_select_from_array_helper(
-            name=f"worker{self.runner.rank} select all logits",
+            name="select all logits",
             source_paddings=self.runner.num_tokens_paddings,
             indices_paddings=index_paddings,
             hidden_dim=hsize,
-            input_sharding=hidden_states_sharding,
+            input_sharding=dp_sharding,
             indices_sharding=dp_sharding if dp_size > 1 else None,
         )
 
         if self.runner.speculative_config:
             vocab_size = self.runner.model_config.get_vocab_size()
             self._precompile_select_from_array_helper(
-                name=
-                f"worker{self.runner.rank} select bonus tokens for spec decoding",
+                name="select bonus tokens for spec decoding",
                 source_paddings=self.runner.num_logits_paddings,
                 indices_paddings=self.runner.num_reqs_paddings,
                 hidden_dim=vocab_size,
@@ -428,8 +353,7 @@ class CompilationManager:
                                              PartitionSpec(None, "model")),
             )
             self._precompile_select_from_array_helper(
-                name=
-                f"worker{self.runner.rank} select target tokens for spec decoding",
+                name="select target tokens for spec decoding",
                 source_paddings=self.runner.num_logits_paddings,
                 indices_paddings=self.runner.num_logits_paddings,
                 hidden_dim=vocab_size,
@@ -452,7 +376,7 @@ class CompilationManager:
                     np.array([num_reqs], dtype=np.int32)):
                 lora_metadata = self.runner.lora_utils.extract_lora_metadata()
                 self._run_compilation(
-                    f"worker{self.runner.rank} compute_logits",
+                    "compute_logits",
                     self.runner.compute_logits_fn,
                     self.runner.state,
                     hidden_states,
@@ -494,7 +418,7 @@ class CompilationManager:
                     do_sampling=do_sampling,
                 )
                 self._run_compilation(
-                    f"worker{self.runner.rank} sample",
+                    "sample",
                     sample,
                     self.runner.rng_params_for_sampling,
                     self.runner.mesh,
@@ -535,7 +459,7 @@ class CompilationManager:
             logits = self._create_dummy_tensor((num_reqs, hsize), jnp.bfloat16)
             token_ids = self._create_dummy_tensor((num_reqs, ), jnp.int32)
             self._run_compilation(
-                f"worker{self.runner.rank} gather_logprobs",
+                "gather_logprobs",
                 self.runner._compute_and_gather_logprobs,
                 logits,
                 token_ids,
@@ -587,7 +511,7 @@ class CompilationManager:
                             do_sampling=do_sampling)
 
                     self._run_compilation(
-                        f"worker{self.runner.rank} {compilation_name}",
+                        compilation_name,
                         self.runner.rejection_sampler,
                         draft_token_ids,
                         num_draft_tokens,
@@ -604,9 +528,7 @@ class CompilationManager:
     def _precompile_eagle3_helpers(self) -> None:
         logger.info(
             "Compiling eagle3 jitted helpers with different input shapes.")
-        target_hidden_size = self.runner.model_config.get_hidden_size()
-        draft_hidden_size = self.runner.speculative_config.draft_model_config.get_hidden_size(
-        )
+        hidden_size = self.runner.model_config.get_hidden_size()
         dtype = self.runner.model_config.dtype
 
         num_kv_cache_groups = len(self.runner.kv_cache_config.kv_cache_groups)
@@ -653,11 +575,10 @@ class CompilationManager:
 
         for num_logits in self.runner.num_logits_paddings:
             hidden_states = self._create_dummy_tensor(
-                (num_logits, draft_hidden_size), jnp.bfloat16)
+                (num_logits, hidden_size), jnp.bfloat16)
             self._run_compilation(
                 "eagle3_get_draft_token_ids",
                 self.runner.drafter._get_draft_token_ids,
-                self.runner.drafter.state,
                 hidden_states,
                 num_logits=num_logits,
             )
@@ -665,8 +586,8 @@ class CompilationManager:
         input_ids_loop = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32,
             NamedSharding(self.runner.mesh, PartitionSpec()))
-        draft_hidden_state_loop = self._create_dummy_tensor(
-            (self.runner.max_num_reqs, draft_hidden_size), dtype,
+        target_hidden_state_loop = self._create_dummy_tensor(
+            (self.runner.max_num_reqs, hidden_size), dtype,
             NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
         next_token_ids = self._create_dummy_tensor(
             (self.runner.max_num_reqs, ), jnp.int32)
@@ -674,12 +595,9 @@ class CompilationManager:
             (self.runner.max_num_reqs, ), jnp.int32)
         for num_tokens in self.runner.num_tokens_paddings:
             aux_hidden_states = [
-                self._create_dummy_tensor((num_tokens, target_hidden_size),
-                                          dtype),
-                self._create_dummy_tensor((num_tokens, target_hidden_size),
-                                          dtype),
-                self._create_dummy_tensor((num_tokens, target_hidden_size),
-                                          dtype),
+                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
+                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
+                self._create_dummy_tensor((num_tokens, hidden_size), dtype),
             ]
 
             positions = self._create_dummy_tensor((num_tokens, ), jnp.int32)
@@ -702,23 +620,23 @@ class CompilationManager:
                 num_reqs,
             ):
                 target_hidden_states, input_ids, last_token_indices, _ = self.runner.drafter._filter_token_and_prepare_initial_inputs(
-                    self.runner.drafter.state, token_indices, query_start_loc,
-                    seq_lens, input_ids, aux_hidden_states, attention_metadata,
-                    next_token_ids, num_reqs)
+                    token_indices, query_start_loc, seq_lens, input_ids,
+                    aux_hidden_states, attention_metadata, next_token_ids,
+                    num_reqs)
                 return target_hidden_states, input_ids, last_token_indices
 
             input_ids = self._create_dummy_tensor((num_tokens, ), jnp.int32)
             aux_hidden_states = [
                 self._create_dummy_tensor(
-                    (num_tokens, target_hidden_size), jnp.bfloat16,
+                    (num_tokens, hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, target_hidden_size), jnp.bfloat16,
+                    (num_tokens, hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
                 self._create_dummy_tensor(
-                    (num_tokens, target_hidden_size), jnp.bfloat16,
+                    (num_tokens, hidden_size), jnp.bfloat16,
                     NamedSharding(self.runner.mesh, PartitionSpec(None,
                                                                   None))),
             ]
@@ -750,17 +668,17 @@ class CompilationManager:
                 state,
                 kv_caches,
                 input_ids,
-                draft_hidden_states,
+                target_hidden_states,
                 attention_metadata,
             ):
                 kv_caches, hidden_states, _ = self.runner.drafter.model_fn(
-                    state, kv_caches, input_ids, draft_hidden_states,
+                    state, kv_caches, input_ids, target_hidden_states,
                     attention_metadata)
                 self.runner.kv_caches = kv_caches
                 return hidden_states
 
-            draft_hidden_states = self._create_dummy_tensor(
-                (num_tokens, draft_hidden_size), dtype,
+            target_hidden_states = self._create_dummy_tensor(
+                (num_tokens, hidden_size), dtype,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, "model")))
             input_ids = self._create_dummy_tensor(
                 (num_tokens, ), jnp.int32,
@@ -771,7 +689,7 @@ class CompilationManager:
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids,
-                draft_hidden_states,
+                target_hidden_states,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
@@ -781,7 +699,6 @@ class CompilationManager:
             self._run_compilation(
                 "eagle3_prepare_hidden_states_and_input_ids",
                 self.runner.drafter._prepare_hidden_states_and_input_ids,
-                self.runner.drafter.state,
                 aux_hidden_states,
                 query_start_loc,
                 target_token_ids,
@@ -804,19 +721,18 @@ class CompilationManager:
                 self.runner.drafter.state,
                 self.runner.kv_caches,
                 input_ids_loop,
-                draft_hidden_state_loop,
+                target_hidden_state_loop,
                 attention_metadata,
                 num_tokens=num_tokens,
             )
 
             hidden_states = self._create_dummy_tensor(
-                (num_tokens, draft_hidden_size), jnp.bfloat16,
+                (num_tokens, hidden_size), jnp.bfloat16,
                 NamedSharding(self.runner.mesh, PartitionSpec(None, None)))
 
             self._run_compilation(
                 "eagle3_select_inputs_for_loop_speculation",
                 self.runner.drafter._select_inputs_for_loop_speculation,
-                self.runner.drafter.state,
                 positions,
                 hidden_states,
                 hidden_states,
@@ -827,7 +743,6 @@ class CompilationManager:
             self._run_compilation(
                 "eagle3_select_draft_token_ids",
                 self.runner.drafter._select_draft_token_ids,
-                self.runner.drafter.state,
                 hidden_states,
                 last_token_indices,
                 num_tokens=num_tokens,

tpu_inference/runner/kv_cache.py

@@ -9,7 +9,7 @@ from torchax.ops.mappings import t2j_dtype
 
 import tpu_inference.kernels.ragged_paged_attention.v3.kernel as rpa
 import tpu_inference.kernels.ragged_paged_attention.v3.kernel_hd64 as rpa_hd64
-from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.jax.sharding import ShardingAxisName
 from tpu_inference.logger import init_logger
 
 logger = init_logger(__name__)
@@ -82,7 +82,7 @@ def create_kv_caches(
                       ShardingAxisName.ATTN_HEAD))
 
     def _allocate() -> jax.Array:
-        return jnp.zeros(
+        return jnp.empty(
             shape=cache_shape,
             dtype=cache_dtype,
         )