[Gemma3] Reduce concatenate and copy in KV cache update to improve perf

[gaurides]

Description of the change

---

This PR fixes issue #2780
When running Gemma3 on CPU, we observe significant % of time is spent on concatenate when doing KV cache update. This PR helps reduce time for KV Cache update and thus improve tokens/sec performance.

More details for the code change

The original code -- what it does and why it's slow

The cache is one big 6D tensor: [B, num_layers, 2, max_length, num_heads, head_dim]

cache shape: [1, 34, 2, 2048, 4, 256] in BF16 = ~136 MB
              B  layers K/V  seq  heads dim

During each token generation step, the flow is:

# 1. Slice out per-layer cache (cheap -- just a view)
current_cache = cache[:, i, ...]  # shape: [B, 2, S, H, D]

# 2. Inside attention: unstack K and V (cheap -- views)
key_cache = cache[:, 0, ...]    # [B, S, H, D]
value_cache = cache[:, 1, ...]  # [B, S, H, D]

# 3. Update single token position (DUS on 4MB tensor -- acceptable)
key = ops.slice_update(key_cache, start, key_update)
value = ops.slice_update(value_cache, start, value_update)

# 4. BOTTLENECK: stack K+V back together (copies ~4 MB per layer)
cache = ops.stack((key, value), axis=1)  # new [B, 2, S, H, D] tensor
#   This happens 34 times (once per layer) = 34 x 4 MB = ~136 MB

# 5. BIGGEST BOTTLENECK: stack all 34 layers together
cache = ops.stack(caches, axis=1)  # new [B, 34, 2, S, H, D] tensor
#   This copies all 34 layer caches into one new 136 MB tensor

Steps 4 and 5 together account for the ~18% concatenate + ~10% copy_bitcast seen in the profile.

This solution -- what changes

The key idea: if the cache is never a single tensor, it never needs to be assembled into one.

Instead of one [B, 34, 2, S, H, D] tensor, the cache becomes:

cache = (
    (key_layer0, value_layer0),   # each is [B, S, H, D] = ~2 MB
    (key_layer1, value_layer1),
    ...
    (key_layer33, value_layer33),
)
# A Python tuple of 34 tuples -- 68 separate tensors, each ~2 MB

The generation flow becomes:

# 1. Get per-layer cache (Python tuple indexing -- zero cost)
current_cache = cache[i]  # returns (key_cache, value_cache) tuple

# 2. Inside attention: unpack tuple (zero cost)
key_cache, value_cache = cache  # just Python variable binding

# 3. Update single token position (DUS on 2MB tensor -- same as before)
key = ops.slice_update(key_cache, start, key_update)     # [B, S, H, D]
value = ops.slice_update(value_cache, start, value_update) # [B, S, H, D]

# 4. Return as tuple (ZERO COST -- no data copy)
cache = (key, value)   # Python tuple, no ops.stack

# 5. Collect all layers (ZERO COST -- no data copy)
cache = tuple(caches)  # Python tuple, no ops.stack

Steps 4 and 5 are now free. Creating a Python tuple doesn't copy any tensor data -- it just creates a container holding references to the existing tensors.

Why does this work with JAX's while_loop?

JAX's while_loop (which the sampler uses) requires that loop-carried variables have the same "pytree structure" across iterations. A pytree is JAX's term for nested containers of tensors -- tuples, lists, dicts are all valid.

Our tuple of tuples:

((tensor, tensor), (tensor, tensor), ..., (tensor, tensor))

is a valid pytree. Each individual tensor [B, S, H, D] maintains the same shape across iterations (only the content at position cache_update_index changes). So JAX is happy.

Summary of data movement per token

Operation                          Original        V3
-------------------------------------------------------
Per-layer slice_update on key      2 MB            2 MB (same)
Per-layer slice_update on value    2 MB            2 MB (same)
Per-layer K/V reassembly           4 MB (stack)    0 (tuple)
Cross-layer reassembly             136 MB (stack)  0 (tuple)
-------------------------------------------------------
Total per token                    ~272 MB         ~136 MB

The slice_update copies remain (unavoidable -- attention needs the full K/V with the new token inserted), but the stacking copies are completely eliminated.

Colab Notebook

There is no change in API or model, the change is in implementation.

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
• [Y] My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• [Y] I have followed the Keras Hub Model contribution guidelines in making these changes.
• [Y] I have followed the Keras Hub API design guidelines in making these changes.
• [Y] I have signed the Contributor License Agreement.

[gemini-code-assist]

Code Review

This pull request refactors the cache mechanism in Gemma 3 models to use a tuple of per-layer (key, value) pairs instead of a stacked dense tensor, avoiding stacking overhead. The feedback highlights that the docstring for call_with_cache in gemma3_causal_lm.py should be updated to accurately reflect this new tuple-of-tuples structure instead of describing it as a dense float Tensor.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The cache format has been successfully updated to a tuple of per-layer (key, value) pairs to avoid stacking overhead. However, please note that the docstring for call_with_cache (around line 148) still documents the cache argument as a dense float Tensor, the cache of key and value.. To prevent confusion for users or developers reading the API documentation, please update the docstring of call_with_cache to accurately describe the new tuple-of-tuples structure.

References

1. Type information is provided in the docstring Args section using the format arg_name: type. description. ^(link)

[gaurides]

Thanks for catching the docstring. I have pushed a commit to fix it.

[laxmareddyp]

@gaurides

Thanks for the detailed write-up. The optimization rationale is clear — eliminating ops.stack reassembly during generation makes sense for CPU.

A few things to address:

Pattern divergence:

Right now every model in the repo uses ops.stack(caches, axis=1) for cache assembly and cache[:, 0, ...] / cache[:, 1, ...] for unpacking. This PR makes Gemma3 the only model with a tuple-of-tuples cache format.

If this optimization is worth doing, we should adopt it as a repo-wide pattern, not a Gemma3-only special case. A one-off divergence means any future sampler or utility that assumes a tensor cache will break silently for Gemma3 only.

GPU/TPU profiling:

The profiling is from CPU. On GPU/TPU, XLA typically optimizes these stacks away via buffer reuse. Do you have numbers showing this helps on GPU/TPU too?

Testing:

The tuple cache format hasn't been used anywhere in the repo before, and it flows through ops.while_loop and tree.map_structure in the samplers. Please confirm existing Gemma3 tests pass across all 3 backends (JAX, TF, PyTorch) — especially generate() with beam and contrastive samplers.

[gaurides]

Thanks @laxmareddyp for your comments; these are good points. I will follow-up on these.

[divyashreepathihalli]

have you benchmarked the performance improvement? if yes, can you share it with us.

[gaurides]

@divyashreepathihalli : not yet, I was vacation. I will get back to it this week.