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

tpu_inference/layers/jax/attention/deepseek_v3_attention.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 InitVar, dataclass
 from typing import Any, Tuple
@@ -6,14 +20,18 @@ import jax
 import jax.numpy as jnp
 from flax import nnx
 from flax.typing import Sharding
-from jax.experimental import shard_map
 from jax.sharding import Mesh
 from jax.sharding import PartitionSpec as P
 
 from tpu_inference import utils
+from tpu_inference.kernels.mla.v1.kernel import mla_ragged_paged_attention
 from tpu_inference.kernels.ragged_paged_attention.v3.kernel import \
     ragged_paged_attention
+from tpu_inference.kernels.ragged_paged_attention.v3.tuned_block_sizes import \
+    get_tuned_block_sizes
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
+from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.layers.jax.base import create_param
 from tpu_inference.layers.jax.layers import RMSNorm
 from tpu_inference.layers.jax.rope import DeepseekScalingRotaryEmbedding
@@ -48,8 +66,9 @@ class MLA(nnx.Module):
     rms_norm_eps: float
 
     # Sharding attributes
-    nhd_sharding: Sharding = ()
+    rd_sharding: Sharding = ()
     q_da_sharding: Sharding = ()
+    ap_sharding: Sharding = ()
     anh_sharding: Sharding = ()
     kv_da_sharding: Sharding = ()
 
@@ -66,6 +85,7 @@ class MLA(nnx.Module):
     rope_input_ordering: str = "split"
     quant: Any | None = None
     rope_mscale_all_dim: float = 1.0
+    use_mla_kernel: bool = False
 
     rngs: InitVar[nnx.Rngs]
 
@@ -77,10 +97,10 @@ class MLA(nnx.Module):
         self.N = self.num_attention_heads
         self.K = self.num_key_value_heads
         self.D = self.hidden_size
-
         self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
 
-        assert self.N == self.K, "N and K must be equal for MLA"
+        if not self.use_mla_kernel:
+            assert self.N == self.K, "N and K must be equal for MLA"
 
         if self.rope_scaling["factor"] <= 1.0:
             yarn_mscale = 1.0
@@ -108,10 +128,10 @@ class MLA(nnx.Module):
                                                   self.q_da_sharding,
                                                   self.dtype,
                                                   random_init=self.random_init)
-        self.kernel_q_up_proj_ANH = create_param(
+        self.kernel_q_up_proj_AP = create_param(
             rngs,
-            (self.q_lora_rank, self.N, self.qk_head_dim),
-            self.anh_sharding,
+            (self.q_lora_rank, self.N * self.qk_head_dim),
+            self.ap_sharding,
             self.dtype,
             random_init=self.random_init,
         )
@@ -122,17 +142,38 @@ class MLA(nnx.Module):
             self.dtype,
             random_init=self.random_init,
         )
-        self.kernel_kv_up_proj_ANH = create_param(
-            rngs,
-            (self.kv_lora_rank, self.N,
-             self.qk_nope_head_dim + self.v_head_dim),
-            self.anh_sharding,
-            self.dtype,
-            random_init=self.random_init,
-        )
-        self.kernel_o_proj_NHD = create_param(
-            rngs, (self.N, self.v_head_dim, self.D),
-            self.nhd_sharding,
+        # NOTE (jacobplatin): we are keeping these variables as 3D because
+        # we would need to reshape them before the below projection,
+        # which caused issues as Qwix wasn't quantizing it correctly
+        # on the abstract pass
+        if self.use_mla_kernel:
+            self.kernel_k_up_proj_ANH = create_param(
+                rngs,
+                (self.kv_lora_rank, self.N, self.qk_nope_head_dim),
+                self.anh_sharding,
+                self.dtype,
+                random_init=self.random_init,
+            )
+            self.kernel_v_up_proj_ANH = create_param(
+                rngs,
+                (self.kv_lora_rank, self.N, self.v_head_dim),
+                self.anh_sharding,
+                self.dtype,
+                random_init=self.random_init,
+            )
+        else:
+            self.kernel_kv_up_proj_AL = create_param(
+                rngs,
+                (self.kv_lora_rank, self.N *
+                 (self.qk_nope_head_dim + self.v_head_dim)),
+                self.
+                ap_sharding,  # NOTE: we use the same sharding for kv_up_proj_AL and kernel_q_up_proj_AP
+                self.dtype,
+                random_init=self.random_init,
+            )
+        self.kernel_o_proj_RD = create_param(
+            rngs, (self.N * self.v_head_dim, self.D),
+            self.rd_sharding,
             self.dtype,
             random_init=self.random_init)
         self.q_rms_norm = RMSNorm(
@@ -188,17 +229,24 @@ class MLA(nnx.Module):
             q_TA = jnp.einsum("TD,DA -> TA", x_q_TD,
                               self.kernel_q_down_proj_DA.value)
             q_TA = self.q_rms_norm(q_TA)
-            # Query up projection.
-            q_TNH = jnp.einsum("TA,ANH -> TNH", q_TA,
-                               self.kernel_q_up_proj_ANH.value)
+            # Query up projection, then reshape to TNH.
+            q_TP = jnp.einsum("TA,AP -> TP", q_TA,
+                              self.kernel_q_up_proj_AP.value)
+            q_TNH = q_TP.reshape(q_TA.shape[0], self.N, self.qk_head_dim)
             # Split the query into nope and rope.
             q_nope_TNH = q_TNH[..., :self.qk_nope_head_dim]
             q_rope_TNH = q_TNH[..., self.qk_nope_head_dim:]
             q_rope_TNH = self.rope.apply_rope(md.input_positions, q_rope_TNH)
-            # Concatenate the nope and rope queries.
-            q_TNH = jnp.concatenate([q_nope_TNH, q_rope_TNH], axis=-1)
-            # Multiple the query by scaling factor
-            q_TNH = nnx.with_sharding_constraint(q_TNH, self.query_tnh)
+            if self.use_mla_kernel:
+                # Absorb the k up-projection matrix into q
+                q_TNA = jnp.einsum("TNH,ANH -> TNA", q_nope_TNH,
+                                   self.kernel_k_up_proj_ANH.value)
+                q_TNA = nnx.with_sharding_constraint(q_TNA, self.query_tnh)
+            else:
+                # Concatenate the nope and rope queries.
+                q_TNH = jnp.concatenate([q_nope_TNH, q_rope_TNH], axis=-1)
+                # Multiply the query by scaling factor
+                q_TNH = nnx.with_sharding_constraint(q_TNH, self.query_tnh)
 
         with jax.named_scope("kv_proj"):
             # KV down projection.
@@ -209,21 +257,30 @@ class MLA(nnx.Module):
             # Reshape k_rope_BSH to include head dimension for RoPE application
             k_rope_SNH = k_rope_SH[..., None, :]
             k_rope_SNH = self.rope.apply_rope(md.input_positions, k_rope_SNH)
-            k_rope_SNH = jnp.broadcast_to(
-                k_rope_SNH,
-                (k_rope_SNH.shape[0], self.N, self.qk_rope_head_dim))
+            assert k_rope_SNH.shape[1] == 1
+            k_rope_SH = k_rope_SNH[:, 0, :]
+
             kv_SA = kv_SA[..., :self.kv_lora_rank]
             kv_SA = self.kv_rms_norm(kv_SA)
-            # KV up projection.
-            kv_nope_SNH = jnp.einsum("SA,ANH -> SNH", kv_SA,
-                                     self.kernel_kv_up_proj_ANH.value)
-            # Split the latent kv vector into k nope vector and v vector.
-            k_nope_SNH = kv_nope_SNH[..., :self.qk_nope_head_dim]
-            v_SNH = kv_nope_SNH[..., self.qk_nope_head_dim:]
-            # Concatenate the key vector.
-            k_SNH = jnp.concatenate([k_nope_SNH, k_rope_SNH], axis=-1)
-            k_SNH = nnx.with_sharding_constraint(k_SNH, self.keyvalue_skh)
-            v_SNH = nnx.with_sharding_constraint(v_SNH, self.keyvalue_skh)
+            kv_SA = nnx.with_sharding_constraint(kv_SA, self.keyvalue_skh)
+
+            if not self.use_mla_kernel:
+                k_rope_SNH = jnp.broadcast_to(
+                    k_rope_SNH,
+                    (k_rope_SNH.shape[0], self.N, self.qk_rope_head_dim))
+                # KV up projection, then reshape to SN(Hk+Hv).
+                kv_SL = jnp.einsum("SA,AL -> SL", kv_SA,
+                                   self.kernel_kv_up_proj_AL.value)
+                kv_nope_SNH = kv_SL.reshape(
+                    kv_SA.shape[0], self.N,
+                    self.qk_nope_head_dim + self.v_head_dim)
+                # Split the latent kv vector into k nope vector and v vector.
+                k_nope_SNH = kv_nope_SNH[..., :self.qk_nope_head_dim]
+                v_SNH = kv_nope_SNH[..., self.qk_nope_head_dim:]
+                # Concatenate the key vector.
+                k_SNH = jnp.concatenate([k_nope_SNH, k_rope_SNH], axis=-1)
+                k_SNH = nnx.with_sharding_constraint(k_SNH, self.keyvalue_skh)
+                v_SNH = nnx.with_sharding_constraint(v_SNH, self.keyvalue_skh)
 
         with jax.named_scope("attn_op"):
             # TODO(wenxindongwork): K and V have different head dimension,
@@ -234,44 +291,67 @@ class MLA(nnx.Module):
             # q, k, v head dimension to be multiple of 128. For now, we will
             # pad the q, k, v dimension to multiple of 128.
             # We should update the MLA kv cache implementation in the future.
-            multiple_of_128 = ((self.qk_head_dim - 1) // 128 + 1) * 128
-            q_TNH = jnp.pad(q_TNH, ((0, 0), (0, 0),
-                                    (0, multiple_of_128 - self.qk_head_dim)))
-            k_SNH = jnp.pad(k_SNH, ((0, 0), (0, 0),
-                                    (0, multiple_of_128 - self.qk_head_dim)))
-            v_SNH = jnp.pad(v_SNH, ((0, 0), (0, 0),
-                                    (0, multiple_of_128 - self.v_head_dim)))
+            if not self.use_mla_kernel:  # MLA kernel handles padding
+                multiple_of_128 = ((self.qk_head_dim - 1) // 128 + 1) * 128
+                q_TNH = jnp.pad(q_TNH,
+                                ((0, 0), (0, 0),
+                                 (0, multiple_of_128 - self.qk_head_dim)))
+                k_SNH = jnp.pad(k_SNH,
+                                ((0, 0), (0, 0),
+                                 (0, multiple_of_128 - self.qk_head_dim)))
+                v_SNH = jnp.pad(v_SNH,
+                                ((0, 0), (0, 0),
+                                 (0, multiple_of_128 - self.v_head_dim)))
+
             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_SNH, v_SNH = utils.quantize_kv(k_SNH, v_SNH,
-                                                 self.kv_cache_quantized_dtype,
-                                                 k_scale, v_scale)
-            new_kv_cache, outputs_TNH = self.attention(
-                is_prefill,
-                kv_cache,
-                q_TNH,
-                k_SNH,
-                v_SNH,
-                attention_metadata,
-                self.mesh,
-                q_scale,
-                k_scale,
-                v_scale,
-            )
-            # TODO(wenxindongwork): For now, unpad the outputs_TNH to match the v_head_dim.
-            # We shall add the MLA kv cache implementation in the future.
-            outputs_TNH = outputs_TNH[..., :self.v_head_dim]
 
-        with jax.named_scope("o_proj"):
-            o_TD = jnp.einsum("TNH,NHD -> TD", outputs_TNH,
-                              self.kernel_o_proj_NHD.value)
-            o_TD = nnx.with_sharding_constraint(
-                o_TD, self.activation_attention_out_td)
-        return new_kv_cache, o_TD
+            # TODO(gpolovets): MLA does not currently support quantized KV!
+            if not self.use_mla_kernel:
+                if self.kv_cache_quantized_dtype:
+                    # TODO(kyuyeunk/jacobplatin): Enable w8a8 when VREG spill issue is resolved.
+                    k_scale = self._k_scale
+                    v_scale = self._v_scale
+                    k_SNH, v_SNH = quantize_kv(self.kv_cache_quantized_dtype,
+                                               k_SNH, v_SNH, k_scale, v_scale)
+
+                new_kv_cache, outputs_TNH = self.attention(
+                    is_prefill,
+                    kv_cache,
+                    q_TNH,
+                    k_SNH,
+                    v_SNH,
+                    attention_metadata,
+                    self.mesh,
+                    q_scale,
+                    k_scale,
+                    v_scale,
+                )
+                # TODO(wenxindongwork): For now, unpad the outputs_TNH to match the v_head_dim.
+                # We shall add the MLA kv cache implementation in the future.
+                outputs_TNH = outputs_TNH[..., :self.v_head_dim]
+
+            else:
+                new_kv_cache, outputs_TNA = self.mla_attention(
+                    kv_cache,
+                    q_TNA,
+                    q_rope_TNH,
+                    kv_SA,
+                    k_rope_SH,
+                    attention_metadata,
+                    self.mesh,
+                )
+                outputs_TNH = jnp.einsum("TNA,ANH -> TNH", outputs_TNA,
+                                         self.kernel_v_up_proj_ANH.value)
+
+            with jax.named_scope("o_proj"):
+                outputs_TNH = nnx.with_sharding_constraint(
+                    outputs_TNH, self.activation_attention_out_td)
+                outputs_TR = outputs_TNH.reshape(outputs_TNH.shape[0],
+                                                 self.N * self.v_head_dim)
+                o_TD = jnp.einsum("TR,RD -> TD", outputs_TR,
+                                  self.kernel_o_proj_RD.value)
+
+            return new_kv_cache, o_TD
 
     def attention(
         self,
@@ -326,21 +406,22 @@ class MLA(nnx.Module):
         out_specs = (self.attn_o_tnh, P(None, None, "model"))
 
         def _ragged_paged_attention(*args):
-            return ragged_paged_attention(
+            outputs = ragged_paged_attention(
                 *args,
                 sm_scale=self.scale,
                 q_scale=q_scale,
                 k_scale=k_scale,
                 v_scale=v_scale,
             )
+            return outputs
 
         output_TNH, kv_cache = jax.jit(
-            shard_map.shard_map(
+            jax.shard_map(
                 _ragged_paged_attention,
                 mesh=mesh,
                 in_specs=in_specs,
                 out_specs=out_specs,
-                check_rep=False,
+                check_vma=False,
             ))(
                 q_TNH,
                 k_SKH,
@@ -352,3 +433,115 @@ class MLA(nnx.Module):
                 md.request_distribution,
             )
         return kv_cache, output_TNH
+
+    def mla_attention(
+        self,
+        kv_cache: KVCache,
+        q_TNA: jax.Array,
+        q_rope_TNH: jax.Array,
+        k_SA: jax.Array,
+        k_rope_SH: jax.Array,
+        attention_metadata: AttentionMetadata,
+        mesh: Mesh,
+    ) -> Tuple[KVCache, jax.Array]:
+        """Performs scaled dot-product attention and updates the KV cache.
+
+        This function handles the core attention logic, which varies between
+        prefill and generation modes. In prefill, it computes self-attention
+        over the input sequence with a causal mask. In generation, it attends
+        to the full history of keys and values stored in the cache.
+
+        Args:
+            kv_cache: The key-value cache to be updated and used.
+            q_TNA: Query tensor of shape `(query_seq, num_attention_heads, lkv_dim)`.
+            q_rope_TNH: Query rope tensor of shape `(query_seq, num_attention_heads, rope_dim)`.
+            k_SA: Key tensor of shape `(kv_seq, lkv_dim)`.
+            k_rope_SH: Key rope tensor of shape `(kv_seq, rope_dim)`.
+            attention_metadata: Metadata containing sequence lengths.
+            mesh: The JAX device mesh (unused in this specific function but
+                kept for potential future use or API consistency).
+            q_scale: Quantization scale for q.
+            k_scale: Quantization scale for k.
+            v_scale: Quantization scale for v.
+
+        Returns:
+            A tuple containing:
+                - The updated KV cache.
+                - The attention output tensor of shape
+                  `(seq, num_q_heads, head_dim)`.
+        """
+        md = attention_metadata
+        in_specs = (
+            self.query_tnh,  # q
+            self.query_tnh,  # q_rope
+            self.keyvalue_skh,  # k
+            self.keyvalue_skh,  # k_rope
+            P(ShardingAxisName.MLP_TENSOR),  # kv_cache
+            P(ShardingAxisName.ATTN_DATA),  # md.seq_lens: Replicated
+            P(ShardingAxisName.ATTN_DATA),  # page_indices_flat: Replicated
+            P(ShardingAxisName.ATTN_DATA),  # query_start_loc: Replicated
+            P(ShardingAxisName.ATTN_DATA),  # distribution: Replicated
+        )
+
+        out_specs = (self.attn_o_tnh, P(ShardingAxisName.MLP_TENSOR))
+
+        def _mla_ragged_paged_attention(q, q_rope, k, k_rope, kv_cache, *args):
+
+            def _initialize_block_sizes():
+                # Set reasonable starting estimates for block sizes. (TODO(gpolovets): update this to use tuned sizes)
+                # Referring to get_tuned_block_sizes() in kernels/ragged_paged_attention/v3/tuned_block_sizes.py: 'TPU v7'/128/'q_bfloat16_kv_bfloat16/q_head-128_kv_head-1_head-128'/4096
+                max_num_tokens = q.shape[0]
+                max_num_seqs = md.seq_lens.shape[0]
+                num_page_indices = md.block_tables.shape[0]
+                assert num_page_indices % max_num_seqs == 0
+                pages_per_seq = num_page_indices // max_num_seqs
+                # num_kv_pages_per_block = min(pages_per_seq, 16)
+                bkv_p, bq_sz = get_tuned_block_sizes(
+                    q.dtype,
+                    kv_cache.dtype,
+                    self.num_attention_heads,
+                    1,
+                    self.qk_nope_head_dim,
+                    kv_cache.shape[1],  # page size
+                    max_num_tokens,
+                    pages_per_seq,
+                )
+                num_kv_pages_per_block = min(min(pages_per_seq, bkv_p), 4)
+                num_queries_per_block = min(min(max_num_tokens, bq_sz),
+                                            4)  # OOMS at 8
+                return num_kv_pages_per_block, num_queries_per_block
+
+            num_kv_pages_per_block, num_queries_per_block = _initialize_block_sizes(
+            )
+            output, kv_cache = mla_ragged_paged_attention(
+                q,
+                q_rope,
+                k,
+                k_rope,
+                kv_cache,
+                *args,
+                sm_scale=self.scale,
+                num_kv_pages_per_block=num_kv_pages_per_block,
+                num_queries_per_block=num_queries_per_block)
+
+            return kv_cache, output
+
+        kv_cache, output_TNH = jax.jit(
+            jax.shard_map(
+                _mla_ragged_paged_attention,
+                mesh=mesh,
+                in_specs=in_specs,
+                out_specs=out_specs,
+                check_vma=False,
+            ), )(
+                q_TNA,
+                q_rope_TNH,
+                k_SA,
+                k_rope_SH,
+                kv_cache,
+                md.seq_lens,
+                md.block_tables,
+                md.query_start_loc,
+                md.request_distribution,
+            )
+        return kv_cache, output_TNH

tpu_inference/layers/jax/attention/gpt_oss_attention.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
 from typing import Tuple
 
@@ -5,7 +19,6 @@ import jax
 import jax.numpy as jnp
 from flax import nnx
 from flax.typing import Sharding
-from jax.experimental import shard_map
 from jax.sharding import Mesh
 from jax.sharding import PartitionSpec as P
 
@@ -13,6 +26,7 @@ from tpu_inference import utils
 from tpu_inference.kernels.ragged_paged_attention.v3.kernel_hd64 import \
     ragged_paged_attention_hd64
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
 from tpu_inference.layers.jax.base import create_param
 from tpu_inference.layers.jax.rope import GptOssRotaryEmbedding
 
@@ -158,17 +172,17 @@ class GptOssAttention(nnx.Module):
     ) -> Tuple[KVCache, jax.Array]:
         """Performs scaled dot-product attention by calling the ragged_paged_attention kernel."""
         md = attention_metadata
-        kv_cache_spec = P(None, None, "model")
+        kv_cache_spec = P("data", None, "model")
 
         in_specs = (
             self.query_tnh,  # q
             self.keyvalue_skh,  # k
             self.keyvalue_skh,  # v
             kv_cache_spec,  # kv_cache
-            P(),  # md.seq_lens: Replicated
-            P(),  # page_indices_flat: Replicated
-            P(),  # query_start_loc: Replicated
-            P(),  # distribution: Replicated
+            P("data"),  # md.seq_lens
+            P("data"),  # page_indices_flat
+            P("data"),  # query_start_loc
+            P("data"),  # distribution
             P(('model')),  # sinks
         )
         out_specs = (self.attn_o_tnh, kv_cache_spec)
@@ -185,12 +199,12 @@ class GptOssAttention(nnx.Module):
             )
 
         output_TNH, kv_cache = jax.jit(
-            shard_map.shard_map(
+            jax.shard_map(
                 _ragged_paged_attention_wrapper,
                 mesh=mesh,
                 in_specs=in_specs,
                 out_specs=out_specs,
-                check_rep=False,
+                check_vma=False,
             ))(
                 q_TNH,
                 k_SKH,
@@ -235,9 +249,8 @@ class GptOssAttention(nnx.Module):
             # q_scale = self._q_scale
             k_scale = self._k_scale
             v_scale = self._v_scale
-            k_TKH, v_TKH = utils.quantize_kv(k_TKH, v_TKH,
-                                             self.kv_cache_quantized_dtype,
-                                             k_scale, v_scale)
+            k_TKH, v_TKH = quantize_kv(self.kv_cache_quantized_dtype, k_TKH,
+                                       v_TKH, k_scale, v_scale)
 
         with jax.named_scope("attn_op"):
             new_kv_cache, attn_out_TNH = self.attention(

tpu_inference/layers/jax/attention/llama4_attention.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
 
 import jax
@@ -5,8 +19,8 @@ import jax.numpy as jnp
 from flax import nnx
 from jax.sharding import Sharding
 
-from tpu_inference import utils
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
 from tpu_inference.layers.jax.attention.attention import Attention, KVCache
 from tpu_inference.layers.jax.rope_interface import apply_rope
 from tpu_inference.logger import init_logger
@@ -114,9 +128,8 @@ class Llama4Attention(Attention):
             # q_scale = self._q_scale
             k_scale = self._k_scale
             v_scale = self._v_scale
-            k_SKH, v_SKH = utils.quantize_kv(k_SKH, v_SKH,
-                                             self.kv_cache_quantized_dtype,
-                                             k_scale, v_scale)
+            k_SKH, v_SKH = quantize_kv(self.kv_cache_quantized_dtype, k_SKH,
+                                       v_SKH, k_scale, v_scale)
 
         with jax.named_scope("attn_op"):
             new_kv_cache, outputs_TNH = self.attention(

tpu_inference/layers/jax/base.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 dataclasses
 from dataclasses import dataclass, fields
 from typing import Any, Callable, Mapping

tpu_inference/layers/jax/constants.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.
 """
 Current Used Abbreviation for Tensor Dimensions:
 B: Batch size

tpu_inference/layers/jax/layers.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
 from typing import Any
 

tpu_inference/layers/jax/misc.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 Tuple
 

tpu_inference/layers/jax/moe/__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.