Porting Gemma 4 in Pure JAX Across TPU and GPU Architectures

JAX compiles Python code to XLA HLO IR. This pipeline promises target-agnostic execution. Write a neural network model once, run it on Google TPU or NVIDIA GPU. Porting Gemma 4 exposes limits in this compiler abstraction. Hardware specifics leak through sharding APIs, memory layout choices, and communication primitives.

Gemma 4 relies on Grouped-Query Attention (GQA), Rotary Position Embeddings (RoPE), and SwiGLU activation layers. Implementing these blocks in pure JAX requires explicit control over memory allocation and tensor layout. XLA compiles these blocks differently depending on the target hardware architecture.

Let's examine a pure JAX implementation of the Gemma 4 transformer layer attention block. The code uses jax.numpy and explicit sharding annotations.

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P
from jax.sharding import NamedSharding
 
def gemma4_attention_layer(q_inputs, kv_inputs, w_q, w_k, w_v, w_o, mesh, num_heads, num_kv_heads, head_dim):
    # q_inputs: [batch, seq_len, embed_dim]
    # kv_inputs: [batch, seq_len, embed_dim]
    # Weights: w_q [embed_dim, num_heads * head_dim]
 
    # Sharding specifications
    q_sharding = NamedSharding(mesh, P("data", None, "model"))
    kv_sharding = NamedSharding(mesh, P("data", None, None))
 
    # Project inputs
    q = jnp.dot(q_inputs, w_q)
    q = q.reshape(q.shape[0], q.shape[1], num_heads, head_dim)
    q = jax.lax.with_sharding_constraint(q, q_sharding)
 
    k = jnp.dot(kv_inputs, w_k)
    k = k.reshape(k.shape[0], k.shape[1], num_kv_heads, head_dim)
    k = jax.lax.with_sharding_constraint(k, kv_sharding)
 
    v = jnp.dot(kv_inputs, w_v)
    v = v.reshape(v.shape[0], v.shape[1], num_kv_heads, head_dim)
    v = jax.lax.with_sharding_constraint(v, kv_sharding)
 
    # Scaled dot-product attention
    scale = 1.0 / jnp.sqrt(head_dim)
    # Repeat K/V heads to match Q heads for GQA
    k_rep = jnp.repeat(k, num_heads // num_kv_heads, axis=2)
    v_rep = jnp.repeat(v, num_heads // num_kv_heads, axis=2)
 
    attn_weights = jnp.einsum("bshd,bthd->bsht", q * scale, k_rep)
    attn_probs = jax.nn.softmax(attn_weights, axis=-1)
    output = jnp.einsum("bsht,bthd->bshd", attn_probs, v_rep)
 
    # Project out
    output = output.reshape(output.shape[0], output.shape[1], -1)
    output = jnp.dot(output, w_o)
    return jax.lax.with_sharding_constraint(output, NamedSharding(mesh, P("data", None, None)))

This code highlights the first major layout leak: GQA head duplication. On TPU v5e, the Matrix Multiply Unit (MXU) favors symmetric dimensions. The jnp.repeat operation maps to hardware-level memory broadcasting without copying data. On NVIDIA H100, XLA compiles this repeat into an explicit memory expansion in Global Memory (HBM3) unless the compiler fuses the operation with the subsequent einsum. If fusion fails, memory bandwidth utilization drops.

Sharding Mechanics and Hardware Interconnects

JAX uses NamedSharding to distribute tensors across devices. The underlying hardware topology dictates how the compiler schedules communication. TPU clusters use Inter-Chip Interconnect (ICI) configured in 2D or 3D torus topologies. This setup provides uniform latency and bandwidth between adjacent chips.

GPU clusters rely on NVLink for intra-node communication and InfiniBand or RoCE for inter-node communication. This creates a non-uniform memory access (NUMA) environment. If you shard Gemma 4 across multiple nodes using a simple 2D grid, the communication cost varies based on the physical location of the devices.

Consider a tensor sharded across 16 devices. On TPU v5e, a 4x4 mesh maps to physical chips with minimal routing hops:

tpu_mesh = Mesh(jax.devices().reshape(4, 4), ("data", "model"))

On GPU, if the 16 devices span two nodes (8 GPUs per node), the same mesh definition can lead to cross-node NVLink bottlenecks. The compiler must route model-parallel updates across the slower inter-node network. To prevent this, you must define hierarchical meshes that match the physical hardware:

gpu_mesh = Mesh(jax.devices().reshape(2, 8), ("data", "model"))

This explicit mesh layout ensures that all-gather and reduce-scatter operations for the attention projection weights remain inside the high-bandwidth NVLink domain.

Memory Layouts and XLA Compilation Leaks

TPU hardware uses Big Endian storage and prefers NHWC (or equivalent channel-last) layouts. The MXU processes 128x128 matrix tiles. GPU Tensor Cores process smaller tiles (typically 16x16 or 32x32) and prefer row-major layouts.

During Gemma 4 execution, Rotary Position Embeddings (RoPE) require slicing and concatenating tensor dimensions.

def apply_rope(x, cos, sin):
    # x shape: [batch, seq_len, heads, dim]
    half_dim = x.shape[-1] // 2
    x1 = x[..., :half_dim]
    x2 = x[..., half_dim:]
    rotated = jnp.concatenate([-x2, x1], axis=-1)
    return x * cos + rotated * sin

On TPU, the compiler fuses this slice-and-concatenate sequence directly into the Vector Processing Unit (VPU) register load operations. No intermediate tensors are written to HBM.

On GPU, the XLA GPU compiler often fails to fuse the concatenation step when it is followed by element-wise multiplication. This failure forces XLA to write the intermediate rotated tensor back to GPU global memory. This write operation increases memory bandwidth consumption and can cause kernel launch bottlenecks. To bypass this, you must write a custom Triton kernel or structure the JAX code to avoid explicit concatenation:

def apply_rope_fused(x, cos, sin):
    # Avoid concatenation by using element-wise operations with sign flips
    # Reshape to pair adjacent elements
    x_paired = x.reshape(x.shape[:-1] + (x.shape[-1] // 2, 2))
    # Rotate: [x1, x2] -> [-x2, x1]
    rotated = jnp.stack([-x_paired[..., 1], x_paired[..., 0]], axis=-1)
    rotated = rotated.reshape(x.shape)
    return x * cos + rotated * sin

This formulation compiles to a single fused CUDA kernel, reducing memory traffic on NVIDIA hardware.

Memory Management and Out-of-Memory Conditions

JAX manages memory allocation via pre-allocation. By default, JAX claims 90% of GPU memory when the process starts. TPU runs pre-allocate nearly 99% of HBM. This behavior prevents memory fragmentation but introduces issues when running Gemma 4 with long context windows.

Gemma 4 uses a dynamic Key-Value (KV) cache to store past keys and values during autoregressive decoding. In long-context tasks (up to 32k tokens), the KV cache size can exceed the remaining free memory. If JAX pre-allocates too much memory for model parameters and activations, the dynamic allocation of the KV cache triggers an OOM error, a frequent bottleneck in llm token cost optimization.

To prevent this on GPU (or when running local AI models on your laptop), you must tune the pre-allocation fraction:

export XLA_PYTHON_CLIENT_MEM_FRACTION=0.80
export XLA_PYTHON_CLIENT_PREALLOCATE=false

Disabling pre-allocation allows the GPU memory manager to allocate memory dynamically. This change introduces a minor allocation overhead during the first few steps but prevents OOMs during long-context generation. On TPU, pre-allocation is managed by the Cloud TPU runtime, and disabling it can degrade performance. You must rely on static shape definitions and compiler-directed memory planning instead.

Collective Communication Performance

The choice of hardware dictates the performance of collective operations. Gemma 4's MLP block uses SwiGLU activation:

def swiglu_mlp(x, w_gate, w_up, w_down, mesh):
    col_sharding = NamedSharding(mesh, P("data", "model"))
    row_sharding = NamedSharding(mesh, P("data", None))
 
    gate = jnp.dot(x, w_gate)
    up = jnp.dot(x, w_up)
    intermediate = (gate * jax.nn.sigmoid(gate)) * up
    intermediate = jax.lax.with_sharding_constraint(intermediate, col_sharding)
 
    output = jnp.dot(intermediate, w_down)
    return jax.lax.with_sharding_constraint(output, row_sharding)

The transition from column-sharded intermediate activation to row-sharded output requires an all-reduce operation.

On TPU, the all-reduce is executed over the ICI links. The TPU compiler schedules this communication to overlap with the computation of the next transformer layer. The compiler uses the dedicated Ring-AllReduce hardware within the TPU pod.

On GPU, the all-reduce uses the NCCL library. If the GPU nodes are connected via a standard Ethernet network instead of InfiniBand, the latency of this all-reduce dominates the execution time. The GPU cores sit idle waiting for the NCCL kernel to finish. To mitigate this, you must enable the XLA GPU latency hiding scheduler.

Tuning Compiler Flags

Achieving parity between TPU and GPU performance requires passing target-specific flags to the XLA compiler. For GPU deployments, these flags optimize kernel launch queues and stream synchronization:

export XLA_FLAGS="-xla_gpu_enable_latency_hiding_scheduler=true \
                  -xla_gpu_enable_highest_priority_async_stream=true \
                  -xla_gpu_all_reduce_combine_threshold_bytes=10485760 \
                  -xla_gpu_enable_triton_gemm=true"

The xla_gpu_all_reduce_combine_threshold_bytes flag groups small all-reduce operations into a single call, reducing kernel launch overhead. The xla_gpu_enable_triton_gemm flag allows XLA to generate custom Triton kernels for matrix multiplications, bypassing slower cuBLAS calls in some contexts.

On TPU, the runtime handles these optimizations automatically. The compiler uses the XLA TPU scheduler to optimize the instruction stream based on the physical topology of the TPU pod.

Profiling and Identifying Bottlenecks

To debug performance disparities, you must capture profiles using the JAX Profiler. The profiler shows step times, memory usage, and execution traces.

On GPU, open the trace in TensorBoard and look for Host-to-Device (H2D) transfers. If you see frequent H2D transfers during the training loop, it means some operations are falling back to the CPU. This fallback often happens when using JAX operations that lack GPU kernel implementations, such as certain dynamic slicing patterns.

On TPU, examine the MXU Utilization metric. If the MXU utilization is below 50%, the TPU is memory-bandwidth bound. This bottleneck is often caused by un-fused element-wise operations or inefficient input data pipelines. Ensure that your input pipeline uses tf.data or Grain with prefetching enabled to keep the TPU fed with data.

Porting Gemma 4 in pure JAX reveals that write-once, run-anywhere is only half the story. While the code remains identical, achieving optimal performance requires configuring sharding meshes, memory allocation strategies, and compiler flags to match the target hardware architecture.