Build1 distinct publisher3 min readPublished
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

Compiled by The EngineerSomething wrong?How this is made
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 [16]. Lower that to a GPU and Pallas goes through Triton, where per-block scratch is shared memory, 64 KiB of it on Turing [17]. Against that ceiling the tile range is 8.6x to 17.6x over budget [23]. 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 [24]. Ada lifts the ceiling and still does not come near a megabyte [17]. `check_w4a16_fits_scoped_memory()` runs that comparison at startup and refuses with the arithmetic attached instead of dying at the first token [18]. The author's one-line version is that Pallas is portable as an API and not portable as a memory model [20]. Measured against the project's own bar, that the only thing changing between rigs is a config file, that is one leak [3].
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 [11]. 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 [13]. The first line the process emits carries the whole decision, including `pallas_interpret=False` [14], which the write-up notes is the difference between serving and silently running a simulator [15]. 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 [9], which is 1.5x the budget [22]. 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 [1][10].
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 [5], and a KV-share map folds 35 layers onto 15 caches, an average of 2.3 layers per cache [6][25]. 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 [7]. 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 [26]. This is one engineer's set of rented rigs, published on dev.to [27], so read the memory ceilings as checkable arithmetic and the lost decode throughput as an unquantified characterisation.
Ranked by verification strength, evidence, and original report placement.
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.
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.
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.
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.
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.
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.
Distinct publishers with included, body-backed reporting in this cluster.
dev.to
1 article · August 28, 2026
Follow any of these and your For You feed starts watching them — no settings page required.
build
Fitting 6.155 GB of Gemma 4 weights lands you on AWS's only Arm GPU family1 distinct publisher
build
Dual 3090s, no NVLink: the serving stack broke long before the model did1 distinct publisher
build
Your .ai viewer is a pdf.js problem, and its worst bugs never throw1 distinct publisher
build
MCP's empty-string cursor: the null check that hides most of a tool catalog1 distinct publisher
Evidence-backed comparisons of source perspectives and observed adoption signals. Read the methodology
Which Builder, Operator, and Investor concerns the observed source mix emphasized—not a truth score.
Evidence, demonstrated adoption, hype gap, incentives, and confidence are assessed independently, each on its own current evidence. How these are measured.
Verbatim traces, one witness
The technical spine is unusually checkable for a blog post: vLLM's backend-selection message, the Triton exception with both numbers (98304 required, 65536 available), the startup policy line, the named guard, and an xprof breakdown showing 0.0% Tensor Core. What none of it has is a second observer — the same person wrote the code, rented the hardware, took the measurements and published them, and the text stops mid-sentence before the profile discussion finishes.
One engineer's rented rigs
Adoption here is three machines the author is paying for. The port serves on TPU v5e, v6e and a T4G, and the operational numbers around it are real, but no other team, product or user appears anywhere in the account — and the split into two checkpoints means even this single deployment is not one artefact in production.
Self-bounding, with one leap
Most of this piece argues against its own headline: the fast path does not port, the quantized kernel is refused at startup, one deployment becomes two checkpoints. That is not how overselling reads. The overreach is narrow and specific — 'Pallas is portable as an API and not portable as a memory model' is drawn from one Turing card plus an asserted Ada ceiling, and generalises from a single W4A16 kernel's tile sizes to the programming model as a whole.
Own-repo showcase, no vendor money visible
The author is publicising his own repository, which is a reason to read the wins generously — the 117-second install, the cache restore, 'the clearest win of the whole exercise'. Against that: no vendor, employer or sponsor appears anywhere, the hardware is self-rented, and the sharpest findings are admissions that his own fast path does not run on GPU and that a bug in his own cache cost a week. Promotional pull exists; it is not steering the technical claims.
Trust the traces, not the generality
We would defend the specifics — the error strings, the tile sizes, the dtype policy, the 0.0% Tensor Core profile — because they are quoted and internally consistent. We would not yet defend what they are said to mean for Pallas, or for any hardware between a 7.5 Turing card and the Ada ceiling the author only asserts. One writer, one repository, three rented machines, no reproduction, and a text that ends before it finishes its last thought.