Repacked QAT Gemma 4 26B-A4B serves on one TPU v6e at 1.9x FP8 throughput

What happened

Google's quantization-aware-trained (QAT) Gemma 4 26B-A4B now runs on a single Google Cloud TPU v6e chip under vLLM, using 17.43 GiB of HBM, holding 53,888 tokens of KV cache and generating 1,283 output tokens per second. According to a dev.to write-up documenting the build, the previous single-chip option, RedHat's FP8 checkpoint, consumes 27.99 GiB, holds 3,456 tokens of cache and serves 668 tokens per second. That is 15.6x the KV cache and 1.9x the throughput, and on a 3,880-record classification suite the two builds score within a point of each other.

The gap at 26B

Gemma 4 26B-A4B is a mixture-of-experts model with 25.8B parameters, 128 experts per layer and 8 active per token. At bf16 it needs 48.07 GiB, while one v6e chip offers 28.74 GiB of usable HBM, so it cannot fit unquantized. Google trained 4-bit QAT variants of every Gemma 4 size, but the dev.to post identifies a hole at 26B for vLLM users: the model card lists a GGUF Q4_0 export aimed at llama.cpp, which vLLM cannot load; an "unquantized" bf16 QAT checkpoint; and Compressed-Tensors W4A16 checkpoints, the format intended for vLLM, released only for the E2B, E4B, 12B and 31B sizes. The only remaining route was the third-party FP8 build, which fits with 0.75 GiB to spare — cache for roughly one and a half 2,048-token requests.

A repack, not a re-quantization

The author filled the gap by repacking Google's "unquantized" QAT export. Its tensors are stored as bf16, but inspection showed that every group of 32 weights along the input dimension already sits on a 16-level Q4_0 grid with levels from −8 to 7, the residue of quantization-aware training. Producing a 4-bit checkpoint is therefore a change of container, not a lossy conversion.

The one trap is the per-group scale. The textbook Q4_0 rule, dividing the group's maximum absolute weight by 8, assumes the peak sits at level ±8. When it sits lower, the rule derives the wrong step and re-rounds the whole group onto a grid that does not contain the original values — about 5% median error per group, with every shape check still passing. The repack script instead tries the maximum divided by each m from 1 to 8, keeps the first step that reproduces the entire group, then refines it by least squares over the 32 values.

The resulting checkpoint is published on Hugging Face under the xbill9 account. Attention, the dense MLP and all 3,840 experts are quantized to int4 with group size 32 in compressed-tensors format; the router, embeddings, norms and vision tower are copied unchanged. Verification found zero groups off the grid, and expert values are 92.6% bit-identical to the source, with the remaining differences caused only by storing the scale at bf16, the precision vLLM's TPU W4A16 layers load it at, for a worst-case relative difference of 1.1e-2.

What changed in vLLM's TPU backend

Serving on TPU required two pieces in tpu-inference, vLLM's JAX-based TPU backend: an already-approved pull request, #3653, adding W4A16 for linear layers, and a new WNA16FusedMoEMethod, #3660, for the experts. The experts method loads the packed int4 weights and 32-wide scales itself, because the default path re-quantizes with one scale per output channel and would discard the trained grid, and routes them to the GMM backend, whose gmm_v2 kernel accepts a scale per 32-wide group and dequantizes each weight tile in fast on-chip memory before the multiply. Activations stay in bf16, so quantization remains weight-only.

The checkpoint is also portable: the same file loads unpatched on vLLM 0.30.0 on an NVIDIA L4, since the layout matches what vLLM already expects for int4 mixture-of-experts checkpoints.

Why it matters

The freed memory is the headline. Dropping weight storage from 27.99 GiB to 17.43 GiB converts directly into KV cache, 15.6x more of it, and cache is the binding constraint on context length and batch size when serving. With 1.9x the throughput and accuracy within a point of FP8 on the test suite, a 26B-parameter model becomes practical on one mid-tier accelerator. The build also demonstrates that QAT exports can be moved between containers essentially losslessly if the trained grid is respected, and the resulting artifact serves on stock vLLM beyond TPU hardware. Every script, log and per-record output is committed to the accompanying repository, so the claims are reproducible end to end.