{
  "id": 4544133,
  "title": "Gemma 4 in Pure JAX: What Changes Between Turing and Ada, and What Doesn't",
  "url": "https://urgent.news/2026/08/31/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what",
  "topic": "tech",
  "section": "Tech",
  "published": "2026-08-31T01:40:09.000Z",
  "source": {
    "name": "Dev.to",
    "slug": "dev-to",
    "url": "https://dev.to/gde/gemma-4-in-pure-jax-what-changes-between-turing-and-ada-and-what-doesnt-4c5e"
  },
  "original_language": "en",
  "account": "This article measures the performance of a Gemma 4 port running in pure JAX across two NVIDIA GPUs, one from the Turing generation (NVIDIA T4G) and one from the Ada generation (NVIDIA L4). The port was executed on two AWS instances, g5g.2xlarge and g6.2xlarge, with identical configurations. The main findings are centered around two leaks and one bug that prevent the port from running optimally on both GPU types.\n\nThe first leak stems from the compute dtype policy, which is determined by the device's compute capability. On both cards, the policy sets the compute_dtype to bfloat16. However, this causes a significant performance hit on the Turing GPU, as XLA routes bfloat16 operations through fp32, resulting in a loss of performance, particularly in the decode stage. The leak is difficult to detect, as the logs do not raise any errors, and the performance degradation is not immediately apparent in the throughput measurements.\n\nThe second leak is related to the Pallas kernel used for the fused weight matrices. Pallas is optimized for a GPU with a 16 MB scratchpad per core, but the model's shapes require tiles to occupy between 550 KiB and 1.1 MiB per block. Turing GPUs provide 64 KiB per block, while Ada GPUs offer significantly more, but not enough to accommodate the entire tile. This leads to Triton sharing memory, which is not sufficient for the Gemma 4 geometry, causing the port to fail on both GPU types.\n\nThe third component discussed is a KV ring cache bug that occurs due to the irregular geometry of Gemma 4. The cache index is supposed to represent an absolute real position, but because of the port's right-padding, it violates this invariant. The bug manifests as a clean HTTP 200 response with status: success, and output that appears normal, making it difficult to detect. The only way to identify the issue is through a degeneracy check on the output, which is run on every response.",
  "summary": "This article is a measurement report on running a hand-written Gemma 4 port in pure JAX across two NVIDIA GPUs a generation apart, and on the two places the \"it's just JAX\" abstraction leaks. One of those leaks costs 87% of decode and nothing in the logs is red. The code is here: https://github.com/xbill9/gemma4-dev What This Article Measures One port, one build, one checkpoint, two cards.…",
  "key_points": [
    "Gemma 4 port runs in pure JAX on Turing and Ada GPUs",
    "Compute dtype policy sets bfloat16 on both GPUs, causing performance hit on Turing",
    "KV ring cache bug due to irregular Gemma 4 geometry prevents optimal port performance"
  ],
  "editors_take": null,
  "illustration": null,
  "coverage": {
    "outlets": 1,
    "also_reported_by": []
  },
  "ai_generated": true,
  "disclaimer": "Summaries, key points and the editor’s take are written by software from other outlets’ reporting and may contain errors — always check the linked original."
}