fix bug after rebase

Author: bzgoogle

tpu_inference/runner/compilation_manager.py

@@ -1,13 +1,14 @@
 import os
 import time
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-import vllm.envs as envs
 from jax.sharding import NamedSharding, PartitionSpec
 
+import vllm.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
@@ -36,9 +37,9 @@ def __init__(self, runner: "TPUModelRunner"):
                               envs.VLLM_XLA_CACHE_PATH)
 
     def _create_dummy_tensor(self,
-                             shape: Tuple[int, ...],
+                             shape: tuple[int, ...],
                              dtype: Any,
-                             sharding: Optional[NamedSharding] = None) -> Any:
+                             sharding: NamedSharding | None = None) -> Any:
         """Helper to create dummy tensors for precompilation."""
         tensor = jnp.ones(shape, dtype=dtype)
         if sharding:
@@ -273,11 +274,11 @@ def _precompile_backbone_with_inputs_embeds(self) -> None:
     def _precompile_select_from_array_helper(
         self,
         name: str,
-        source_paddings: List[int],
-        indices_paddings: List[int],
+        source_paddings: list[int],
+        indices_paddings: list[int],
         hidden_dim: int,
-        input_sharding: Optional[NamedSharding] = None,
-        indices_sharding: Optional[NamedSharding] = None,
+        input_sharding: NamedSharding | None = None,
+        indices_sharding: NamedSharding | None = None,
         only_equal_paddings: bool = False,
         check_should_skip_padding: bool = True,
     ) -> None:
@@ -351,16 +352,18 @@ def _precompile_select_from_array(self) -> None:
                 source_paddings=self.runner.num_logits_paddings,
                 indices_paddings=self.runner.num_reqs_paddings,
                 hidden_dim=vocab_size,
-                input_sharding=NamedSharding(self.runner.mesh,
-                                             PartitionSpec(None, ('model', 'expert')),
+                input_sharding=NamedSharding(
+                    self.runner.mesh, PartitionSpec(None,
+                                                    ('model', 'expert'))),
             )
             self._precompile_select_from_array_helper(
                 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,
-                input_sharding=NamedSharding(self.runner.mesh,
-                                             PartitionSpec(None, ('model', 'expert')),
+                input_sharding=NamedSharding(
+                    self.runner.mesh, PartitionSpec(None,
+                                                    ('model', 'expert'))),
                 only_equal_paddings=True,
             )
 
@@ -392,7 +395,7 @@ def _precompile_sampling(self) -> None:
         for num_reqs in self.runner.num_reqs_paddings:
             logits_sharding = NamedSharding(
                 self.runner.mesh,
-                PartitionSpec(ShardingAxisName.ATTN_DATA, ('model', 'expert'))
+                PartitionSpec(ShardingAxisName.ATTN_DATA, ('model', 'expert')))
             dp_size = self.runner.vllm_config.sharding_config.total_dp_size
             sampling_metadata_sharding = NamedSharding(
                 self.runner.mesh, PartitionSpec(
@@ -481,8 +484,8 @@ def _precompile_rejection_sampler(self) -> None:
         vocab_size = self.runner.model_config.get_vocab_size()
         for num_logits in self.runner.num_logits_paddings:
             for num_reqs in self.runner.num_reqs_paddings:
-                sharding = NamedSharding(self.runner.mesh,
-                                         PartitionSpec(None, ('model', 'expert')))
+                sharding = NamedSharding(
+                    self.runner.mesh, PartitionSpec(None, ('model', 'expert')))
                 target_probs = self._create_dummy_tensor(
                     (num_logits, vocab_size), jnp.bfloat16, sharding)
                 draft_token_ids = self._create_dummy_tensor((num_logits, ),

tpu_inference/runner/kv_cache.py

@@ -1,4 +1,4 @@
-from typing import Any, List
+from typing import Any
 
 import jax
 import jax.numpy as jnp
@@ -46,9 +46,9 @@ def create_kv_caches(
     num_kv_heads: int,
     head_size: int,
     mesh: Mesh,
-    layer_names: List[str],
+    layer_names: list[str],
     cache_dtype: jnp.dtype = DEFAULT_KV_CACHE_DTYPE,
-) -> List[jax.Array]:
+) -> list[jax.Array]:
     """
     Creates a list of KV cache where each array mapps to single attention layer.
 
@@ -78,8 +78,7 @@ def create_kv_caches(
 
     sharding = NamedSharding(
         mesh,
-        PartitionSpec(ShardingAxisName.ATTN_DATA, None,
-                      ('model', 'expert'))
+        PartitionSpec(ShardingAxisName.ATTN_DATA, None, ('model', 'expert')))
 
     def _allocate() -> jax.Array:
         return jnp.empty(

tpu_inference/runner/tpu_runner.py

@@ -2,23 +2,56 @@
 import functools
 import os
 import random
+from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple, cast
+from typing import Any, cast
 
 import jax
 import jax.numpy as jnp
 import jaxtyping
 import numpy as np
 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 j2t_dtype
+
+import vllm.envs as envs
+from tpu_inference import utils as common_utils
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+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 ShardingAxisName, ShardingConfigManager
+from tpu_inference.logger import init_logger
+from tpu_inference.models.common.model_loader import get_model
+from tpu_inference.models.jax.utils.weight_utils import (
+    shard_put,
+    transfer_state_with_mappings,
+)
+from tpu_inference.runner import utils as runner_utils
+from tpu_inference.runner.compilation_manager import CompilationManager
+from tpu_inference.runner.input_batch_jax import CachedRequestState, InputBatch
+from tpu_inference.runner.kv_cache_manager import KVCacheManager
+from tpu_inference.runner.lora_utils import LoraUtils
+from tpu_inference.runner.multimodal_manager import MultiModalManager
+from tpu_inference.runner.persistent_batch_manager import PersistentBatchManager
+from tpu_inference.runner.speculative_decoding_manager import (
+    SpecDecodeMetadata,
+    SpeculativeDecodingManager,
+)
+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
 from vllm.config import VllmConfig
-from vllm.distributed.kv_transfer import (get_kv_transfer_group,
-                                          has_kv_transfer_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
@@ -31,8 +64,7 @@
                              LogprobsTensors, ModelRunnerOutput)
 from vllm.v1.request import Request
 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
-from vllm.v1.worker.kv_connector_model_runner_mixin import \
-    KVConnectorModelRunnerMixin
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 
 from tpu_inference import utils as common_utils
@@ -108,7 +140,7 @@ def __init__(
         next_tokens: jax.Array,
         num_reqs: int,
         discard_sampled_tokens_req_indices: list[int],
-        logits_indices_selector: Optional[List[int]] = None,
+        logits_indices_selector: list[int] | None = None,
     ):
         self._model_runner_output = model_runner_output
         self._next_tokens = next_tokens
@@ -137,7 +169,7 @@ class AsyncPreResults:
     request_seq_lens: list[tuple[int, CachedRequestState, int]]
     discard_sampled_tokens_req_indices: list[int]
     placeholder_req_id_to_index: dict[str, int]
-    logits_indices_selector: Optional[List[int]] = None
+    logits_indices_selector: list[int] | None = None
 
 
 @dataclass
@@ -147,7 +179,7 @@ class ExecuteModelState:
 
     scheduler_output: "VllmSchedulerOutput"
     attn_metadata: AttentionMetadata
-    input_ids: Optional[jax.Array]
+    input_ids: jax.Array | None
     hidden_states: jax.Array
     logits: jax.Array
     aux_hidden_states: Optional[jax.Array]
@@ -552,7 +584,7 @@ def capture_model(self) -> None:
     def execute_model(
         self,
         scheduler_output: "VllmSchedulerOutput",
-        intermediate_tensors: Optional[IntermediateTensors] = None,
+        intermediate_tensors: IntermediateTensors | None = None,
     ) -> ModelRunnerOutput | None:
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called "
@@ -797,7 +829,7 @@ def _sample_from_logits(
         self,
         scheduler_output: "VllmSchedulerOutput",
         attn_metadata: AttentionMetadata,
-        input_ids: Optional[jax.Array],
+        input_ids: jax.Array | None,
         hidden_states: jax.Array,
         logits: jax.Array,
         aux_hidden_states: Optional[jax.Array],
@@ -1617,26 +1649,26 @@ def _get_input_ids_embeds(self, input_ids: jax.Array,
         else:
             return input_ids, None
 
-    def take_draft_token_ids(self) -> Optional[DraftTokenIds]:
+    def take_draft_token_ids(self) -> DraftTokenIds | None:
         return self.speculative_decoding_manager.take_draft_token_ids()
 
     ###### Local disagg utilities ######
 
     def get_kv_cache_for_block_ids(
         self,
-        block_ids: List[int],
-    ) -> List[jax.Array]:
+        block_ids: list[int],
+    ) -> list[jax.Array]:
         return self.kv_cache_manager.get_kv_cache_for_block_ids(block_ids)
 
     def transfer_kv_cache(self,
-                          kv_cache_slices: List[jax.Array]) -> List[jax.Array]:
+                          kv_cache_slices: list[jax.Array]) -> list[jax.Array]:
         return self.kv_cache_manager.transfer_kv_cache(kv_cache_slices)
 
     def insert_request_with_kv_cache(
         self,
         request: "Request",
-        kv_cache_slices: List[jax.Array],
-        block_ids: List[List[int]],
+        kv_cache_slices: list[jax.Array],
+        block_ids: list[list[int]],
     ):
         return self.kv_cache_manager.insert_request_with_kv_cache(
             request, kv_cache_slices, block_ids)
@@ -1646,8 +1678,8 @@ def insert_request_with_kv_cache(
     def _sync_weights(
         self,
         updated_weights: jaxtyping.PyTree,
-        mappings: Dict[str, Tuple[str, Tuple[str]]],
-        transpose_keys: Dict[str, Tuple[int]],
+        mappings: dict[str, tuple[str, tuple[str]]],
+        transpose_keys: dict[str, tuple[int]],
         reshard_fn: Callable[[jaxtyping.PyTree, jaxtyping.PyTree],
                              jaxtyping.PyTree] = None
     ) -> None: