tpu-inference 0.11.1.dev202511270815__py3-none-any.whl → 0.13.0rc2.post7__py3-none-any.whl

tpu_inference/layers/jax/moe/deepseek_v3_moe.py

@@ -1,3 +1,17 @@
+# 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 enum
 from dataclasses import InitVar, dataclass
 from functools import partial
@@ -14,8 +28,8 @@ from qwix._src.providers import ptq
 
 from tpu_inference.layers.jax.base import create_param
 from tpu_inference.layers.jax.layers import FlaxUtils
-from tpu_inference.layers.jax.moe.moe import MoE
-from tpu_inference.models.jax.utils.quantization.quantization_utils import (
+from tpu_inference.layers.jax.moe.moe import CombineExperts, MoE
+from tpu_inference.models.jax.utils.qwix.qwix_utils import (
     manually_quantize_qwix_activation, manually_quantize_qwix_weight)
 
 modeling_flax_utils = FlaxUtils()
@@ -150,6 +164,7 @@ class SparseMoE(MoE):
 
     def __post_init__(self, rngs: nnx.Rngs):
         super().__post_init__(rngs)
+        self.combine_experts = CombineExperts(dtype=self.dtype)
 
         # Derive the expert sharding
         self.expert_axis_name = self.edf_sharding[0]
@@ -331,15 +346,7 @@ class SparseMoE(MoE):
                 processed_tokens, jnp.argsort(sort_indices))
             reshaped_tokens_TXD = unsorted_tokens_tD.reshape(
                 -1, self.num_experts_per_tok, self.hidden_size)
-        with jax.named_scope("combine_weights"):
-            output_TD = jnp.einsum(
-                "TXD,TX -> TD",
-                reshaped_tokens_TXD.astype(jnp.float32),
-                router_weights_TX.astype(jnp.float32),
-                precision='float32',
-            )
-
-        return output_TD.astype(self.dtype)
+        return self.combine_experts(reshaped_tokens_TXD, router_weights_TX)
 
     def _gmm(self, inputs, kernel, group_sizes):
         """Performs Grouped Matrix Multiply."""
@@ -575,11 +582,11 @@ class SparseMoE(MoE):
         )
         out_specs = PartitionSpec(*self.activation_ffw_td)
 
-        mapped_moe_fwd = partial(jax.experimental.shard_map.shard_map,
+        mapped_moe_fwd = partial(jax.shard_map,
                                  mesh=self.mesh,
                                  in_specs=in_specs,
                                  out_specs=out_specs,
-                                 check_rep=False)(
+                                 check_vma=False)(
                                      SparseMoE._distributed_sparse_moe_fwd)
 
         kernel_gating_EDF = self.kernel_gating_EDF.value

tpu_inference/layers/jax/moe/gpt_oss_moe.py

@@ -1,3 +1,17 @@
+# 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 dataclasses import InitVar, dataclass
 
 import jax

tpu_inference/layers/jax/moe/moe.py

@@ -1,3 +1,17 @@
+# 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 dataclasses import InitVar, dataclass
 
 import jax
@@ -12,6 +26,29 @@ from tpu_inference.layers.jax.layers import FlaxUtils
 modeling_flax_utils = FlaxUtils()
 
 
+@dataclass(kw_only=True)
+class CombineExperts(nnx.Module):
+    """Combines expert outputs with router weights.
+
+    Supports `TED,TE -> TD` when passed expert outputs, using float32
+    accumulation for numerical stability, then casting back to the target
+    dtype.
+    """
+
+    dtype: jnp.dtype
+
+    def __call__(self, expert_outputs_TED: Float, weights_TE: Float) -> Float:
+        with jax.named_scope("combine_experts"):
+            output_TD = jnp.einsum(
+                "TED,TE -> TD",
+                expert_outputs_TED.astype(jnp.float32),
+                weights_TE.astype(jnp.float32),
+                precision="float32",
+            )
+
+        return output_TD.astype(self.dtype)
+
+
 @dataclass(kw_only=True)
 class Router(nnx.Module):
     """Router module for Mixture-of-Experts (MoE) layers.
@@ -139,6 +176,9 @@ class MoE(nnx.Module):
                                                  sharding=self.efd_sharding,
                                                  random_init=self.random_init)
 
+        # Shared combine module for combine path
+        self.combine_experts = CombineExperts(dtype=self.dtype)
+
     def _moe_fwd_preapply_router_weights(self, x_TD: jax.Array, weights_TE):
         """Performs the forward pass of the MoE experts with router weights pre-applied to the inputs.
 
@@ -204,6 +244,6 @@ class MoE(nnx.Module):
         with jax.named_scope("down_projection"):
             down_proj_TED = jnp.einsum('TEF,EFD -> TED', fuse_TEF,
                                        self.kernel_down_proj_EFD.value)
-        with jax.named_scope("sum"):
-            output_TD = jnp.einsum('TED,TE -> TD', down_proj_TED, weights)
-        return output_TD.astype(self.dtype)
+        # Combine across experts
+        output_TD = self.combine_experts(down_proj_TED, weights)
+        return output_TD

tpu_inference/layers/jax/pp_utils.py

@@ -0,0 +1,53 @@
+# 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, Protocol
+
+from flax import nnx
+from vllm.distributed import get_pp_group
+from vllm.distributed.utils import get_pp_indices
+
+
+class PPMissingLayer(nnx.Module):
+    """
+    A placeholder layer for missing layers in a pipeline parallel model.
+    """
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __call__(self, *args, **kwargs):
+        """Return the first arg from args or the first value from kwargs."""
+        return args[0] if args else next(iter(kwargs.values()))
+
+
+class LayerFn(Protocol):
+
+    def __call__(self) -> nnx.Module:
+        ...
+
+
+def make_layers(
+    num_hidden_layers: int,
+    layer_fn: LayerFn,
+) -> tuple[int, int, List[nnx.Module]]:
+    start_layer, end_layer = get_pp_indices(num_hidden_layers,
+                                            get_pp_group().rank_in_group,
+                                            get_pp_group().world_size)
+
+    layers = [PPMissingLayer() for _ in range(start_layer)] \
+        + [layer_fn() for _ in range(start_layer, end_layer)] \
+        + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)]
+
+    return start_layer, end_layer, layers

tpu_inference/layers/jax/rope.py

@@ -1,3 +1,17 @@
+# 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 math
 from dataclasses import dataclass, field
 from typing import Optional, Tuple

tpu_inference/layers/jax/rope_interface.py

@@ -1,3 +1,17 @@
+# 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 math
 from typing import Any, Dict
 

tpu_inference/layers/jax/sample/__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/layers/jax/sample/rejection_sampler.py

@@ -1,3 +1,16 @@
+# 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.
 """
 JAX-based rejection sampler for speculative decoding on TPU.
 

tpu_inference/layers/jax/sample/sampling.py

@@ -1,3 +1,17 @@
+# 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 functools
 
 import jax
@@ -28,7 +42,7 @@ def sample(
     if tpu_sampling_metadata.do_sampling:
         # Unshard the logits explicity to avoid latency increase.
         logits = jax.lax.with_sharding_constraint(
-            logits, NamedSharding(mesh, P(ShardingAxisName.ATTN_DATA, None)))
+            logits, NamedSharding(mesh, P(ShardingAxisName.MLP_DATA, None)))
     greedy_sampled = jnp.argmax(logits, axis=-1)
     if not tpu_sampling_metadata.do_sampling:
         return greedy_sampled

tpu_inference/layers/jax/sample/sampling_metadata.py

@@ -1,3 +1,17 @@
+# 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 functools
 from dataclasses import dataclass
 from typing import Optional

tpu_inference/layers/jax/transformer_block.py

@@ -1,3 +1,17 @@
+# 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 dataclasses import dataclass
 from typing import Any, Optional, Tuple
 

tpu_inference/layers/vllm/__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/layers/vllm/attention.py

@@ -15,6 +15,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
 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.logger import init_logger
 from tpu_inference.models.vllm.vllm_model_wrapper_context import \
     get_vllm_model_wrapper_context
@@ -117,10 +118,9 @@ class PallasAttentionBackendImpl(AttentionImpl):
         query, key, value = jax_view(query), jax_view(key), jax_view(value)
         q_scale = k_scale = v_scale = None
         if self.kv_cache_quantized_dtype:
-            key, value = utils.quantize_kv(key, value,
-                                           self.kv_cache_quantized_dtype,
-                                           layer._k_scale_float,
-                                           layer._v_scale_float)
+            key, value = quantize_kv(self.kv_cache_quantized_dtype, key, value,
+                                     layer._k_scale_float,
+                                     layer._v_scale_float)
             # TODO(kyuyeunk): Enable w8a8 when VREG spill issue is resolved.
             # q_scale = layer._q_scale_float
             k_scale = layer._k_scale_float