A Pallas kernel tiled for TPU VMEM splits one Gemma 4 port into two checkpoints
One hand-written JAX port runs on TPU v5e, v6e and a Turing T4G, but the quantized fast path stays on TPU and the compute dtype has to be read off the live device, because neither leak ever shows up as red in a log.
The Engineer · Build desk
What happened
- A hand-written pure-JAX Gemma 4 port runs the same code on Cloud TPU v5e, on TPU v6e, and on an NVIDIA T4G attached to an AWS Graviton2 host.
- Gemma 4 makes any port carry two attention geometries, head_dim 256 on sliding layers and 512 on global ones, where most inference stacks assume a single head dimension per model.
- On the vLLM path those mixed head dimensions rule out FA4 and force the Triton attention backend, a selection driven by the model config rather than by the operator.
- Because the fused W4A16 Pallas kernel is tiled for TPU VMEM, the TPU rigs serve the -qat-w4a16-ct export while the GPU rigs serve the dense reference checkpoint at 16-bit.
- Rather than trust a config file, the port reads compute capability off the live device and picks float16 on pre-Ampere hardware and bfloat16 everywhere else.
Why it matters
- What it limits constraint Weight parity is what the one-port claim gives up, so any quality or latency comparison run across the TPU and GPU rigs is comparing two weight sets and not one port on two chips.
- What it costs cost The silent bf16 emulation bills you in decode throughput with no log line to show for it, which means the cost lands on whoever is later asked why the rig feels slow.
- Who has to decide decision For models with irregular head dimensions, picking a kernel-tiled serving stack over plain XLA becomes a hardware-eligibility decision, because one path carries a per-block memory ceiling and the other does not.
- What it sets up precedent Any kernel advertised as portable while tied to VMEM sizing now has a reasonable expectation attached to it: compute the requirement at boot, or have the portability claim tested at the first token.
The fused W4A16 path is written in Pallas and tiled for TPU VMEM, which gives 16 MB per core, and at this model's shapes the tiles want 550 KiB to 1.1 MiB per block [17]. Lower that to a GPU and Pallas goes through Triton, where per-block scratch is shared memory, 64 KiB of it on Turing [18]. Against that ceiling the tile range is 8.6x to 17.6x over budget [25]. The two budgets are not even the same unit of accounting, since VMEM is counted per core and shared memory per block, and the nominal gap between the two figures is 256x [26]. Ada lifts the ceiling and still does not come near a megabyte [18]. `check_w4a16_fits_scoped_memory()` runs that comparison at startup and refuses with the arithmetic attached instead of dying at the first token [19]. The author's one-line version is that Pallas is portable as an API and not portable as a memory model [3]. Measured against the project's own bar, that the only thing changing between rigs is a config file, that is one leak [5].
The second leak does not announce itself. A wrong compute dtype does not raise; it emulates, and XLA routes bfloat16 on a pre-Ampere GPU through fp32 with nothing red in the logs [2]. On the SM 7.5 card the policy resolves to float16, Turing's only real 16-bit datapath, since the part has neither bf16 nor fp8 [14]. The first line the process emits carries the whole decision, including `pallas_interpret=False` [15], which the write-up notes is the difference between serving and silently running a simulator [16]. A simulator that answers your HTTP requests is the most courteous way to be wrong.
What kills the vLLM path on Turing is a shared-memory request of 98,304 bytes against a 65,536-byte hardware limit [11], which is 1.5x the budget [24]. That figure travels only if your rig makes the same backend selection at the same shapes on a card with 64 KiB of shared memory per block. Somewhere with a higher ceiling the same code may fit and tell you nothing about the T4G. In this port there is no vLLM in the path at all, and attention is ordinary XLA rather than a hand-tiled kernel, so the attention path has no per-block ceiling to exceed [4][12].
The geometry is still work for whoever ports next. The 8:1 MQA means the KV budget bears no relation to what the parameter count implies [7], and a KV-share map folds 35 layers onto 15 caches, an average of 2.3 layers per cache [8][27]. On top of that sit a 512-slot sliding ring and a 4.70 GB per-layer-embedding table that is quantized to 4 bits on load [9]. The write-up reports a week lost to a padding-eviction bug in that ring, under an invariant that a cache index is an absolute real position and that padding never occupies an index a real position uses [21]. The text breaks off mid-sentence there [22]. This is one engineer's set of rented rigs, published on dev.to [23], so read the memory ceilings as checkable arithmetic and the lost decode throughput as an unquantified characterisation.
What to watch
- The write-up breaks off inside the KV ring padding-eviction section, and the invariant test that closes it is the part worth reading when it lands.
- A W4A16 kernel retiled against shared-memory blocks rather than VMEM would put one checkpoint back on both rigs.
- A vLLM attention backend that accepts heterogeneous head dimensions without falling back to Triton would remove the Turing failure entirely.
Clarity's read
What the record supports and how the coverage leans. The claims behind it follow.
Reality
- Evidence54
- Adoption14
- Hype gap+8
- Incentives32
- Confidence45
Claim ledger
Ranked by verification strength, evidence, and original report placement.
- [1]
The same hand-written pure-JAX Gemma 4 port code runs on Cloud TPU v5e and v6e and on an NVIDIA T4G attached to an AWS Graviton2 host.
- [2]
A wrong compute dtype does not raise, it emulates: bfloat16 on a pre-Ampere GPU is routed through fp32 by XLA, most of decode is lost to conversion, and nothing in the logs is red.
- [3]
The author's stated conclusion is that Pallas is portable as an API and not portable as a memory model, because a kernel written against VMEM does not describe a GPU.
- [4]
The Gemma 4 port lives in ports/gemma4/ and is driven by a generation loop behind an OpenAI-compatible server, with no PyTorch, no vLLM and no torch_xla.
- [5]
The project's premise is that if the port is really portable, the only thing that should change between those rigs is a config file, and the question is settled by measurement rather than by reading docs.
- [6]
Gemma 4 has two attention geometries: sliding layers use head_dim=256 and global layers use 512, while most inference stacks assume one head dimension per model.
- [7]
The model uses 8:1 MQA, so the KV budget is nothing like what the parameter count would suggest.
- [8]
A KV-share map collapses 35 layers onto 15 caches.
- [9]
The port must carry a 512-slot sliding ring plus per-layer embeddings held in a 4.70 GB table that gets quantized to 4 bits on load.
- [10]
On the vLLM path the heterogeneous head dimensions {'sliding_attention': 256, 'full_attention': 512} force the TRITON_ATTN backend because FA4 is not available.
- [11]
On a Turing GPU the Triton attention backend raises triton.runtime.errors.OutOfResources for shared memory, required 98304, hardware limit 65536.
- [12]
In the JAX port attention is ordinary XLA rather than a hand-tiled kernel, so there is no per-block shared-memory ceiling in the attention path at all.
- [13]
The port does not take the compute dtype from a config file; it reads the live compute capability off the device and sets COMPUTE_DTYPE = float16 if IS_PRE_AMPERE else bfloat16.
- [14]
That policy resolves to bfloat16 on TPU, where the MXUs run it natively, bfloat16 on an SM 8.9 Ada card, and float16 on the SM 7.5 Turing card, which is Turing's only real 16-bit datapath since it has neither bf16 nor fp8.
- [15]
The first line the process emits states the decision: platform=gpu compute_capability=7.5 compute_dtype=float16 pallas_interpret=False.
- [16]
According to the write-up, pallas_interpret=False matters just as much as the dtype, because it is the difference between serving and silently running a simulator.
- [17]
The fused W4A16 kernel is written in Pallas and tiled for TPU VMEM, which gives 16 MB per core, and at this model's shapes the tiles want 550 KiB to 1.1 MiB per block.
- [18]
On a GPU Pallas lowers through Triton and those tiles become shared memory; Turing gives 64 KiB per block, and Ada raises the ceiling but nowhere near a megabyte.
- [19]
The same Pallas kernel cannot run on either GPU, and the rig computes the requirement at startup and refuses with the arithmetic attached via check_w4a16_fits_scoped_memory(), rather than dying as a cryptic OutOfResources at the first token.
- [20]
The practical consequence is that the GPU rigs serve the dense reference checkpoint at 16-bit while the TPU rigs serve the -qat-w4a16-ct export: same port, same model family, different weights.
- [21]
A padding-eviction bug in the KV ring cache cost a week; the invariant is that a cache index is an absolute real position and padding never occupies an index a real position uses.
- [22]
The source text breaks off mid-sentence while describing a port that right-pads, so the resolution of the KV cache bug is not stated.
- [23]
The account is a single write-up published on dev.to describing the author's own rented accelerators.
- [24]
The Triton attention backend's shared-memory request on the Turing card is 1.5x what the hardware allows per block.
- [25]
The Pallas W4A16 tiles are between 8.6x and 17.6x the 64 KiB per-block shared memory Turing provides.
- [26]
The nominal scratch budgets differ by 256x: 16 MB of TPU VMEM per core against 64 KiB of GPU shared memory per block.
- [27]
The KV-share map averages about 2.3 layers per cache.
Get this ledger inside ChatGPT. Ask about any of our stories and the answer brings its figure, each claim’s verdict and the sources. Free, no account.
Sources
1 independent publisher whose own reporting we read for this story.
- dev.toGemma 4 in Pure JAX: What Ports from TPU to GPU, and What Doesn't
1 article · August 28, 2026
Topics and entities
Follow any of these and your For You feed starts watching them — no settings page required.