tpu-inference 0.11.1__py3-none-any.whl

tpu_inference/worker/tpu_worker_jax.py

@@ -0,0 +1,321 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import tempfile
+from typing import Callable, Dict, Optional, Tuple, Union
+
+import jax
+import jax.numpy as jnp
+import jaxtyping
+import vllm.envs as envs
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,
+                                          has_kv_transfer_group)
+from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
+                                             init_distributed_environment)
+from vllm.lora.request import LoRARequest
+from vllm.tasks import SupportedTask
+from vllm.v1.core.kv_cache_utils import get_num_blocks, get_uniform_page_size
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
+
+from tpu_inference import utils
+from tpu_inference.di.abstracts import (AbstractKVCacheConfig,
+                                        AbstractLoRARequest,
+                                        AbstractSchedulerOutput)
+from tpu_inference.di.interfaces import HostInterface
+from tpu_inference.distributed.utils import (get_host_ip, get_kv_transfer_port,
+                                             get_node_id)
+from tpu_inference.logger import init_logger
+from tpu_inference.runner.kv_cache import get_rpa_page_size_bytes
+from tpu_inference.runner.tpu_jax_runner import TPUModelRunner
+from tpu_inference.worker._temporary_vllm_compat import (
+    adapt_kv_cache_config_if_needed, adapt_lora_request_if_needed,
+    adapt_scheduler_output_if_needed)
+from tpu_inference.worker.base import AbstractTpuWorker
+
+logger = init_logger(__name__)
+
+_DTYPE: dict[str, jnp.dtype] = {
+    "bfloat16": jnp.bfloat16,
+    "float": jnp.float32,
+    "float32": jnp.float32,
+}
+
+
+class TPUWorker(AbstractTpuWorker):
+
+    def __init__(self,
+                 vllm_config: VllmConfig,
+                 local_rank: int,
+                 rank: int,
+                 distributed_init_method: str,
+                 is_driver_worker: bool = False,
+                 devices=None,
+                 host_interface: Optional[HostInterface] = None):
+        super().__init__(host_interface)
+
+        # If we use vLLM's model implementation in PyTorch, we should set it
+        # with torch version of the dtype.
+        impl = os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower()
+        if impl != "vllm":  # vllm-pytorch implementation does not need this conversion
+
+            # NOTE(wenlong): because sometimes mm needs to use torch for preprocessing
+            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)
+
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.parallel_config = vllm_config.parallel_config
+        self.cache_config = vllm_config.cache_config
+        self.local_rank = local_rank
+        self.rank = rank
+        self.distributed_init_method = distributed_init_method
+        self.is_driver_worker = is_driver_worker
+        self.devices = devices if devices is not None else []
+
+        if self.model_config.trust_remote_code:
+            # note: lazy import to avoid importing torch before initializing
+            from vllm.utils import init_cached_hf_modules
+
+            init_cached_hf_modules()
+
+        # Delay profiler initialization to the start of the profiling.
+        # This is because in vLLM V1, MP runtime is initialized before the
+        # TPU Worker is initialized. The profiler server needs to start after
+        # MP runtime is initialized.
+        self.profile_dir = None
+        if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
+            # For TPU, we can only have 1 active profiler session for 1 profiler
+            # server. So we only profile on rank0.
+            self.profile_dir = envs.VLLM_TORCH_PROFILER_DIR
+            logger.info("Profiling enabled. Traces will be saved to: %s",
+                        self.profile_dir)
+
+        use_jax_profiler_server = os.getenv("USE_JAX_PROFILER_SERVER", False)
+        # Only one instance of profiler is allowed
+        if use_jax_profiler_server and jax.devices()[0] == self.devices[0]:
+            jax_profiler_server_port = int(
+                os.getenv("JAX_PROFILER_SERVER_PORT", 9999))
+            logger.info(
+                f"Starting JAX profiler server on port {jax_profiler_server_port}"
+            )
+            jax.profiler.start_server(jax_profiler_server_port)
+
+    def initialize_cache(self, num_gpu_blocks: int,
+                         num_cpu_blocks: int) -> None:
+        self.cache_config.num_gpu_blocks = num_gpu_blocks
+        self.cache_config.num_cpu_blocks = num_cpu_blocks
+
+    def init_device(self):
+        if not self.devices:
+            try:
+                device_indexes = self.vllm_config.additional_config[
+                    "sharding"]["sharding_strategy"]["device_indexes"]
+                self.devices = [jax.devices()[i] for i in device_indexes]
+            except KeyError:
+                tp = self.parallel_config.tensor_parallel_size
+                self.devices = jax.devices()[:tp]
+
+        # Initialize the vLLM distribution layer as a single chip environment,
+        # we'll swap the model's parallel modules with TPU SPMD equivalents.
+        with set_current_vllm_config(self.vllm_config):
+            temp_file = tempfile.mkstemp()[1]
+            init_distributed_environment(
+                world_size=1,
+                rank=0,
+                local_rank=0,
+                distributed_init_method=f"file://{temp_file}",
+                backend="gloo",
+            )
+            ensure_model_parallel_initialized(
+                tensor_model_parallel_size=1,
+                pipeline_model_parallel_size=1,
+            )
+        ensure_kv_transfer_initialized(self.vllm_config)
+        self.model_runner = TPUModelRunner(self.vllm_config, self.devices)
+        logger.info(f"Init worker | "
+                    f"rank={self.rank} | "
+                    f"node_id={get_node_id()} | "
+                    f"is_driver_worker={self.is_driver_worker} | "
+                    f"hbm={utils.hbm_usage_gb(self.devices)}GiB")
+
+    def determine_available_memory(self) -> int:
+        gpu_memory_utilization = self.cache_config.gpu_memory_utilization
+        hbm_usage = utils.hbm_usage_bytes(self.devices)
+        total_hbm_limit = total_hbm_used = 0
+        for used, limit in hbm_usage:
+            total_hbm_used += used
+            total_hbm_limit += limit
+
+        total_hbm_limit_cap = total_hbm_limit * gpu_memory_utilization
+        total_hbm_avail = int(total_hbm_limit_cap - total_hbm_used)
+
+        total_hbm_limit_gb = round(total_hbm_limit / utils.GBYTES, 2)
+        total_hbm_limit_cap_gb = round(total_hbm_limit_cap / utils.GBYTES, 2)
+        total_hbm_used_gb = round(total_hbm_used / utils.GBYTES, 2)
+        total_hbm_avail_gb = round(total_hbm_avail / utils.GBYTES, 2)
+
+        logger.info(f"Memory statistics | "
+                    f"{total_hbm_limit_gb=}GiB | "
+                    f"{total_hbm_limit_cap_gb=}GiB | "
+                    f"{total_hbm_used_gb=}GiB | "
+                    f"{total_hbm_avail_gb=}GiB")
+
+        if total_hbm_avail <= 0:
+            raise ValueError(f"{total_hbm_used_gb=}GiB exceeds "
+                             f"{total_hbm_limit_cap_gb=}GiB by "
+                             f"{-total_hbm_avail_gb}GiB. Please consider "
+                             f"increasing --gpu-memory-utilization from "
+                             f"{gpu_memory_utilization} to a larger value.")
+        return total_hbm_avail
+
+    def execute_model(
+        self,
+        scheduler_output: Union[AbstractSchedulerOutput, SchedulerOutput],
+    ) -> Optional[ModelRunnerOutput]:
+        # NOTE: This method intentionally returns a concrete vLLM type, which
+        # violates the pure abstract contract of the base class. This is a
+        # deliberate, temporary compromise for the same reasons outlined in
+        # the `get_kv_cache_spec` method.
+
+        # Adapt the input if necessary (temporary compatibility layer)
+        adapted_scheduler_output = adapt_scheduler_output_if_needed(
+            scheduler_output)
+
+        # Unwrap the adapter to get the concrete vLLM object
+        vllm_scheduler_output = adapted_scheduler_output.vllm_scheduler_output
+        output = self.model_runner.execute_model(vllm_scheduler_output)
+
+        # With a connector, the scheduler expects output from all workers
+        if has_kv_transfer_group():
+            return output
+
+        return output if self.is_driver_worker else None
+
+    def take_draft_token_ids(self) -> Optional[DraftTokenIds]:
+        return self.model_runner.take_draft_token_ids()
+
+    def add_lora(
+        self,
+        lora_request: Union[AbstractLoRARequest, LoRARequest],
+    ) -> bool:
+        # Adapt the input if necessary (temporary compatibility layer)
+        adapted_lora_request = adapt_lora_request_if_needed(lora_request)
+
+        # Unwrap the adapter to get the concrete vLLM object
+        vllm_lora_request = adapted_lora_request.vllm_lora_request  # noqa: F841
+
+        raise NotImplementedError(
+            "LoRA is not supported by the JAX worker yet.")
+
+    def profile(self, is_start: bool = True):
+        if is_start:
+            options = jax.profiler.ProfileOptions()
+            options.python_tracer_level = os.getenv("PYTHON_TRACER_LEVEL", 0)
+            jax.profiler.start_trace(self.profile_dir,
+                                     profiler_options=options)
+        else:
+            jax.profiler.stop_trace()
+
+    def load_model(self) -> None:
+        self.model_runner.load_model()
+
+    def compile_or_warm_up_model(self) -> None:
+        self.model_runner.capture_model()
+        # Reset the seed to ensure that the random state is not affected by
+        # the model initialization and profiling.
+        self.model_runner._init_random()
+
+    def reset_mm_cache(self) -> None:
+        pass
+
+    def get_model(self):
+        return self.model_runner.get_model()
+
+    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
+        return self.model_runner.get_supported_tasks()
+
+    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
+        # NOTE: This method intentionally returns a concrete vLLM type, which
+        # violates the pure abstract contract of the base class. This is a
+        # deliberate, temporary compromise.
+        #
+        # The vLLM executor that calls this method expects the concrete
+        # `vllm.KVCacheSpec` object to perform its own internal logic. If we
+        # returned an abstract adapter, the vLLM code would break.
+        #
+        # The ideal long-term solution is for the vLLM DI container to be
+        # responsible for this translation. When vLLM can be modified, this
+        # method should be changed to return `dict[str, AbstractKVCacheSpec]`,
+        # and the vLLM side should be updated to handle the translation.
+        kv_cache_specs = self.model_runner.get_kv_cache_spec()
+
+        if len(kv_cache_specs) == 0:
+            return kv_cache_specs
+
+        # TODO(kyuyeunk): Instead of checking page_size_bytes here, introduce
+        # feature that allows overriding page_size_bytes of KVCacheSpec.
+        vllm_page_size_bytes = get_uniform_page_size(kv_cache_specs)
+        rpa_page_size_bytes = get_rpa_page_size_bytes(self.model_runner.mesh,
+                                                      kv_cache_specs)
+
+        if vllm_page_size_bytes != rpa_page_size_bytes:
+            logger.info(
+                f"KV cache page size calculated by vLLM "
+                f"({vllm_page_size_bytes} Bytes) does not match with actual "
+                f"page size used by RPA kernel ({rpa_page_size_bytes} Bytes). "
+                f"Recalculating number of KV blocks using actual page size.")
+
+            available_memory = self.determine_available_memory()
+            num_blocks = get_num_blocks(self.vllm_config, len(kv_cache_specs),
+                                        available_memory, rpa_page_size_bytes)
+
+            cache_config = self.vllm_config.cache_config
+            cache_config.num_gpu_blocks_override = num_blocks
+
+        return kv_cache_specs
+
+    def initialize_from_config(
+        self,
+        kv_cache_config: Union[AbstractKVCacheConfig, KVCacheConfig],
+    ) -> None:
+        """Allocate GPU KV cache with the specified kv_cache_config."""
+        adapted_kv_cache_config = adapt_kv_cache_config_if_needed(
+            kv_cache_config)
+        vllm_kv_cache_config = adapted_kv_cache_config.vllm_kv_cache_config
+        self.model_runner.initialize_kv_cache(vllm_kv_cache_config)
+
+    def get_node_kv_ip_port(self) -> tuple[int, str, int]:
+        node_id = get_node_id()
+        ip = get_host_ip()
+        port = get_kv_transfer_port()
+        return (int(node_id), ip, int(port))
+
+    def check_health(self) -> None:
+        # worker will always be healthy as long as it's running.
+        return
+
+    def sync_weights(
+        self,
+        updated_weights: jaxtyping.PyTree,
+        mappings: Dict[str, Tuple[str, Tuple[str]]],
+        transpose_keys: Dict[str, Tuple[int]],
+        reshard_fn: Callable[[jaxtyping.PyTree, jaxtyping.PyTree],
+                             jaxtyping.PyTree] = None
+    ) -> None:
+        """Sync the updated weights to the model runner."""
+        return self.model_runner._sync_weights(updated_weights=updated_weights,
+                                               mappings=mappings,
+                                               transpose_keys=transpose_keys,
+                                               reshard_fn=reshard_fn)
+
+    def shutdown(self) -> None:
+        return

tpu_inference-0.11.1.dist-info/METADATA

@@ -0,0 +1,101 @@
+Metadata-Version: 2.4
+Name: tpu_inference
+Version: 0.11.1
+Author: tpu_inference Contributors
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tpu-info==0.4.0
+Requires-Dist: yapf==0.43.0
+Requires-Dist: pytest
+Requires-Dist: pytest-mock
+Requires-Dist: absl-py
+Requires-Dist: numpy
+Requires-Dist: google-cloud-storage
+Requires-Dist: jax==0.7.2
+Requires-Dist: jaxlib==0.7.2
+Requires-Dist: libtpu==0.0.23
+Requires-Dist: jaxtyping
+Requires-Dist: flax==0.11.1
+Requires-Dist: torchax==0.0.7
+Requires-Dist: qwix==0.1.1
+Requires-Dist: torchvision==0.23.0
+Requires-Dist: pathwaysutils
+Requires-Dist: parameterized
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+
+<p align="center">
+   <!-- This image will ONLY show up in GitHub's dark mode -->
+  <img src="docs/assets/tpu_inference_dark_mode_short.png#gh-dark-mode-only" alt="vLLM TPU" style="width: 86%;">
+    <!-- This image will ONLY show up in GitHub's light mode (and on other platforms) -->
+  <img src="docs/assets/tpu_inference_light_mode_short.png#gh-light-mode-only" alt="vLLM TPU" style="width: 86%;">
+</p>
+
+<p align="center">
+| <a href="https://tpu.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://discuss.vllm.ai/c/hardware-support/google-tpu-support/27"><b>User Forum</b></a> | <a href="https://join.slack.com/share/enQtOTY2OTUxMDIyNjY1OS00M2MxYWQwZjAyMGZjM2MyZjRjNTA0ZjRkNjkzOTRhMzg0NDM2OTlkZDAxOTAzYmJmNzdkNDc4OGZjYTUwMmRh"><b>Developer Slack</b></a> |
+</p>
+
+---
+
+_Upcoming Events_ 🔥
+
+- Join us at the [PyTorch Conference, October 22-23](https://events.linuxfoundation.org/pytorch-conference/) in San Francisco!
+- Join us at [Ray Summit, November 3-5](https://www.anyscale.com/ray-summit/2025) in San Francisco!
+- Join us at [JAX DevLab on November 18th](https://rsvp.withgoogle.com/events/devlab-fall-2025) in Sunnyvale!
+
+_Latest News_ 🔥
+
+- [2025/10] vLLM TPU: A New Unified Backend Supporting PyTorch and JAX on TPU
+
+<details>
+<summary><i>Previous News</i> 🔥</summary>
+
+</details>
+
+---
+## About
+
+vLLM TPU is now powered by `tpu-inference`, an expressive and powerful new hardware plugin unifying JAX and PyTorch under a single lowering path within the vLLM project. The new backend now provides a framework for developers to:
+
+- Push the limits of TPU hardware performance in open source.
+- Provide more flexibility to JAX and PyTorch users by running PyTorch model definitions performantly on TPU without any additional code changes, while also extending native support to JAX.
+- Retain vLLM standardization: keep the same user experience, telemetry, and interface.
+
+## Recommended models and features
+
+Although vLLM TPU’s new unified backend makes out-of-the-box high performance serving possible with any model supported in vLLM, the reality is that we're still in the process of implementing a few core components.
+
+For this reason, we’ve provided a list of recommended [models](https://github.com/vllm-project/tpu-inference/blob/main/support_matrices/model_support_matrix.csv) and [features](https://github.com/vllm-project/tpu-inference/blob/main/support_matrices/feature_support_matrix.csv) that are validated for accuracy and stress-tested for performance.
+
+## Get started
+
+Get started with vLLM on TPUs by following the [quickstart guide](https://github.com/vllm-project/tpu-inference/tree/main/docs/getting_started/quickstart.md).
+
+Visit our [documentation](https://github.com/vllm-project/tpu-inference/tree/main/docs) to learn more.
+
+## Contribute
+
+We're always looking for ways to partner with the community to accelerate vLLM TPU development. If you're interested in contributing to this effort, check out the [Contributing guide](https://github.com/vllm-project/tpu-inference/blob/main/CONTRIBUTING.md) and [Issues](https://github.com/vllm-project/tpu-inference/issues) to start. We recommend filtering Issues on the [**good first issue** tag](https://github.com/vllm-project/tpu-inference/issues?q=is%3Aissue+state%3Aopen+label%3A%22good+first+issue%22) if it's your first time contributing.
+
+## Contact us
+
+- For technical questions and feature requests, open a GitHub [Issue](https://github.com/vllm-project/tpu-inference/issues)
+- For feature requests, please open one on Github [here](https://github.com/vllm-project/tpu-inference/issues/new/choose)
+- For discussing with fellow users, use the [TPU support topic in the vLLM Forum](https://discuss.vllm.ai/c/hardware-support/google-tpu-support/27)
+- For coordinating contributions and development, use the [Developer Slack](https://join.slack.com/share/enQtOTY2OTUxMDIyNjY1OS00M2MxYWQwZjAyMGZjM2MyZjRjNTA0ZjRkNjkzOTRhMzg0NDM2OTlkZDAxOTAzYmJmNzdkNDc4OGZjYTUwMmRh)
+- For collaborations and partnerships, contact us at [vllm-tpu@google.com](mailto:vllm-tpu@google.com)

tpu_inference-0.11.1.dist-info/RECORD

@@ -0,0 +1,168 @@
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/test_base.py,sha256=Ct5WFRMHL7IHEIxk8FrzAvO8m0xFuDpzDBKkAKKAL2Q,7341
+tests/test_quantization.py,sha256=tmHBwpAh1Lz4cSB15fwnvmbA1TZ_zM_I1iP99hhGaEk,34444
+tests/test_tpu_info.py,sha256=ZrwlMsp8ffITkS_b8Q1t_QG-a-WVAd4NUcjHhGibcsI,4670
+tests/test_utils.py,sha256=JFxlYnIddw8t096smLEs_PTycocVVzMGDBgZv5YUlnc,7763
+tests/tpu_backend_test.py,sha256=1_rEUA2XGsDCbZVX5KFOQ00OyTF4YnKRtNmk6ctbKXc,2462
+tests/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/core/test_adapters.py,sha256=HcZHf0GTwfHtW1rhcvAb1A3ezejQpYzzMuJhUvIsDo4,2927
+tests/core/test_core_tpu.py,sha256=n6IPk3VzaFYgm3LDeDp1qoKgRN5ysL7JidFOex2lIDg,22342
+tests/core/test_disagg_executor.py,sha256=QdE2YZs08EyDDCmSjhiXkXqQ9BJTgO6csr_E1xkkfSg,2256
+tests/core/test_disagg_utils.py,sha256=alktTGppaGdg-_un0Amz8Y0IDQz-xNJN0dXG-YApEmY,1955
+tests/core/test_init.py,sha256=NEFI5A9eKGu4rmeJ2iqd0EmhlA3bzbVkXmMi1PV1b9U,1687
+tests/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/kernels/quantized_matmul_kernel_test.py,sha256=od5-zXFjcsc_gWGRDrREL8E_ftymNniQVTzgtkBo_Gc,5679
+tests/kernels/ragged_kv_cache_update_v2_test.py,sha256=6-HjP5CoUG-kcuP8MS-JJVMiBnPRo_zadS3VInnO0D4,10821
+tests/kernels/ragged_paged_attention_kernel_v2_test.py,sha256=pWqo9UYF0tzwgBKO_xYw-TYSPrtAsKcMK5Haj8hFG7I,11340
+tests/kernels/ragged_paged_attention_kernel_v3_test.py,sha256=Hrd8iUkS1pS3rxeTyY53aYRg_ZL_d3NqgBXvOgnigSU,14838
+tests/lora/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/lora/test_lora.py,sha256=nBbwnmvNgTWjqjKXJ0o_n5k7IksXMt5I9SDbpe6IsfM,4168
+tpu_inference/__init__.py,sha256=5hJ_YCx4yQJ3HH2BruqWaOtnYi_IapS9no7l62foFFo,1096
+tpu_inference/backend.py,sha256=V0DveQe4maWGz_hRD4bivwTXIQsANZkEj63_0m7U6nA,2552
+tpu_inference/logger.py,sha256=HQCz7NefmbturuhOC7-3Ixbtcdgoz4g9FHh2RB6o8cc,334
+tpu_inference/tpu_info.py,sha256=9UohshkndR6dZpGWpWXfTD4qvIVdVgHf0yOoSEkLTrw,2276
+tpu_inference/utils.py,sha256=M1JMLFtd_5_za7XAQi2ENY8d7aRC-S7wbpYpLh42tyQ,9533
+tpu_inference/adapters/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/adapters/vllm_adapters.py,sha256=n_iJ-BM4aGlnuf6Qhgye6u-H9dkzZP4SPufjspqw-dk,1412
+tpu_inference/adapters/vllm_config_adapters.py,sha256=V9sNdkKYHJpK-OKaaMYXZZP-IhZW6MOe7fJSwQbJngE,4076
+tpu_inference/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/core/adapters.py,sha256=dTZV95MDUdORJbcdYf1JYNTnNDVmv38V22GE7hiQJmo,4484
+tpu_inference/core/core_tpu.py,sha256=_iGVOp30qEGJX3MTYFRXRTsCdbdtd6vwtBtIeFA0sy8,32609
+tpu_inference/core/disagg_executor.py,sha256=dM0cvw2uS-jDlfG4BtsmGAa6hKyhhQ1H-ZQVvn65Xb0,4597
+tpu_inference/core/disagg_utils.py,sha256=ufWNFWQ5n4YnZpPOtoReHlYo4dlN7AbIqCyqS4an0t4,1572
+tpu_inference/di/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/di/abstracts.py,sha256=pMC-wD9aVoCmD3RXW5A4oHZ9Islu2C6huG9HYmQvxeY,541
+tpu_inference/di/host.py,sha256=FKRd5Xs1BVzbXku8A35tZmJDwPVg_66drdpAdSzJ5VI,2601
+tpu_inference/di/interfaces.py,sha256=LFlfXHWK61apIBb2nEBNjuAsdLLmnxTtrkVGslEKTj8,1524
+tpu_inference/distributed/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/distributed/tpu_connector.py,sha256=l_5l44BVIIClz4hrv5kWtctoUELHtvEXdqfypXlQh3I,28499
+tpu_inference/distributed/utils.py,sha256=8AOevmxJi7o9hLXyAydcYh-WaWGS6-BKJpV8kW6-P6E,1494
+tpu_inference/executors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/executors/ray_distributed_executor.py,sha256=VzAPBVb7c8zwGZFtn1OxnwxQTiZMfLnzeI1P7M69d5k,14888
+tpu_inference/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/experimental/llama3_jax_stashed.py,sha256=YK1oSIfto9ALo-HB45XfSrbq9XgVbE4m2C-9zRwmSzI,10913
+tpu_inference/interfaces/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/interfaces/cache.py,sha256=ZSNvjYpRjmm3RlsqENa3R9oHP0L7W4zv8nhCvpNGJLA,813
+tpu_inference/interfaces/config.py,sha256=f0fJBbp5FAWjfS-gC5UK2ptrn42f-DMLK9y_QB9Hm_U,1211
+tpu_inference/interfaces/config_parts.py,sha256=QuqV6LH8rPXmWuxraVFKmY08aTaZCy62ndM_vK9JkKQ,2105
+tpu_inference/interfaces/engine.py,sha256=Z1Vxmf5tiKTT3LMMpMX73urMqK3Uc4ZvO0UW8oXAsxE,1446
+tpu_inference/interfaces/outputs.py,sha256=ay9DXf_9JnaPc5kPJg3MYbO8frIQTJZxiugP7ow0gUI,580
+tpu_inference/interfaces/params.py,sha256=Cp8MtBj3LW8-4h9J23AJO4wGvG3aOuFIq3YFS-OG8zA,364
+tpu_inference/interfaces/platform.py,sha256=_EVTdilqpXJX2rRdypANuojOhDO0BCkUwekaxXQqDvQ,1833
+tpu_inference/interfaces/request.py,sha256=DRkjdWo5wmkVwQlq9DqpMDPeVPmQd6dfyhN2_k8tezw,950
+tpu_inference/interfaces/scheduler.py,sha256=cFBRkqVNXHrn-08Zvr9B23YTJUzSehy1rE-Fy2V5nvg,816
+tpu_inference/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/collectives/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/collectives/all_gather_matmul.py,sha256=0OYLLjlDmkRYScl7lHRi0o___5I5iMiW1gso-_dWSbc,27255
+tpu_inference/kernels/collectives/all_gather_matmul_tuned_block_sizes.py,sha256=KdaOIzTfIgUR0CcUTA46tpYH-cxPNoJx2cTMEvHx-Ac,1399
+tpu_inference/kernels/collectives/util.py,sha256=LbLD6lOxuszbUsykF89gWQqEJUICCZsfzam3EJDPnFE,1859
+tpu_inference/kernels/flash_attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/flash_attention/kernel.py,sha256=n8gmAFVfchMXlyaSEj8xXJm6AadFt26edQihPRdithY,25897
+tpu_inference/kernels/quantized_matmul/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/quantized_matmul/kernel.py,sha256=4oEVUXgWOeOY-PfySHf-iEuUSd9J7GQk_rDSbxa7CXg,14086
+tpu_inference/kernels/quantized_matmul/tuned_block_sizes.py,sha256=3zhIm73JEE8qOty2_0v3AJlVz13k6qMB5wlXBDyC1EM,35130
+tpu_inference/kernels/quantized_matmul/util.py,sha256=rf6nIiAj9I2cj4LDvtaZGhcLXEc94o2xgMWasnFaREM,1943
+tpu_inference/kernels/ragged_paged_attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/ragged_paged_attention/v2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/ragged_paged_attention/v2/kernel.py,sha256=OiQGAHhyggbp1PeuasPymopFohKOJjGXcpq9p_S8UWA,34940
+tpu_inference/kernels/ragged_paged_attention/v2/ragged_kv_cache_update.py,sha256=vGp2ZWODTbjyG9z2z0Qf_BX-wYHd5bUybnc_DtOz0nI,10995
+tpu_inference/kernels/ragged_paged_attention/v2/tuned_block_sizes.py,sha256=mw80bXBGenroGdrITV0F_EaI2s-Z9KWwqU9WodvJg14,97919
+tpu_inference/kernels/ragged_paged_attention/v3/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/kernels/ragged_paged_attention/v3/kernel.py,sha256=zc-re4Knsdcfvt2oRO5KGD9-dJs0P8GVJ3yGtclHU2A,54740
+tpu_inference/kernels/ragged_paged_attention/v3/tuned_block_sizes.py,sha256=KR2UFpCWjsXCmfMcxxV3yV2DVJp5xcEomOtOKYnSL78,131402
+tpu_inference/kernels/ragged_paged_attention/v3/util.py,sha256=5ij66Rl7YsjTCH1UERP1W-XXC57sL6ZVPQdTLhMtKHQ,1010
+tpu_inference/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/common/attention_metadata.py,sha256=St8ZatbY1D7xQACKJH459jMgp3oTP3AQ36mi9FZdrPU,850
+tpu_inference/layers/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/jax/attention_interface.py,sha256=bXBD8C8RTYTyLJOIGcKd1jH_ZruM0jabLj4n98RIKSA,12003
+tpu_inference/layers/jax/base.py,sha256=Vhts6ZMwNCZ8LbnEXeB0rl3nHdS5hDJWX7HEa7Fl7yE,5775
+tpu_inference/layers/jax/binary_search.py,sha256=ZQi-z1wG6WTcfVQXeTGOZokX4K1DSf9kCzqfrhEU8lk,12320
+tpu_inference/layers/jax/constants.py,sha256=NcYg0zAf3ClfP7YMYdYu_F1GngOzZaIxIAHBZDunKw4,2755
+tpu_inference/layers/jax/layers.py,sha256=yv_lC2tbJuzVL-OaXYooX82Ys8hWZATeH9M78coJ3VI,10633
+tpu_inference/layers/jax/misc.py,sha256=znKv1Nuq_LgYpaIu0qlzUVDgQWnjjG7aqPJGM8kuwcw,566
+tpu_inference/layers/jax/rope.py,sha256=3ZyR06vwliipkynHHrvcK-Q_aRhvQKDYBOqBYr3oWM8,7029
+tpu_inference/layers/jax/rope_interface.py,sha256=X0SruXizlCHGnssFujC1pL07UC4Vsp7-gdBy_Q7JZhI,8375
+tpu_inference/layers/jax/sharding.py,sha256=L0Uh92oLaXFNNQ0qqzNtBD3x3wnTRexQt8GzsCvqH1k,17874
+tpu_inference/layers/jax/transformer_block.py,sha256=MBN4_hYCGq_-eyomGVUqplBZugZ2LBWUFOgM1UtUxFY,2952
+tpu_inference/layers/jax/attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/jax/attention/attention.py,sha256=KsGuQpOu7yUpimIr5XBniHKaa2ohx_Ke2YaCOvAG3jc,9837
+tpu_inference/layers/jax/attention/deepseek_v3_attention.py,sha256=YlagoBMwINv2KRH1dr4oEcH_cQ9QMPB55nO2FQZsWs0,14010
+tpu_inference/layers/jax/attention/llama4_attention.py,sha256=VvUmfBxQEbHf3F2BrcYDUnq5abj7CSDYeRsNx_eVAh0,6162
+tpu_inference/layers/jax/moe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/jax/moe/deepseek_v3_moe.py,sha256=Q6CuwwiZtWYm6iUee1wJoDJrwJE6_bcznTK2HrtXb0M,26089
+tpu_inference/layers/jax/moe/moe.py,sha256=cA8R1rjbBwNEoNlsPWjeIBB9nvaRDwlEdwQTVg6lTpY,8762
+tpu_inference/layers/jax/sample/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/jax/sample/rejection_sampler.py,sha256=IRfVWjkbVXp9Sv1YrGMMh-LYx1AwbY-3FTXEO1-Ue9g,20423
+tpu_inference/layers/jax/sample/sampling.py,sha256=-47SC7AqU4UgyO91zAdYXTgrBfdlQ9I89HFZKwU0eQA,3223
+tpu_inference/layers/jax/sample/sampling_metadata.py,sha256=c3jHNjh1hkFJ5gxGTEk0qBOZnICeY3EELViF5Omp_Nc,2252
+tpu_inference/layers/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/vllm/attention.py,sha256=UVhuNCCrz6jdLNotjGtgaR4CVZ4zNmq5VhsiuOTi6_I,6649
+tpu_inference/layers/vllm/fused_moe.py,sha256=ld_-sIHRdUY2tTTHrzHzCahVxH4P0sZVZrxYBQYSJhE,17455
+tpu_inference/layers/vllm/linear_common.py,sha256=_YlJtbdaYcck_j-gFLos_k0ycktVWxT8Qo57tR2YqJ8,7749
+tpu_inference/layers/vllm/sharding.py,sha256=Ck2OzNiucHtrEutDqPQNteu8MEm6isIkE8U5ziowHgM,5779
+tpu_inference/layers/vllm/quantization/__init__.py,sha256=UGv9cJftrBNoC0pU8SLnTLq3zvqMcolN5YJ6n_J5jf4,1392
+tpu_inference/layers/vllm/quantization/awq.py,sha256=78H4AYgbvLCrW-5bGbn9_WM1J8KnRzVOInfKSW_QmzQ,8476
+tpu_inference/layers/vllm/quantization/common.py,sha256=wm3pge6XMTMsLK7_SSdgBP0PvQzz-1mrqN2I6xMqzrc,4218
+tpu_inference/layers/vllm/quantization/unquantized.py,sha256=QIN6lWfVhN4ikUQlDbD8GhkZcLp1-s1Zi66aqKenmeo,10062
+tpu_inference/layers/vllm/quantization/compressed_tensors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py,sha256=ifC6UsCY0tB6BO7X-PWtw-ikUc5IhcPcLvo0_RFrEsM,5253
+tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py,sha256=6sQvsxiWdi5Vte8V9vrQ2abaqGqWpq-mtzU7lGAo-ac,8759
+tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py,sha256=4y7lYgybpXszpCAtxGFhR8LDEbEoCCeo3DfUSOXxhaQ,5202
+tpu_inference/lora/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/lora/torch_lora_ops.py,sha256=pr3N7DVfkn3ANijUC6dBoiCtIJW4fdJpKdC3zWBUsxE,3121
+tpu_inference/lora/torch_punica_tpu.py,sha256=ZfwWpPhkz4VQyxX9KeClx1hhchglsCWl0xpcGZsuMG0,12522
+tpu_inference/mock/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/mock/vllm_config_utils.py,sha256=FlQshLjoHdgs3C66tYHYbKFUjbk9DhUwY-7HibZk0fI,878
+tpu_inference/mock/vllm_envs.py,sha256=hHtbFOM45T5EB2tEGecMGbJA0qOI9dmNYcjANgtah98,51477
+tpu_inference/mock/vllm_logger.py,sha256=vUGnN5nKT--ZvU15YCzODUM_FGiXKhcrrjDGjeN00RQ,7297
+tpu_inference/mock/vllm_logging_utils.py,sha256=TEUmKj3xHiLzHBnFqAujcxH0t2hBQ04sUaho2RyORnk,486
+tpu_inference/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/models/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/models/common/model_loader.py,sha256=kOwc5Dyn433U0F-qZU1D0_k5USkMTY5Em0_WvQfjIYc,17661
+tpu_inference/models/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/models/jax/deepseek_v3.py,sha256=735PSgqxxrYL9JIsohhUXimjSNYMeNlepfRLrYHZ9us,40038
+tpu_inference/models/jax/llama3.py,sha256=bi-wIgZxR9h_DwoYHczPZXqrcvbzCVwnANuKnak6HcI,13024
+tpu_inference/models/jax/llama4.py,sha256=WMs4gQxbkEZXo7beVJSwPNyZX0AR6prpSE7RGVb9U74,21733
+tpu_inference/models/jax/llama_eagle3.py,sha256=STUkAK6XEA7JM3i_Lx36-t5BhkAGeW_xYiq3zYhHP1A,12297
+tpu_inference/models/jax/phi3.py,sha256=Oz68PE2Z1t8wTed95_w0KMIXfnfV72ZwXugNOdWOV5w,13576
+tpu_inference/models/jax/qwen2.py,sha256=RYb0hMKzPnFOAyhqbztoNlSrFIlRa74fYqSNecA2VOY,13354
+tpu_inference/models/jax/qwen2_5_vl.py,sha256=GrUlM16EWsaGPpSnn1KhjcrAHfeJeC1Z3cVefw0-ynQ,38522
+tpu_inference/models/jax/qwen3.py,sha256=SOL-Pvp56IrMxqXpPf5EFacBI6AJNlqf4Zrr1pkabGw,10994
+tpu_inference/models/jax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/models/jax/utils/file_utils.py,sha256=NOuSC3YFnZpf3CZgYdghbbiNYJt42zgjlEYbOZIVct4,2840
+tpu_inference/models/jax/utils/multi_modal_utils.py,sha256=huW_yfntOJ_3ZXYUN1tJtmeK7EMoOBZExTZQtvfHOdk,6189
+tpu_inference/models/jax/utils/weight_utils.py,sha256=lZIW-39BA6GzdMZ_nr-CapBttLsfEajJvMJo8ykr0B0,19507
+tpu_inference/models/jax/utils/quantization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/models/jax/utils/quantization/quantization_utils.py,sha256=hpzEzosiGi_02bgBXzW-AwZnKEiP_NPiKvpLSIPNjD4,24519
+tpu_inference/models/vllm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/models/vllm/vllm_model_wrapper.py,sha256=CyA9Gk8rmL1_FmIJ0NQcsutkwZn_DBZlzwuib2M2HuI,11141
+tpu_inference/models/vllm/vllm_model_wrapper_context.py,sha256=yxlJHPmRQIAwlb1MmHK3xfXokgIkJ-evNU4PgyoJUdg,1187
+tpu_inference/platforms/__init__.py,sha256=2m4E-nxkBhYZFG23Ni4_AFpZe8xQTimdRltkrNzp7WA,69
+tpu_inference/platforms/tpu_jax.py,sha256=oKQFXjNF6cK2QZT7bqgb50oBwr-FN4VO0VdQXl1TQmE,9941
+tpu_inference/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/runner/block_table_jax.py,sha256=HCjrOMpsWk_x3lW--AOPPUHBHplIzGioMTuHKFxtr6A,4164
+tpu_inference/runner/compilation_manager.py,sha256=16n36Ne4LbmPei8UIAnUMw4TrLcBpe7a5Kvc3oibqcA,30904
+tpu_inference/runner/input_batch_jax.py,sha256=lqFGhZ3w92MPzpiGJ6bNUsQC8X1AP8JpKiWMeXs5tto,18260
+tpu_inference/runner/kv_cache.py,sha256=dU7DRJn0--qgPLV00jCIw4sabSf007mO5kCWnNrNeDI,3952
+tpu_inference/runner/kv_cache_manager.py,sha256=bDkbfpQ41L-n6R-LrseZE85DIuTtu4vbt4mCj1MJa48,21467
+tpu_inference/runner/lora_utils.py,sha256=XFNHPJvZe4e87tbyyKpOY9Vb28M9Rza3HXHNsem7jVg,3872
+tpu_inference/runner/multimodal_manager.py,sha256=2-QQcLuWikP7JgmC3tGovNDYfvikZl3tWyAiX4x8YDc,9283
+tpu_inference/runner/persistent_batch_manager.py,sha256=Zo8w2EdFZTSQtx6DCl57P8kQWkebquXl22RGIX2yqec,11160
+tpu_inference/runner/speculative_decoding_manager.py,sha256=_2oAwo_8e4N-FJXjC9oR-fsO8WjukCdvQhPH4R8B-c4,10274
+tpu_inference/runner/structured_decoding_manager.py,sha256=0SIoa5orxDcx76ziatKJ-GfnTAVIPCPTaMS15nxRR5U,3673
+tpu_inference/runner/tpu_jax_runner.py,sha256=HA-PBThgXv0GHfFxA9ltQr7fQFDw8rwE4SKMsYV0zMI,34285
+tpu_inference/runner/utils.py,sha256=5QcZW8an8EHs_zHKzIGqIf3ltAevusdwgaLLFrB9rc8,17131
+tpu_inference/spec_decode/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/spec_decode/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/spec_decode/jax/eagle3.py,sha256=PgIAJMuEyy61Tz4SQ6QZqB-B4t4-RYDmUIoHDyOHEjA,15204
+tpu_inference/worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tpu_inference/worker/_temporary_vllm_compat.py,sha256=GpF8TuPMDbc0fvIxe7XWEe69FES_F-jJnmcaTgf2dO8,5182
+tpu_inference/worker/base.py,sha256=0Dd3CKk3e7DgvzhfH4M-9-MEQNyYh4zUWSO4tnHFd6s,3140
+tpu_inference/worker/tpu_worker_jax.py,sha256=7b2QVTSbveifm9_BgNnVGwEvh5zPrEi1qiXXTwFFODc,14093
+tpu_inference-0.11.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+tpu_inference-0.11.1.dist-info/METADATA,sha256=uKyRzPptKu13NN6_lYOinPlLYk57ZUFleECr2JDgLrs,5393
+tpu_inference-0.11.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tpu_inference-0.11.1.dist-info/top_level.txt,sha256=gb1hRIQ3DOawUfVzvPL2E__2KPIl9I0vb5r0xcRBGYQ,20
+tpu_inference-0.11.1.dist-info/RECORD,,

tpu_inference-0.11.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+