diff --git a/tpu_inference/models/jax/gemma4.py b/tpu_inference/models/jax/gemma4.py
index 4fc04d1d8e..88e94e9121 100644
--- a/tpu_inference/models/jax/gemma4.py
+++ b/tpu_inference/models/jax/gemma4.py
@@ -33,8 +33,7 @@
 from tpu_inference.layers.jax import JaxModule
 from tpu_inference.layers.jax.embed import JaxEmbed
 from tpu_inference.layers.jax.linear import (JaxEinsum, JaxLinear, JaxLmHead,
-                                             JaxMergedColumnParallelLinear,
-                                             JaxQKVParallelLinear)
+                                             JaxMergedColumnParallelLinear)
 from tpu_inference.layers.jax.moe.moe import JaxRoutedExperts
 from tpu_inference.layers.jax.norm import JaxRmsNorm
 from tpu_inference.layers.jax.pp_utils import PPMissingLayer, make_layers
@@ -309,23 +308,36 @@ def __init__(self,
                            None) if _shard_kv_on_k else (None, None, "model")
         _kv_bias_spec = ("model", None) if _shard_kv_on_k else (None, "model")
 
+        self.q_proj = JaxEinsum(
+            "TD,DNH->TNH",
+            (self.hidden_size, self.num_heads, self.head_dim),
+            bias_shape=(self.num_heads,
+                        self.head_dim) if config.attention_bias else None,
+            param_dtype=dtype,
+            kernel_init=nnx.with_partitioning(init_fn, (None, "model", None)),
+            bias_init=nnx.with_partitioning(init_fn, ("model", None))
+            if config.attention_bias else None,
+            rngs=rng,
+            quant_config=quant_config,
+            prefix=prefix + ".q_proj",
+        )
+        self.k_proj = JaxEinsum(
+            "TD,DKH->TKH",
+            (self.hidden_size, self.num_kv_heads, self.head_dim),
+            bias_shape=(self.num_kv_heads,
+                        self.head_dim) if config.attention_bias else None,
+            param_dtype=dtype,
+            kernel_init=nnx.with_partitioning(init_fn, _kv_kernel_spec),
+            bias_init=nnx.with_partitioning(init_fn, _kv_bias_spec)
+            if config.attention_bias else None,
+            rngs=rng,
+            quant_config=quant_config,
+            prefix=prefix + ".k_proj",
+        )
         if use_k_eq_v:  # TODO: Add QKV fusion logic for k == v case.
-            self.qkv_proj = None
-            self.q_proj = JaxEinsum(
-                "TD,DNH->TNH",
-                (self.hidden_size, self.num_heads, self.head_dim),
-                bias_shape=(self.num_heads,
-                            self.head_dim) if config.attention_bias else None,
-                param_dtype=dtype,
-                kernel_init=nnx.with_partitioning(init_fn,
-                                                  (None, "model", None)),
-                bias_init=nnx.with_partitioning(init_fn, ("model", None))
-                if config.attention_bias else None,
-                rngs=rng,
-                quant_config=quant_config,
-                prefix=prefix + ".q_proj",
-            )
-            self.k_proj = JaxEinsum(
+            self.v_proj = None
+        else:
+            self.v_proj = JaxEinsum(
                 "TD,DKH->TKH",
                 (self.hidden_size, self.num_kv_heads, self.head_dim),
                 bias_shape=(self.num_kv_heads,
@@ -336,24 +348,8 @@ def __init__(self,
                 if config.attention_bias else None,
                 rngs=rng,
                 quant_config=quant_config,
-                prefix=prefix + ".k_proj",
+                prefix=prefix + ".v_proj",
             )
-            self.v_proj = None
-        else:
-            self.qkv_proj = JaxQKVParallelLinear(
-                hidden_size=self.hidden_size,
-                num_heads=self.num_heads,
-                num_kv_heads=self.num_kv_heads,
-                head_dim=self.head_dim,
-                use_bias=config.attention_bias,
-                dtype=dtype,
-                rngs=rng,
-                quant_config=quant_config,
-                prefix=prefix,
-            )
-            self.q_proj = None
-            self.k_proj = None
-            self.v_proj = None
 
         self.q_norm = JaxRmsNorm(
             self.head_dim,
@@ -428,13 +424,10 @@ def __call__(
         attention_metadata: AttentionMetadata,
     ) -> Tuple[jax.Array, jax.Array]:
         md = attention_metadata
-        if self.qkv_proj is not None:
-            q, k, v = self.qkv_proj(x)
-        else:
-            k = self.k_proj(x)
-            v = k
-            # q: (T, N, H)
-            q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x) if self.v_proj is not None else k
+        # q: (T, N, H)
+        q = self.q_proj(x)
         # Q norm (always applied)
         q = self.q_norm(q)
 
@@ -1008,12 +1001,8 @@ def __call__(
 
 
 class Gemma4ForCausalLM(JaxModule, LoadableWithIterator):
+    # qkv_proj packing is removed in PR 3376 for performance gain
     packed_modules_mapping = {
-        "qkv_proj": [
-            "q_proj",
-            "k_proj",
-            "v_proj",
-        ],
         "gate_up_proj": [
             "gate_proj",
             "up_proj",