fix bug after rebase

Author: bzgoogle

tpu_inference/runner/compilation_manager.py

@@ -1,18 +1,20 @@
 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.jax.sample.sampling import sample
-from tpu_inference.layers.jax.sample.sampling_metadata import \
-    TPUSupportedSamplingMetadata
+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.utils import device_array
@@ -36,9 +38,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:
@@ -272,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:
@@ -348,16 +350,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,
             )
 
@@ -389,7 +393,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(
@@ -478,8 +482,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_jax_runner.py

@@ -2,66 +2,73 @@
 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.sharding import NamedSharding, PartitionSpec
 from torchax.ops.mappings import j2t_dtype
-from vllm.config import VllmConfig
-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,
-                             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.lora_model_runner_mixin import LoRAModelRunnerMixin
 
+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.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)
+    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.persistent_batch_manager import PersistentBatchManager
 from tpu_inference.runner.speculative_decoding_manager import (
-    SpecDecodeMetadata, SpeculativeDecodingManager)
-from tpu_inference.runner.structured_decoding_manager import \
-    StructuredDecodingManager
+    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 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.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,
+    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.lora_model_runner_mixin import LoRAModelRunnerMixin
 
 logger = init_logger(__name__)
 
@@ -105,7 +112,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
@@ -133,7 +140,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
@@ -143,13 +150,13 @@ 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]
-    spec_decode_metadata: Optional[SpecDecodeMetadata]
-    kv_connector_output: Optional[KVConnectorOutput]
-    logits_indices_selector: Optional[List[int]] = None
+    aux_hidden_states: jax.Array | None
+    spec_decode_metadata: SpecDecodeMetadata | None
+    kv_connector_output: KVConnectorOutput | None
+    logits_indices_selector: list[int] | None = None
 
 
 @functools.partial(jax.jit, donate_argnums=(0, 1, 2))
@@ -192,7 +199,7 @@ class TPUModelRunner(KVConnectorModelRunnerMixin, LoRAModelRunnerMixin):
     def __init__(
         self,
         vllm_config: VllmConfig,
-        devices: List[Any],
+        devices: list[Any],
     ):
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
@@ -266,8 +273,8 @@ def _init_mesh(self) -> None:
             axis_names = ("data", "attn_dp", "expert", "model")
             mesh_shape = (sharding_strategy.model_dp_size,
                           sharding_strategy.attn_dp_size,
-                          sharding_strategy.expert_size,
-                          sharding_strategy.tp_size)
+                          sharding_strategy.expert_size, 2)
+            print(f"DEBUG: {sharding_strategy}")
 
         else:
             axis_names = ("data", "model")
@@ -462,7 +469,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 "
@@ -691,13 +698,13 @@ 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],
-        spec_decode_metadata: Optional[SpecDecodeMetadata],
-        kv_connector_output: Optional[KVConnectorOutput],
-        logits_indices_selector: Optional[List[int]] = None,
+        aux_hidden_states: jax.Array | None,
+        spec_decode_metadata: SpecDecodeMetadata | None,
+        kv_connector_output: KVConnectorOutput | None,
+        logits_indices_selector: list[int] | None = None,
     ) -> ModelRunnerOutput | AsyncTPUModelRunnerOutput:
         padded_num_reqs = runner_utils.get_padded_num_reqs_with_upper_limit(
             self.input_batch.num_reqs, self.max_num_reqs)
@@ -1493,26 +1500,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)
@@ -1522,8 +1529,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: