From c0693dfb4d69ad76f76a5db9a451bb33c4aac103 Mon Sep 17 00:00:00 2001 From: patrickji Date: Tue, 11 Aug 2026 21:59:30 +0000 Subject: [PATCH 1/2] separate q,k,v projection in sliding attention layers Signed-off-by: patrickji --- tpu_inference/models/jax/gemma4.py | 82 +++++++++++++----------------- 1 file changed, 35 insertions(+), 47 deletions(-) diff --git a/tpu_inference/models/jax/gemma4.py b/tpu_inference/models/jax/gemma4.py index 4fc04d1d8e..52e04f9245 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) @@ -1009,11 +1002,6 @@ def __call__( class Gemma4ForCausalLM(JaxModule, LoadableWithIterator): packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], "gate_up_proj": [ "gate_proj", "up_proj", From 451b533a55db036856ed2567615f10b6f7e7e6ee Mon Sep 17 00:00:00 2001 From: patrickji Date: Wed, 12 Aug 2026 18:17:48 +0000 Subject: [PATCH 2/2] add comment for tracking qkv_proj removal Signed-off-by: patrickji --- tpu_inference/models/jax/gemma4.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tpu_inference/models/jax/gemma4.py b/tpu_inference/models/jax/gemma4.py index 52e04f9245..88e94e9121 100644 --- a/tpu_inference/models/jax/gemma4.py +++ b/tpu_inference/models/jax/gemma4.py @@ -1001,6 +1001,7 @@ def __call__( class Gemma4ForCausalLM(JaxModule, LoadableWithIterator): + # qkv_proj packing is removed in PR 3376 for performance gain packed_modules_mapping = { "gate_up_proj": [ "gate_proj",