tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tpu_inference/models/jax/qwen3.py

@@ -0,0 +1,318 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jax.sharding import Mesh
+from transformers import Qwen3Config
+from vllm.config import VllmConfig
+
+from tpu_inference import utils
+from tpu_inference.layers.common.attention_interface import attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
+from tpu_inference.layers.jax.rope_interface import apply_rope
+from tpu_inference.logger import init_logger
+from tpu_inference.models.jax.qwen2 import Qwen2DecoderLayer
+from tpu_inference.models.jax.qwen2 import Qwen2MLP as Qwen3MLP
+from tpu_inference.models.jax.qwen2 import Qwen2Model
+from tpu_inference.models.jax.utils.weight_utils import (get_default_maps,
+                                                         load_hf_weights)
+
+logger = init_logger(__name__)
+
+init_fn = nnx.initializers.uniform()
+
+
+class Qwen3Attention(nnx.Module):
+
+    def __init__(self, config: Qwen3Config, dtype: jnp.dtype, rng: nnx.Rngs,
+                 mesh: Mesh, kv_cache_dtype: str):
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.rope_theta = config.rope_theta
+        self.rope_scaling = getattr(config, "rope_scaling", None)
+        self.rms_norm_eps = config.rms_norm_eps
+
+        self.head_dim_original = getattr(config, "head_dim",
+                                         self.hidden_size // self.num_heads)
+        self.head_dim = utils.get_padded_head_dim(self.head_dim_original)
+
+        sharding_size = mesh.shape["model"]
+        self.num_heads = utils.get_padded_num_heads(self.num_heads,
+                                                    sharding_size)
+        self.num_kv_heads = utils.get_padded_num_heads(self.num_kv_heads,
+                                                       sharding_size)
+
+        self.mesh = mesh
+
+        self.q_proj = nnx.Einsum(
+            "TD,DNH->TNH",
+            (self.hidden_size, self.num_heads, self.head_dim),
+            param_dtype=dtype,
+            kernel_init=nnx.with_partitioning(init_fn, (None, "model", None)),
+            rngs=rng,
+        )
+        self.q_norm = nnx.RMSNorm(
+            self.head_dim,
+            epsilon=self.rms_norm_eps,
+            param_dtype=dtype,
+            scale_init=nnx.with_partitioning(init_fn, (None, )),
+            rngs=rng,
+        )
+        self.k_proj = nnx.Einsum(
+            "TD,DKH->TKH",
+            (self.hidden_size, self.num_kv_heads, self.head_dim),
+            param_dtype=dtype,
+            kernel_init=nnx.with_partitioning(init_fn, (None, "model", None)),
+            rngs=rng,
+        )
+        self.k_norm = nnx.RMSNorm(
+            self.head_dim,
+            epsilon=self.rms_norm_eps,
+            param_dtype=dtype,
+            scale_init=nnx.with_partitioning(init_fn, (None, )),
+            rngs=rng,
+        )
+        self.v_proj = nnx.Einsum(
+            "TD,DKH->TKH",
+            (self.hidden_size, self.num_kv_heads, self.head_dim),
+            param_dtype=dtype,
+            kernel_init=nnx.with_partitioning(init_fn, (None, "model", None)),
+            rngs=rng,
+        )
+        self.o_proj = nnx.Einsum(
+            "TNH,NHD->TD",
+            (self.num_heads, self.head_dim, self.hidden_size),
+            param_dtype=dtype,
+            kernel_init=nnx.with_partitioning(init_fn, ("model", None, None)),
+            rngs=rng,
+        )
+
+        self._q_scale = 1.0
+        self._k_scale = 1.0
+        self._v_scale = 1.0
+        self.kv_cache_quantized_dtype = None
+        if kv_cache_dtype != "auto":
+            self.kv_cache_quantized_dtype = utils.get_jax_dtype_from_str_dtype(
+                kv_cache_dtype)
+
+    def __call__(
+        self,
+        kv_cache: Optional[jax.Array],
+        x: jax.Array,
+        attention_metadata: AttentionMetadata,
+    ) -> Tuple[jax.Array, jax.Array]:
+        md = attention_metadata
+        # q: (T, N, H)
+        q = self.q_proj(x)
+        q = self.q_norm(q)
+        q = apply_rope(q, md.input_positions, self.head_dim_original,
+                       self.rope_theta, self.rope_scaling)
+
+        # k: (T, K, H)
+        k = self.k_proj(x)
+        k = self.k_norm(k)
+        k = apply_rope(k, md.input_positions, self.head_dim_original,
+                       self.rope_theta, self.rope_scaling)
+
+        # v: (T, K, H)
+        v = self.v_proj(x)
+        # o: (T, N, H)
+        q_scale = k_scale = v_scale = None
+        if self.kv_cache_quantized_dtype:
+            # TODO(kyuyeunk/jacobplatin): Enable w8a8 when VREG spill issue is resolved.
+            # q_scale = self._q_scale
+            k_scale = self._k_scale
+            v_scale = self._v_scale
+            k, v = quantize_kv(self.kv_cache_quantized_dtype, k, v, k_scale,
+                               v_scale)
+        new_kv_cache, outputs = attention(
+            kv_cache,
+            q,
+            k,
+            v,
+            attention_metadata,
+            self.mesh,
+            self.head_dim_original,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+        )
+        # (T, D)
+        o = self.o_proj(outputs)
+        return new_kv_cache, o
+
+
+class Qwen3DecoderLayer(Qwen2DecoderLayer):
+
+    def __init__(self, config: Qwen3Config, dtype: jnp.dtype, rng: nnx.Rngs,
+                 mesh: Mesh, kv_cache_dtype: str):
+        rms_norm_eps = config.rms_norm_eps
+        hidden_size = config.hidden_size
+
+        self.input_layernorm = nnx.RMSNorm(
+            hidden_size,
+            epsilon=rms_norm_eps,
+            param_dtype=dtype,
+            scale_init=nnx.with_partitioning(init_fn, (None, )),
+            rngs=rng,
+        )
+        self.self_attn = Qwen3Attention(config=config,
+                                        dtype=dtype,
+                                        rng=rng,
+                                        mesh=mesh,
+                                        kv_cache_dtype=kv_cache_dtype)
+        self.post_attention_layernorm = nnx.RMSNorm(
+            hidden_size,
+            epsilon=rms_norm_eps,
+            param_dtype=dtype,
+            scale_init=nnx.with_partitioning(init_fn, (None, )),
+            rngs=rng,
+        )
+        self.mlp = Qwen3MLP(
+            config=config,
+            dtype=dtype,
+            rng=rng,
+        )
+
+
+class Qwen3Model(Qwen2Model):
+
+    def __init__(self, vllm_config: VllmConfig, rng: nnx.Rngs,
+                 mesh: Mesh) -> None:
+        model_config = vllm_config.model_config
+        hf_config = model_config.hf_config
+        vocab_size = model_config.get_vocab_size()
+        dtype = model_config.dtype
+        rms_norm_eps = hf_config.rms_norm_eps
+        hidden_size = hf_config.hidden_size
+
+        self.embed = nnx.Embed(
+            num_embeddings=vocab_size,
+            features=hidden_size,
+            param_dtype=dtype,
+            embedding_init=nnx.with_partitioning(init_fn, ("model", None)),
+            rngs=rng,
+        )
+        self.layers = [
+            Qwen3DecoderLayer(
+                config=hf_config,
+                dtype=dtype,
+                rng=rng,
+                mesh=mesh,
+                # TODO (jacobplatin): we should refactor this to pass a dtype (or config) directly
+                kv_cache_dtype=vllm_config.cache_config.cache_dtype)
+            for _ in range(hf_config.num_hidden_layers)
+        ]
+        self.norm = nnx.RMSNorm(
+            hidden_size,
+            epsilon=rms_norm_eps,
+            param_dtype=dtype,
+            scale_init=nnx.with_partitioning(init_fn, (None, )),
+            rngs=rng,
+        )
+        if model_config.hf_config.tie_word_embeddings:
+            self.lm_head = self.embed.embedding
+        else:
+            self.lm_head = nnx.Param(
+                init_fn(rng.params(), (hidden_size, vocab_size), dtype),
+                sharding=(None, "model"),
+            )
+
+
+class Qwen3ForCausalLM(nnx.Module):
+
+    def __init__(self, vllm_config: VllmConfig, rng_key: jax.Array,
+                 mesh: Mesh) -> None:
+        self.vllm_config = vllm_config
+        self.rng = nnx.Rngs(rng_key)
+        self.mesh = mesh
+
+        self.model = Qwen3Model(
+            vllm_config=vllm_config,
+            rng=self.rng,
+            mesh=mesh,
+        )
+
+    def __call__(
+        self,
+        kv_caches: List[jax.Array],
+        input_ids: jax.Array,
+        attention_metadata: AttentionMetadata,
+        *args,
+    ) -> Tuple[List[jax.Array], jax.Array, List[jax.Array]]:
+        kv_caches, x = self.model(
+            kv_caches,
+            input_ids,
+            attention_metadata,
+        )
+        return kv_caches, x, []
+
+    def compute_logits(self, hidden_states: jax.Array) -> jax.Array:
+        if self.vllm_config.model_config.hf_config.tie_word_embeddings:
+            logits = jnp.dot(hidden_states, self.model.lm_head.value.T)
+        else:
+            logits = jnp.dot(hidden_states, self.model.lm_head.value)
+        return logits
+
+    def load_weights(self, rng_key: jax.Array):
+        # NOTE: Since we are using nnx.eval_shape to init the model,
+        # we have to pass dynamic arrays here for __call__'s usage.
+        self.rng = nnx.Rngs(rng_key)
+
+        # Key: path to a HF layer weight
+        # Value: path to a nnx layer weight
+        mappings = {
+            "model.embed_tokens": "model.embed.embedding",
+            "model.layers.*.input_layernorm":
+            "model.layers.*.input_layernorm.scale",
+            "model.layers.*.mlp.down_proj":
+            "model.layers.*.mlp.down_proj.kernel",
+            "model.layers.*.mlp.gate_proj":
+            "model.layers.*.mlp.gate_proj.kernel",
+            "model.layers.*.mlp.up_proj": "model.layers.*.mlp.up_proj.kernel",
+            "model.layers.*.post_attention_layernorm":
+            "model.layers.*.post_attention_layernorm.scale",
+            "model.layers.*.self_attn.k_norm":
+            "model.layers.*.self_attn.k_norm.scale",
+            "model.layers.*.self_attn.k_proj":
+            "model.layers.*.self_attn.k_proj.kernel",
+            "model.layers.*.self_attn.o_proj":
+            "model.layers.*.self_attn.o_proj.kernel",
+            "model.layers.*.self_attn.q_norm":
+            "model.layers.*.self_attn.q_norm.scale",
+            "model.layers.*.self_attn.q_proj":
+            "model.layers.*.self_attn.q_proj.kernel",
+            "model.layers.*.self_attn.v_proj":
+            "model.layers.*.self_attn.v_proj.kernel",
+            "model.norm": "model.norm.scale",
+        }
+
+        # Add lm_head mapping only if it's not tied to embeddings
+        if not self.vllm_config.model_config.hf_config.tie_word_embeddings:
+            mappings.update({
+                "lm_head": "model.lm_head",
+            })
+
+        metadata_map = get_default_maps(self.vllm_config.model_config,
+                                        self.mesh, mappings)
+        load_hf_weights(vllm_config=self.vllm_config,
+                        model=self,
+                        metadata_map=metadata_map,
+                        mesh=self.mesh)

tpu_inference/models/jax/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tpu_inference/models/jax/utils/file_utils.py

@@ -0,0 +1,110 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import hashlib
+import os
+import shutil
+import subprocess
+from typing import List, Optional
+
+import filelock
+import huggingface_hub.constants
+from huggingface_hub import HfFileSystem, snapshot_download
+from tqdm.auto import tqdm
+
+from tpu_inference.logger import init_logger
+
+logger = init_logger(__name__)
+# Do not set the HuggingFace token here, it should be set via the env `HF_TOKEN`.
+hfs = HfFileSystem()
+
+LOCK_DIR = "/tmp/lock"
+
+#####  Local file utils  #####
+
+
+def run_cmd(cmd: str, *args, **kwargs) -> subprocess.CompletedProcess:
+    return subprocess.run(cmd.split(), *args, **kwargs)
+
+
+def delete_file(path: str) -> None:
+    if os.path.isfile(path):
+        os.remove(path)
+    else:
+        logger.error(f"Trying to delete non-existing file: {path}")
+
+
+def list_files(dir: str, pattern: str = "*") -> List[str]:
+    files = glob.glob(os.path.join(dir, pattern))
+    return files
+
+
+def get_lock(model_name_or_path: str):
+    lock_dir = LOCK_DIR
+    model_name_or_path = str(model_name_or_path)
+    os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
+    model_name = model_name_or_path.replace("/", "-")
+    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
+    # add hash to avoid conflict with old users' lock files
+    lock_file_name = hash_name + model_name + ".lock"
+    # mode 0o666 is required for the filelock to be shared across users
+    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name),
+                             mode=0o666)
+    return lock
+
+
+def get_free_disk_size(path: str = "/") -> int:
+    free_bytes = shutil.disk_usage(path)[2]
+    return free_bytes
+
+
+#####  HuggingFace file utils  #####
+
+
+def is_hf_repo(repo_id: str) -> bool:
+    return hfs.exists(repo_id)
+
+
+def list_hf_repo(repo_id: str, pattern: str = "**") -> List[str]:
+    repo_files = hfs.glob(os.path.join(repo_id, pattern))
+    return repo_files
+
+
+def get_hf_model_weights_size(repo_id: str, weights_format: str) -> int:
+    weights_paths = list_hf_repo(repo_id, weights_format)
+    weights_size = 0
+    for weights_path in weights_paths:
+        weights_size += int(hfs.info(weights_path)["size"])
+    return weights_size
+
+
+class DisabledTqdm(tqdm):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, disable=True)
+
+
+def download_model_weights_from_hf(model_path: str, cache_dir: Optional[str],
+                                   weights_format: str) -> str:
+    with get_lock(model_path):
+        local_dir = snapshot_download(
+            model_path,
+            cache_dir=cache_dir,  # can be specified by HF_HOME or HF_HUB_CACHE
+            allow_patterns=weights_format,
+            tqdm_class=DisabledTqdm,
+            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+        )
+    local_files = list_files(local_dir, weights_format)
+    return local_files

tpu_inference/models/jax/utils/multi_modal_utils.py

@@ -0,0 +1,177 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Union
+
+import jax
+import jax.numpy as jnp
+from typing_extensions import TypeAlias
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+NestedTensors: TypeAlias = Union[list["NestedTensors"], list["jax.Array"],
+                                 "jax.Array", tuple["jax.Array", ...]]
+"""
+Uses a list instead of a tensor if the dimensions of each element do not match.
+"""
+
+MultiModalEmbeddings = Union[list[jax.Array], jax.Array, tuple[jax.Array, ...]]
+"""
+The output embeddings must be one of the following formats:
+
+- A list or tuple of 2D tensors, where each tensor corresponds to
+    each input multimodal data item (e.g, image).
+- A single 3D tensor, with the batch dimension grouping the 2D tensors.
+"""
+
+
+def sanity_check_mm_encoder_outputs(
+    mm_embeddings: MultiModalEmbeddings,
+    expected_num_items: int,
+) -> None:
+    """
+    Perform sanity checks for the result of
+    [`vllm.model_executor.models.SupportsMultiModal.get_multimodal_embeddings`][].
+    """
+    assert isinstance(mm_embeddings, (list, tuple, jax.Array)), (
+        "Expected multimodal embeddings to be a list/tuple of 2D tensors, "
+        f"or a single 3D tensor, but got {type(mm_embeddings)} "
+        "instead. This is most likely due to incorrect implementation "
+        "of the model's `get_multimodal_embeddings` method.")
+
+    assert len(mm_embeddings) == expected_num_items, (
+        "Expected number of multimodal embeddings to match number of "
+        f"input items: {expected_num_items}, but got {len(mm_embeddings)=} "
+        "instead. This is most likely due to incorrect implementation "
+        "of the model's `get_multimodal_embeddings` method.")
+
+    assert all(e.ndim == 2 for e in mm_embeddings), (
+        "Expected multimodal embeddings to be a sequence of 2D tensors, "
+        f"but got tensors with shapes {[e.shape for e in mm_embeddings]} "
+        "instead. This is most likely due to incorrect implementation "
+        "of the model's `get_multimodal_embeddings` method.")
+
+
+def flatten_embeddings(embeddings: NestedTensors) -> jax.Array:
+    """
+    Recursively flattens and concatenates NestedTensors on all but the last
+    dimension.
+    """
+
+    if isinstance(embeddings, jax.Array):
+        return embeddings.reshape(-1, embeddings.shape[-1])
+
+    return jnp.concatenate([flatten_embeddings(t) for t in embeddings], axis=0)
+
+
+def _embedding_count_expression(embeddings: NestedTensors) -> str:
+    """
+    Constructs a debugging representation of the number of embeddings in the
+    NestedTensors.
+    """
+
+    if isinstance(embeddings, jax.Array):
+        return " x ".join([str(dim) for dim in embeddings.shape[:-1]])
+
+    return " + ".join(
+        _embedding_count_expression(inner) for inner in embeddings)
+
+
+def _merge_multimodal_embeddings(
+    inputs_embeds: jax.Array,
+    is_multimodal: jax.Array,
+    multimodal_embeddings: jax.Array,
+) -> jax.Array:
+    """
+    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
+    positions in ``inputs_embeds`` corresponding to placeholder tokens in
+    ``input_ids``.
+        This returns a new array with the updated values.
+    Note:
+        This returns a new array with the updated values.
+    """
+    # The check for matching number of tokens is removed as it is not
+    # JIT-compatible. If the shapes mismatch, JAX will raise an error
+    # during execution anyway. The user-friendly error message is
+    # sacrificed for JIT compatibility.
+
+    # JIT-compatible implementation using jnp.where to avoid
+    # NonConcreteBooleanIndexError.
+    # Create a dummy row to handle indices for non-multimodal tokens.
+    # The content of the dummy row does not matter as it will be masked out.
+    dummy_row = jnp.zeros_like(multimodal_embeddings[0:1])
+
+    # Prepend the dummy row to the flattened embeddings.
+    flattened_padded = jnp.concatenate([dummy_row, multimodal_embeddings],
+                                       axis=0)
+
+    # Create gather indices. For each token in the input sequence, this gives
+    # the index into `flattened_padded`.
+    # For non-multimodal tokens, the index will be 0 (pointing to the dummy
+    # row). For the k-th multimodal token, the index will be k.
+    gather_indices = jnp.cumsum(is_multimodal)
+
+    # Gather the embeddings to be placed.
+    update_values = flattened_padded[gather_indices]
+
+    # Use jnp.where to select between original and new embeddings.
+    condition = jnp.expand_dims(is_multimodal, axis=-1)
+    return jnp.where(condition, update_values, inputs_embeds)
+
+
+def merge_multimodal_embeddings(
+    input_ids: jax.Array,
+    inputs_embeds: jax.Array,
+    multimodal_embeddings: jax.Array,
+    placeholder_token_id: Union[int, list[int]],
+) -> jax.Array:
+    """
+    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
+    positions in ``inputs_embeds`` corresponding to placeholder tokens in
+    ``input_ids``.
+
+    ``placeholder_token_id`` can be a list of token ids (e.g, token ids
+    of img_start, img_break, and img_end tokens) when needed: This means
+    the order of these tokens in the ``input_ids`` MUST MATCH the order of
+    their embeddings in ``multimodal_embeddings`` since we need to
+    slice-merge instead of individually scattering.
+
+    For example, if input_ids is "TTTTTSIIIBIIIBIIIETTT", where
+    - T is text token
+    - S is image start token
+    - I is image embedding token
+    - B is image break token
+    - E is image end token.
+
+    Then the image embeddings (that correspond to I's) from vision encoder
+    must be padded with embeddings of S, B, and E in the same order of
+    input_ids for a correct embedding merge.
+
+        This returns a new array with the updated values.
+    """
+    if isinstance(placeholder_token_id, list):
+        placeholder_token_id = jnp.array(placeholder_token_id)
+
+        return _merge_multimodal_embeddings(
+            inputs_embeds,
+            jnp.isin(input_ids, placeholder_token_id),
+            multimodal_embeddings,
+        )
+
+    return _merge_multimodal_embeddings(
+        inputs_embeds,
+        (input_ids == placeholder_token_id),
+        multimodal_embeddings,
+    )

tpu_inference/models/jax/utils/qwix/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.