tpu-inference 0.11.1rc2__py3-none-any.whl → 0.11.1rc3__py3-none-any.whl

tpu_inference/models/common/model_loader.py

@@ -0,0 +1,433 @@
+import functools
+import os
+from typing import Any, Optional
+
+import jax
+import torch
+from flax import nnx
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torchax.ops.mappings import j2t_dtype
+from transformers import PretrainedConfig
+from vllm.config import VllmConfig
+from vllm.utils import supports_kw
+
+from tpu_inference.logger import init_logger
+from tpu_inference.models.jax.utils.quantization.quantization_utils import (
+    apply_qwix_on_abstract_model, apply_qwix_quantization,
+    load_random_weights_into_qwix_abstract_model)
+
+logger = init_logger(__name__)
+
+_MODEL_REGISTRY = {}
+
+
+class UnsupportedArchitectureError(ValueError):
+    """Raised when a model architecture is not supported in the registry."""
+    pass
+
+
+def _get_model_architecture(config: PretrainedConfig) -> nnx.Module:
+    # NOTE: Use inline imports here, otherwise the normal imports
+    # would cause JAX init failure when using multi hosts with Ray.
+
+    from tpu_inference.models.jax.deepseek_v3 import DeepSeekV3
+    from tpu_inference.models.jax.llama4 import Llama4ForCausalLM
+    from tpu_inference.models.jax.llama_eagle3 import EagleLlama3ForCausalLM
+    from tpu_inference.models.jax.phi3 import Phi3ForCausalLM
+    from tpu_inference.models.jax.qwen2 import Qwen2ForCausalLM
+    from tpu_inference.models.jax.qwen2_5_vl import \
+        Qwen2_5_VLForConditionalGeneration
+    from tpu_inference.models.jax.qwen3 import Qwen3ForCausalLM
+
+    if os.getenv("NEW_MODEL_DESIGN", False):
+        from tpu_inference.experimental.llama3_jax_stashed import \
+            LlamaForCausalLM
+    else:
+        from tpu_inference.models.jax.llama3 import LlamaForCausalLM
+
+    _MODEL_REGISTRY["Llama4ForCausalLM"] = Llama4ForCausalLM
+    _MODEL_REGISTRY["DeepseekV3ForCausalLM"] = DeepSeekV3
+    _MODEL_REGISTRY["LlamaForCausalLM"] = LlamaForCausalLM
+    _MODEL_REGISTRY["Qwen2ForCausalLM"] = Qwen2ForCausalLM
+    _MODEL_REGISTRY["Qwen3ForCausalLM"] = Qwen3ForCausalLM
+    _MODEL_REGISTRY[
+        "Qwen2_5_VLForConditionalGeneration"] = Qwen2_5_VLForConditionalGeneration
+    _MODEL_REGISTRY["Phi3ForCausalLM"] = Phi3ForCausalLM
+    _MODEL_REGISTRY["Eagle3LlamaForCausalLM"] = EagleLlama3ForCausalLM
+
+    architectures = getattr(config, "architectures", [])
+    for arch in architectures:
+        if arch in _MODEL_REGISTRY:
+            return _MODEL_REGISTRY[arch]
+    raise UnsupportedArchitectureError(
+        f"Model architectures {architectures} are not supported for now. "
+        f"Supported architectures: {list(_MODEL_REGISTRY.keys())}")
+
+
+def _get_nnx_model(
+    model_class: Any,
+    vllm_config: VllmConfig,
+    rng: jax.Array,
+    mesh: Mesh,
+) -> nnx.Module:
+
+    def create_abstract_model() -> nnx.Module:
+        """
+        Helper class to create an abstract model for `nnx.eval_shape`.
+
+        Returns:
+            An abstract model function.
+        """
+        return model_class(vllm_config, rng, mesh)
+
+    @nnx.jit(donate_argnums=(0, ),
+             static_argnames=('use_qwix_on_abstract_model', ))
+    def create_jit_model(
+            model: nnx.Module,
+            use_qwix_on_abstract_model: bool = False) -> nnx.Module:
+        """
+        Create a jit model.
+
+        Args:
+            model: The model to jit.
+            use_qwix_on_abstract_model: Whether to apply Qwix on the abstract model.
+
+        Returns:
+            The jitted model.
+        """
+        state = nnx.state(model)
+        nnx.update(model, state)
+        if not use_qwix_on_abstract_model:
+            # NOTE: if Qwix is not configured, this will be a no-op
+            model = apply_qwix_quantization(vllm_config,
+                                            model,
+                                            rng,
+                                            mesh,
+                                            apply_to_abstract_model=False)
+        return model
+
+    if os.getenv("JAX_RANDOM_WEIGHTS", False):
+        # Create a sharded model with random inited weights.
+        # TODO: currently Qwen2ForCausalLM is using legacy model implementation
+        # will merge the random init logic when all model are migrated to new model implementation
+
+        # Handle the case where we want to load in random weights to a Qwix-quantized model.  Here, we
+        # need to run an abstract pass for Qwix first and then load in the random weights.
+        if apply_qwix_on_abstract_model(vllm_config):
+            abstract_model_fn = apply_qwix_quantization(
+                vllm_config,
+                create_abstract_model,
+                rng,
+                mesh,
+                apply_to_abstract_model=True)
+
+            model = nnx.eval_shape(abstract_model_fn)
+            quantization_config = vllm_config.model_config.hf_config.quantization_config if hasattr(
+                vllm_config.model_config.hf_config,
+                "quantization_config") else {}
+            load_random_weights_into_qwix_abstract_model(
+                rng, model, mesh, quantization_config)
+            with mesh:
+                jit_model = create_jit_model(model,
+                                             use_qwix_on_abstract_model=True)
+            return jit_model
+
+        @nnx.jit
+        def create_sharded_model():
+            model = model_class(vllm_config, rng, mesh)
+            state = nnx.state(model)
+            pspecs = nnx.get_partition_spec(state)
+            sharded_state = jax.lax.with_sharding_constraint(state, pspecs)
+            nnx.update(model, sharded_state)
+            # NOTE: we don't support quantization for the old Qwen2ForCausalLM implementation
+            return model
+
+        with mesh:
+            jit_model = create_sharded_model()
+            # In this case, we are applying Qwix quantization to the true, concrete model
+            jit_model = apply_qwix_quantization(vllm_config,
+                                                jit_model,
+                                                rng,
+                                                mesh,
+                                                apply_to_abstract_model=False)
+            if hasattr(jit_model, 'initialize_cache'):
+                jit_model.initialize_cache()
+    else:
+        # We first create an abstract model without allocating any weights,
+        # then fill in its weigths during load_weights from HF.
+        # This shows 2 advantages than the normal way:
+        # 1. The model weights will only be allocated once. Otherwise the normal way
+        #    will random-init the model weights first, then load the real weights.
+        #    The two pass weights allocation causes model loading slow.
+        # 2. The model loading won't be OOM. Otherwise the normal way will hold
+        #    a full model weights after random-init, then duplicate a layer during
+        #    the load_weights. This would be easy to OOM if the layer is super large.
+        abstract_model_fn = create_abstract_model
+        # NOTE: only one of the abstract (this) or or concrete Qwix quantization paths should
+        # be taken
+        if should_apply_qwix_on_abstract_model := apply_qwix_on_abstract_model(
+                vllm_config):
+            # NOTE: if Qwix is not configured, this will return `create_abstract_model` and
+            # thus be a no-op
+            abstract_model_fn = apply_qwix_quantization(
+                vllm_config,
+                create_abstract_model,
+                rng,
+                mesh,
+                apply_to_abstract_model=True)
+        model = nnx.eval_shape(abstract_model_fn)
+        # Although the created model can already work, we still need to jit
+        # the model creation again, otherwise the model forward will have
+        # non-trivial overhead in PjitFunction.
+        with mesh:
+            model.load_weights(rng)
+            jit_model = create_jit_model(
+                model,
+                use_qwix_on_abstract_model=should_apply_qwix_on_abstract_model)
+    return jit_model
+
+
+# TODO(pooyam): We need to refactor this. This is returning a bunch of functions that do not work with all models and this is not very easy to see from the code.
+def get_flax_model(
+    vllm_config: VllmConfig,
+    rng: jax.Array,
+    mesh: Mesh,
+    is_draft_model: bool = False,
+) -> nnx.Module:
+    if is_draft_model:
+        model_class = _get_model_architecture(
+            vllm_config.speculative_config.draft_model_config.hf_config)
+    else:
+        model_class = _get_model_architecture(
+            vllm_config.model_config.hf_config)
+    jit_model = _get_nnx_model(model_class, vllm_config, rng, mesh)
+    kv_cache_sharding = NamedSharding(mesh, PartitionSpec(None, None, "model"))
+    hidden_states_sharding = NamedSharding(mesh, PartitionSpec(None,
+                                                               None))  # (T, D)
+
+    # For performance consideration, refer to:
+    # https://flax.readthedocs.io/en/latest/guides/performance.html
+    graphdef, state = nnx.split(jit_model)
+
+    @functools.partial(
+        jax.jit,
+        out_shardings=(
+            kv_cache_sharding,
+            hidden_states_sharding,
+            hidden_states_sharding,  # aux hidden states
+        ),
+        donate_argnums=2,  # 0 is graphdef, 1 is state, 2 is kv_cache
+        static_argnums=6,  #6 is layer_name_to_kvcache_index
+    )
+    def run_model(graphdef, state, *args):
+        model = nnx.merge(graphdef, state)
+        return model(*args)
+
+    logits_sharding = NamedSharding(mesh, PartitionSpec(None, "model"))
+
+    @functools.partial(
+        jax.jit,
+        out_shardings=(logits_sharding),
+    )
+    def run_compute_logits(graphdef, state, *args):
+        model = nnx.merge(graphdef, state)
+        hidden_state, *_ = args
+        return model.compute_logits(hidden_state)
+
+    # Multi-modal support only
+    # This function calculates the image token's embeddings by VIT
+    def run_get_multimodal_embeddings(graphdef, state, image_grid_thw,
+                                      **kwargs):
+        model = nnx.merge(graphdef, state)
+        return model.get_multimodal_embeddings(image_grid_thw, **kwargs)
+
+    # This function will calculates the embeddings of input texts and then merge with the image embeddings
+    @functools.partial(
+        jax.jit,
+        out_shardings=(logits_sharding),
+    )
+    def run_get_input_embeddings(graphdef, state, *args, **kwargs):
+        model = nnx.merge(graphdef, state)
+        return model.get_input_embeddings(*args, **kwargs)
+
+    # For models that want to work with EAGLE-3 speculative decoding
+    @functools.partial(
+        jax.jit,
+        out_shardings=(logits_sharding),
+    )
+    def combine_hidden_states(graphdef, state, hidden_states):
+        model = nnx.merge(graphdef, state)
+        return model.combine_hidden_states(hidden_states)
+
+    model_fn = functools.partial(run_model, graphdef)
+    compute_logits_fn = functools.partial(run_compute_logits, graphdef)
+    get_multimodal_embeddings_fn = functools.partial(
+        run_get_multimodal_embeddings, graphdef)
+    get_input_embeddings_fn = functools.partial(run_get_input_embeddings,
+                                                graphdef)
+    lora_manager, model = None, None
+    combine_hidden_states_fn = functools.partial(combine_hidden_states,
+                                                 graphdef)
+
+    get_mrope_input_positions_fn = None if not hasattr(
+        model_class,
+        "get_mrope_input_positions") else model_class.get_mrope_input_positions
+
+    return model_fn, compute_logits_fn, combine_hidden_states_fn, get_multimodal_embeddings_fn, get_input_embeddings_fn, get_mrope_input_positions_fn, state, lora_manager, model
+
+
+def get_vllm_model(
+    vllm_config: VllmConfig,
+    rng: jax.Array,
+    mesh: Mesh,
+):
+    from tpu_inference.models.vllm.vllm_model_wrapper import VllmModelWrapper
+
+    model = VllmModelWrapper(
+        vllm_config=vllm_config,
+        rng=rng,
+        mesh=mesh,
+    )
+    params, lora_manager = model.load_weights()
+
+    jit_model = model.jit_step_func()
+    compute_logits_fn = model.jit_compute_logits_func()
+    # the model needs to be returned because lora weights are neither torch.nn.parameter nor torch.nn.buffer. After we load the lora weights and set it to the torch.nn.Module, we can shard it and move it to TPU.
+    combine_hidden_states_fn = None
+    return jit_model, compute_logits_fn, combine_hidden_states_fn, None, None, None, params, lora_manager, model
+
+
+def get_model(
+    vllm_config: VllmConfig,
+    rng: jax.Array,
+    mesh: Mesh,
+    is_draft_model: bool = False,
+) -> Any:
+    impl = os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower()
+    logger.info(f"Loading model with MODEL_IMPL_TYPE={impl}")
+
+    if impl == "flax_nnx":
+        try:
+            # Try to load the flax model first
+            return get_flax_model(vllm_config, rng, mesh, is_draft_model)
+        except UnsupportedArchitectureError as e:
+            # Convert the error message to a string to check its contents
+            error_msg = str(e)
+
+            logger.warning(f"Flax model failed with: '{error_msg}'. "
+                           "Falling back to vLLM implementation.")
+            # Fall back to the vLLM model and updating the dtype accordingly
+            vllm_config.model_config.dtype = j2t_dtype(
+                vllm_config.model_config.dtype.dtype)
+            return get_vllm_model(vllm_config, rng, mesh)
+    elif impl == "vllm":
+        return get_vllm_model(vllm_config, rng, mesh)
+    else:
+        raise NotImplementedError("Unsupported MODEL_IMPL_TYPE")
+
+
+def _validate_model_interface(model: Any) -> None:
+    """Validates that the model class has the required methods and signatures.
+
+    A valid model must have:
+    - An __init__ method that accepts a 'vllm_config' keyword argument.
+    - A __call__ method that accepts 'kv_caches', 'input_ids', and
+      'attention_metadata' keyword arguments.
+
+    Args:
+        model: The model class to validate.
+
+    Raises:
+        TypeError: If the model does not meet the interface requirements.
+    """
+    # Check for __init__ with vllm_config
+    model_init = getattr(model, "__init__", None)
+    if not callable(model_init):
+        raise TypeError(
+            f"Model {model.__name__} must have an __init__ method.")
+
+    if not supports_kw(model_init, "vllm_config"):
+        raise TypeError(
+            f"Model {model.__name__} __init__ method must accept a "
+            "'vllm_config' keyword argument.")
+
+    # Check for __call__ with required arguments
+    model_call = getattr(model, "__call__", None)
+    # A class object is always callable (it produces an instance).
+    # We need to check if the class _explicitly_ defines a __call__ method for its
+    # instance, which is different from `type.__call__`.
+    has_defined_call = False
+    if isinstance(model, type):
+        if any("__call__" in C.__dict__ for C in model.__mro__):
+            has_defined_call = True
+    elif callable(model_call):
+        # For an instance, a simple callable check is sufficient.
+        has_defined_call = True
+
+    if not has_defined_call:
+        raise TypeError(f"Model {model.__name__} must have a __call__ method.")
+
+    required_call_args = ("kv_caches", "input_ids", "attention_metadata")
+    missing_args = tuple(arg for arg in required_call_args
+                         if not supports_kw(model_call, arg))
+
+    if missing_args:
+        raise TypeError(
+            f"Model {model.__name__} __call__ method is missing required "
+            f"keyword arguments: {missing_args}")
+
+
+def register_model(arch: str, model: Any) -> None:
+    """
+    Registers a model class for a given architecture name.
+
+    This function registers the model with both the tpu_inference registry
+    and the vLLM registry. For vLLM, it creates a compatible wrapper
+    around the JAX model.
+
+    Args:
+        arch: The name of the architecture (e.g., "LlamaForCausalLM").
+        model: The JAX model class to register (e.g., a flax.nnx.Module).
+    """
+    _validate_model_interface(model)
+
+    # Register with tpu_inference registry for the JAX backend
+    _MODEL_REGISTRY[arch] = model
+
+    # Create a vLLM-compatible wrapper for the JAX model class.
+    # This wrapper inherits from the JAX model and torch.nn.Module
+    # to pass vLLM's type checks. It is not meant to be instantiated
+    # or executed by vLLM's PyTorch backend.
+    def unimplemented_forward(
+        self,
+        input_ids: "torch.Tensor",
+        positions: "torch.Tensor",
+        intermediate_tensors: Optional[Any] = None,
+        inputs_embeds: Optional["torch.Tensor"] = None,
+    ) -> None:
+        raise NotImplementedError(
+            "This is a JAX model and does not implement the PyTorch forward method."
+        )
+
+    # We need a custom __init__ that only calls torch.nn.Module's init,
+    # to avoid triggering JAX logic when vLLM inspects the class.
+    def wrapper_init(self, *args, **kwargs):
+        torch.nn.Module.__init__(self)
+
+    # Dynamically create the wrapper class that is a subclass of both the
+    # JAX model and torch.nn.Module.
+    VllmCompatibleModel = type(
+        f"VllmCompatible{model.__name__}",
+        (model, torch.nn.Module),
+        {
+            "__init__": wrapper_init,
+            "forward": unimplemented_forward,
+            # Prevent vLLM from trying to load weights into this dummy class.
+            "load_weights": lambda self, *args, **kwargs: None,
+        })
+
+    # Register the wrapped model with vLLM's registry.
+    from vllm.model_executor.models.registry import ModelRegistry
+    ModelRegistry.register_model(arch, VllmCompatibleModel)
+    logger.info(
+        f"Registered JAX model {arch} with tpu_inference and vLLM registries.")

{tpu_inference-0.11.1rc2.dist-info → tpu_inference-0.11.1rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tpu_inference
-Version: 0.11.1rc2
+Version: 0.11.1rc3
 Author: tpu_inference Contributors
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
@@ -21,9 +21,9 @@ Requires-Dist: pytest-mock
 Requires-Dist: absl-py
 Requires-Dist: numpy
 Requires-Dist: google-cloud-storage
-Requires-Dist: jax
-Requires-Dist: jaxlib
-Requires-Dist: libtpu
+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
@@ -47,7 +47,7 @@ Dynamic: requires-python
 </p>
 
 <p align="center">
-| <a href="https://github.com/vllm-project/tpu-inference/tree/main/docs"><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> |
+| <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>
 
 ---
@@ -90,7 +90,7 @@ Visit our [documentation](https://github.com/vllm-project/tpu-inference/tree/mai
 
 ## 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%20state%3Aopen%20label%3A%22good%20first%20issue%22) if it's your first time contributing.
+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
 

{tpu_inference-0.11.1rc2.dist-info → tpu_inference-0.11.1rc3.dist-info}/RECORD

@@ -52,6 +52,10 @@ tpu_inference/interfaces/platform.py,sha256=_EVTdilqpXJX2rRdypANuojOhDO0BCkUweka
 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
@@ -67,6 +71,45 @@ tpu_inference/kernels/ragged_paged_attention/v3/__init__.py,sha256=47DEQpj8HBSa-
 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
@@ -76,6 +119,8 @@ tpu_inference/mock/vllm_envs.py,sha256=hHtbFOM45T5EB2tEGecMGbJA0qOI9dmNYcjANgtah
 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
@@ -116,8 +161,8 @@ tpu_inference/worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 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.1rc2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-tpu_inference-0.11.1rc2.dist-info/METADATA,sha256=IY3K7w2UtJtLo8TmAYWCjzLpQvEGsMVkshtWhDGiBqs,5423
-tpu_inference-0.11.1rc2.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-tpu_inference-0.11.1rc2.dist-info/top_level.txt,sha256=gb1hRIQ3DOawUfVzvPL2E__2KPIl9I0vb5r0xcRBGYQ,20
-tpu_inference-0.11.1rc2.dist-info/RECORD,,
+tpu_inference-0.11.1rc3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+tpu_inference-0.11.1rc3.dist-info/METADATA,sha256=eqqOrhIU_Ub_SoR33_ZJWJzszDrpuLZoi0iTBhcpfUo,5396
+tpu_inference-0.11.1rc3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tpu_inference-0.11.1rc3.dist-info/top_level.txt,sha256=gb1hRIQ3DOawUfVzvPL2E__2KPIl9I0vb5r0xcRBGYQ,20
+tpu_inference-0.11.1rc3.dist-info/RECORD,,

{tpu_inference-0.11.1rc2.dist-info → tpu_inference-0.11.1rc3.dist-info}/WHEEL

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