A Measurement Study of Large Language Model Inference Optimization on TPU v5e: Throughput, latency, energy, and cost results for Gemma 2B served with vLLM, with observations on the JAX and Pallas compilation stack
A single TPU v5e chip is Google’s cost-optimized AI accelerator designed for efficient inference and medium-scale training. It features 1 TensorCore , delivers up to 393 trillion Int8 operations per second (TOPS), and includes 16 GB of memory. It provides up to a 2.5x increase in price-performance over previous generations, however there more advanced chips lately.
Key Hardware Specifications
- Compute Performance: 197 BF16 TFLOPS and 393 Int8 TOPS.
- Memory Capacity: 16 GB of high-bandwidth memory (HBM2E).
- Memory Bandwidth: 819 GB/s.
- Architecture: 1 TensorCore with four matrix-multiply units (MXUs).
Availability & Usage
- Cloud Access: Rented directly on Google Cloud (via Compute Engine or Google Kubernetes Engine) and also available at Colab.
- Flexible Slices: While single chips can be utilized, they are commonly provisioned in slices of 1, 4, or 8 chips depending on your model scale.
- Cost: Single-chip VMs typically cost around $0.35 per hour on-demand.
- Framework Support: Compatible with major AI frameworks like PyTorch, JAX, and TensorFlow.
How to Think About TPUs | How to Scale Your Model. Source: Google DeepMind.
Executive Summary
This article comes from a whitepaper I recently wrote, and reports a measurement study of large language model inference on a single Google TPU v5e chip. The subject model was Gemma 2B, served through the vLLM inference engine on the JAX based TPU software stack. The study asked one question: how much of the chip’s potential does a default serving configuration leave unused, and how far can disciplined measurement recover it without changing the model or hardware.
The answer was a factor of ten. The starting configuration produced 1,540 tokens per second. The final configurations produced 13,470 tokens per second at the recommended operating point and 15,338 at the throughput maximum. The gain came from two interventions, each identified by profiling the chip rather than by guesswork: a change to how output tokens are selected, worth 2.3 times on its own, and a systematic increase in the number of requests processed together, worth a further 4.3 times. No model weights were changed, no accuracy was sacrificed, and no additional hardware was used. Because the chip and its power envelope stayed constant, every efficiency metric improved by the same factor: the upper bound on energy per token fell from roughly 128 millijoules to 13, and at published on demand pricing the cost of generation fell from roughly 4.6 million tokens per dollar to 46 million.
The study is equally defined by what did not work, and those results are reported with the same care. Speculative decoding, a widely promoted technique, reduced throughput by a factor of six in this workload’s operating regime. Weight and cache quantization, the standard next step, turned out to be unavailable on this software stack for this model, a finding documented at the source level. A precision optimization that looked certain on paper reduced throughput by 14 percent when applied, and was reverted. Each negative result was verified before being accepted, and each carries a practical lesson for teams running similar systems.
This document describes what was done, why, and with what results. The specific procedures, code, and integration techniques are not disclosed yet.
Background and Motivation
Serving a language model is a different engineering problem from training one. Training is dominated by arithmetic. Serving, and in particular the token by token generation phase called decode, is dominated by data movement. For every token generated, the chip must read essentially the entire model from its high bandwidth memory. On the chip studied here, that memory can deliver 819 gigabytes per second, and the model occupies roughly five gigabytes in its standard 16 bit form. Those two numbers set a hard ceiling on generation speed that no amount of clever software can exceed. The question for any serving deployment is how close to that ceiling it operates.
Most deployments do not know the answer, because the default configurations of serving frameworks are tuned for generality rather than for a specific model on a specific chip. This study set out to find the answer for one concrete case: Gemma 2B, a two and a half billion parameter open model from Google, served with vLLM, the most widely used open source inference engine, on a single TPU v5e chip, the efficiency oriented member of Google’s TPU family. The environment was a standard hosted notebook with one chip, which makes the study reproducible by any practitioner with the same access.
The choice of a small model on a single chip was deliberate. Small models are increasingly the economic workhorses of production systems, handling classification, extraction, routing, and drafting tasks where a frontier model would be wasteful. For these workloads, tokens per second per dollar decides profitability, and a single chip is the unit at which that metric is set; whatever efficiency is achieved on one chip multiplies directly across a fleet. The study also had a second purpose. The TPU serving stack for vLLM is new, built on a hardware plugin that unifies JAX and PyTorch execution paths and compiles model code through XLA, with performance critical kernels written in Pallas, a JAX embedded language for custom TPU and GPU kernels. New stacks have gaps, and part of the value of a measurement campaign is mapping where the advertised surface and the implemented surface differ. Several such gaps were found and are reported here.
What Was Done
The method of the study can be stated in one sentence: measure first, predict second, change third, and verify fourth. Every intervention followed the same loop. A hardware profile of the running system identified where device time was actually spent, at the level of individual compiled operations. From that profile, a specific change was predicted to produce a specific gain. The change was applied, one variable at a time, and the result was compared against the prediction. Changes that failed their prediction were reverted and recorded, not rationalized.
Three categories of change were explored. The first was the token selection path, the code that turns the model’s raw output scores into chosen tokens. The second was the batching configuration, the number of independent requests the engine processes concurrently. The third was the byte reduction category, covering weight quantization, cache quantization, and numeric precision in the output pipeline. In parallel, a custom attention kernel written in Pallas was integrated into the production serving path and validated against the stock implementation, establishing a harness for future kernel work.
Two disciplines were applied throughout. First, every measured number required a paired verification: a second counter, a repeated run, or a trace level confirmation that the intended change had actually taken effect. On three occasions this caught a configuration that had been requested but silently not applied, which would otherwise have produced false conclusions. Second, headline numbers were always taken with profiling disabled, since tracing adds overhead; profiled runs were used only to rank operations. The workload for all measurements was a fixed set of short prompts generating 128 tokens each under deterministic selection, executed as an offline batch. This is a decode heavy, short context workload; its limitations are discussed in Section 9.
Results: The Throughput Campaign
Profiling the starting configuration produced a surprise. The attention computation, the usual suspect in transformer performance work, accounted for only a small share of device time. The largest consumer, at roughly 55 percent, was the token selection machinery, executing an expensive search procedure on every step to support a sampling feature the workload was not meaningfully using. Reconfiguring the request so that selection took the direct path removed that cost entirely: throughput rose from roughly 1,540 to 3,577 tokens per second, a factor of 2.3, with no change to the model.
With the selection cost gone, the remaining time was dominated by operations whose duration is set by memory traffic: reading the feed forward weights and the large output projection over the model’s 256,000 entry vocabulary. Such operations have a decisive property: their cost per step is fixed regardless of how many requests share the step, so doubling the concurrent requests nearly doubles throughput. The study followed this logic up a doubling ladder, verifying at each rung that the predicted operations stayed fixed in cost while others grew.
Table 1. Batch scaling results. All rows are clean timing runs with profiling disabled. Latency per token is the average interval between successive tokens of one request. The watt figure uses the vendor reported 197 watt chip rating as an upper bound and is discussed in Section 7.
The scaling factors tell the story precisely. Doubling from 32 to 64 requests returned 1.81 times the throughput, nearly the ideal 2.0, confirming the chip had been waiting on memory rather than computing. Each further doubling returned less: 1.50, then 1.38, then 1.13. The diminishing curve was predicted in advance, because two categories of work do grow with the number of requests: the attention computation, whose cache reads are proportional to requests served, and the output post processing. At 512 concurrent requests the crossover was measured directly: the attention family had become the single largest consumer of the step, and further batching bought almost nothing. The recommended operating point is 256 concurrent requests, which captures 88 percent of maximum throughput at 17.7 milliseconds per token, against 31.2 at the maximum. Throughput oriented deployments can run at 512; interactive services will prefer 256 or below. The value of the table is that the trade is quantified rather than assumed.
The overall arc
Table 2. The campaign in four rows. Two interventions account for the entire gain. The chip, the model, and the model’s outputs for this workload are unchanged throughout.
A separate verification deserves mention because it anchors the credibility of every number above. The 256 request configuration was measured in three independent sessions, on rebuilt virtual machines, and returned 13,517, 13,470, and 13,484 tokens per second, a spread of one third of one percent. One outlier session returned 12,991, a 3.6 percent dip attributed to environment variance, and is reported here rather than discarded. The headline is therefore quoted as approximately 13.5 thousand tokens per second with a known spread, not as a single lucky run.
What the Measurements Revealed About the Chip
Three utilization questions matter for a serving deployment, and the study answered each with a direct measurement rather than an estimate.
First, is the chip idle. Trace analysis found roughly 85 microseconds of device silence per 11.9 millisecond generation step at 128 concurrent requests, which is 0.7 percent. The scheduler and host side machinery keep the accelerator more than 99 percent occupied even at high concurrency; host overhead is not a bottleneck in this stack, and the study can say so with a number rather than an impression.
Second, how hard is the memory system working. The two dominant weight reading operations were measured at 91 and 84 percent of the chip’s theoretical memory bandwidth while executing. These operations are finished as optimization targets in their current numeric format; only reducing the bytes themselves could make them faster. Averaged across a whole step, memory utilization is roughly 40 percent, because part of each step goes to operations limited by other factors.
Third, how hard is the arithmetic hardware working. Arithmetic utilization during generation is roughly 17 percent of peak. Quoted alone, that would suggest waste. It does not. Generation is memory bound by nature; the arithmetic units spend most of their time waiting for weights to arrive, and low arithmetic utilization alongside high memory utilization is the expected signature of a well tuned decode workload. During the prompt processing phase, which is arithmetic bound, the same chip measured roughly 47 percent of arithmetic peak, confirming the hardware performs when the workload demands it.
The trace work also produced a finding of general value: which operations are free to batch and which are not. Across an eightfold increase in concurrent requests, the feed forward weight read measured 183, 183, and then 375 microseconds, flat until the arithmetic itself began to matter at the largest batch. The attention operation over the same range measured 89.6, 179.1, and 714.4 microseconds, growth of 7.97 times against an 8 times batch increase, linear to within measurement error, while the output projection stayed flat. This comparison, predicted in advance and confirmed at three operating points, explains the entire shape of the scaling curve in Table 1 and generalizes to any transformer on any memory bound accelerator: batching is nearly free until attention and per request bookkeeping catch up, and the crossover is discoverable by measurement.
Arithmetic intensity
The single quantity behind all of these regimes is arithmetic intensity, defined for any operation as
Arithmetic Intensity = Computation FLOPs ÷ Communication Bytes.
The accelerator has an intensity of its own: peak arithmetic rate over memory bandwidth, here 1.97e14 FLOPs per second over 8.19e11 bytes per second, about 240 FLOPs per byte. An operation below 240 cannot keep the arithmetic units fed and is memory bound; above 240 it is compute bound. Because a decode matmul multiplies a batch of B token vectors against a fixed weight matrix, its intensity is approximately B, which turns the batching campaign into one prediction: the weight streaming operations should stay memory bound, and flat in cost, until the batch approaches 240 tokens per step, and become compute bound beyond it. Table 3 computes the intensities from the model’s dimensions and checks them against the measured times.
Table 3. Arithmetic intensity in FLOPs per byte, derived from the model’s published dimensions in bfloat16, against the chip’s critical intensity of about 240. The measured times of Section 4 sit within ten percent of the corresponding roofline on both sides of the threshold.
The table explains three campaign results at once. The weight streaming operations cross the 240 threshold between batch 256 and 512, exactly where their measured cost stopped being flat, which is why batch 256 is the recommended operating point: the last rung of the ladder on the memory bound side of the roofline. Attention’s intensity is about eight regardless of batch, because eight query heads share one key value head and each request reads its own cache; an operation that can never amortize its bytes across the batch is precisely one whose cost grows linearly with it. And the same threshold predicts where quantization would pay: halving the weight bytes doubles an operation’s intensity, moving the crossover to roughly batch 120 and halving the memory bound cost below it.
What Did Not Work, and Why That Matters
Speculative decoding
Speculative decoding is a technique in which a cheap mechanism drafts several tokens ahead and the model verifies them in one pass. It is frequently promoted with claims of three to four times speedups. Applied to this workload at 32 concurrent requests, it reduced throughput from roughly 3,555 to 607 tokens per second, a loss of a factor of six. The explanation is regime dependence. Speculation spends spare compute to reduce the latency of individual requests, a good trade when the chip has idle capacity, as it does at very small batch sizes. At meaningful batch sizes the chip has no spare capacity; every verified draft displaces real work, and rejected drafts are pure waste. It is a latency tool being sold as a throughput tool. The lesson: published speedups carry an implicit operating regime, and the only defense is measuring in one’s own regime, which takes a single run.
Quantization
With the bandwidth ceiling measured and nearly reached, the remaining structural lever was reducing the bytes themselves: storing weights in 8 bits instead of 16, halving traffic on exactly the operations measured at 84 to 91 percent of bandwidth. The study attempted this through every route the stack advertises, and the results form a finding of their own. The integer 8 bit route was rejected with an explicit error: accepted by the platform’s configuration surface, not implemented in its model loader for this model family. The floating point 8 bit route for weights requires a checkpoint quantized in advance and fails unhelpfully when given a standard one; there is no on the fly path. The 8 bit route for the attention cache was the most instructive: the flag was accepted end to end, the engine ran, and the cache was allocated in 16 bits anyway, with no warning. Only a direct check of the allocated cache size revealed that nothing had happened. Three attempts produced one loud failure, one unhelpful failure, one silent failure, and zero functioning quantization paths for this model on this stack version.
This is reported not as criticism but as a map; the stack is new and openly labels parts of this machinery as incomplete. The general lesson is sharp: on a fast moving inference stack, a configuration flag being accepted is not evidence that it did anything. Verify the allocation, the memory footprint, or the trace, not the flag. Teams that quote benchmark improvements from a quantization flag without such verification may be measuring nothing.
The precision experiment
Trace analysis identified a genuine inefficiency in the output pipeline: on every generation step, the model’s output scores are converted to 32 bit precision and copied, a cost that grows with batch size and reached roughly nine percent of the step at the largest batch. Under deterministic selection this precision is mathematically unnecessary, and the selection code demonstrably computes its answer before the conversion happens. Removing it looked certain to return several percent.
It did the opposite: applied as a minimal, correctness preserving change, it reduced throughput by 14 percent and multiplied time to first token by five, reproducibly, and was reverted with the baseline reconfirmed. The most probable explanation is that the conversion, wasteful in isolation, was load bearing for the surrounding compiled pipeline, which had been built and specialized around 32 bit outputs; changing the type at one point forced the runtime off its optimized path downstream. The lesson is among the most valuable in the study: on a just in time compiled stack, operation level arithmetic does not compose across compilation boundaries, and a change that removes a measured cost can add a larger unmeasured one. The inefficiency is real and has been documented for the stack’s maintainers, where it can be fixed at the pipeline level.
Custom Kernels, Pallas, and the Compilation Stack
Part of the study integrated a custom attention kernel, written in Pallas, into the live serving path of the engine, replacing the stock implementation. The integration was verified twice over: by instrumentation proving the custom code executed inside the engine’s core process, and by hardware traces attributing the attention operation’s device time to the custom kernel’s source file across multiple sessions and batch sizes. The measured result of the substitution was, by design, no change in performance: the custom kernel reproduced the stock kernel’s behavior, and throughput at the validation point was statistically identical to the stock baseline. That equality was the success condition, because it proves the harness is faithful; a substitution that changed results before any intentional modification would be measuring integration error, not kernel quality.
To be direct about what is and is not claimed: this study does not claim a throughput gain from custom Pallas kernels, and no comparison here should be read as Pallas outperforming the XLA compiled path. What the kernel work established is the platform for such gains. The traces show the attention kernel is the one major operation with substantial headroom: it runs at roughly one quarter of memory bandwidth in this short context regime, dominated by per request overhead rather than data movement, and grows to the largest consumer of the step at scale, roughly 41 percent of device time at 512 concurrent requests. A validated, hot swappable custom kernel sitting on exactly that operation is the study’s main asset for future work. The substitution and verification technique is part of the undisclosed methodology.
Observations on the compilation model
The TPU serving stack compiles the model ahead of time for a fixed menu of input sizes, called buckets, covering the possible counts of requests and tokens per step. This design delivers the near zero host overhead measured in Section 4, and it worked correctly under scrutiny: a hypothesis that the engine was running steps in an oversized bucket and wasting half its feed forward compute on padding was raised from an early reading of the traces, tested by resolving every traced operation to its true tensor shape, and found to be false. The engine had selected the correct bucket throughout. This dead hypothesis is reported deliberately, because the diagnostic that killed it, attributing every operation to its actual shape rather than trusting names or averages, is reusable and cheap.
The costs of the compilation model are also worth stating. Engine startup, which includes compiling the full bucket menu, ranged from roughly two to five minutes depending on cache state, and each new maximum batch size adds compilation work; irrelevant for long running servers, a real consideration for elastic infrastructure. The compile cache reduces repeat costs within an environment but does not persist across rebuilt environments in the hosted setting used here.
The precision experiment in Section 5.3 illustrates the model’s rigidity: the specialization that makes steady state execution fast also makes point changes to the compiled pipeline hazardous. Optimization on this stack is best pursued above the compiler, through configuration and workload shape, or below it, through whole kernels, rather than by editing the middle.
Five specific gaps between the stack’s advertised and implemented surface were documented, from the quantization findings of Section 5.2 to inefficiencies in the selection and output pipeline, including the observation that requesting the permissive value of a sampling parameter still pays that parameter’s full computational cost. All five have been prepared as reports for the maintainers; a new stack improves at the rate its users measure it.
Compile time: Pallas versus XLA on the same operation
A separate microbenchmark compared the compilation cost of the two routes to a TPU kernel. The same attention computation was prepared twice: once expressed as high level operations and compiled by XLA, and once as a hand written Pallas kernel lowered through the Mosaic compiler, with equivalent results and equivalent execution times once compiled. Over five repeated trials in fresh processes with the compilation cache disabled, the XLA route compiled in a median of 297.5 milliseconds and the Pallas route in a median of 50.7 milliseconds, a factor of 5.9 and a reduction of 83 percent, with narrow spreads on both sides.
Table 3. Compilation cost of the same attention computation by route, single shape, five trials. Execution times after compilation are equivalent; the difference is compile cost only.
The plausible mechanism is that XLA must search a large optimization and fusion space over the operation graph, while a Pallas kernel encodes the implementation decisions in advance, leaving only the lowering work. The scope of the claim is narrow, one operation at one shape, and it concerns developer iteration cost, not serving throughput. It matters nonetheless: kernel development traverses the edit, compile, and test loop hundreds of times, so a several fold reduction in the loop’s fixed cost is a material productivity difference, and it lowers the price of pursuing the custom kernel opportunity identified.
Energy
The chip’s power envelope did not change during the study, so the tenfold growth in output translates directly into a tenfold reduction in energy per token. Because the hosted environment exposes no power telemetry, the study uses the vendor reported chip rating of 197 watts as a deliberate upper bound; a peer reviewed survey both sources this figure and notes that real workload draw typically runs 30 to 60 percent below rating. All figures below are therefore conservative ceilings, and the improvement ratios are exact regardless of true wattage, since the constant cancels.
Table 4. Energy per generated token at the chip’s rated power. The tenfold improvement holds for any true wattage because the same chip produced every row.
A cross check against independent literature supports the scale of these figures. A published survey estimates roughly 94 joules per thousand tokens for a 7 billion parameter model on this chip class at a batch of 32; this study measured roughly 20 joules per thousand tokens for a 2.5 billion parameter model at a batch of 128. The model here streams roughly 3.5 times fewer weight bytes per token and runs at four times the batch, and those two factors together predict almost exactly the observed gap, the kind of external validation single environment benchmarks rarely provide. The result also carries a point for sustainability planning: batching is an energy optimization, not only a throughput one. The same silicon at the same rated power produced 4.3 times the tokens between the first and last rows of the batching campaign; fleets sized on per token energy budgets can be several times smaller than naive per chip figures suggest, provided the serving configuration is measured rather than defaulted.
Economic Implications
The commercial meaning of the study is easiest to see in tokens per dollar. The chip class used here is publicly priced at 1.20 dollars per chip hour on demand, so throughput converts directly into unit cost.
Table 5. Generation cost at 1.20 dollars per chip hour, on demand pricing. Committed use and spot pricing would lower every row proportionally.
Read as a business statement: the identical hardware bill buys ten times the product. A service generating one billion tokens per day would need roughly 7.5 chips around the clock at the starting configuration and 0.8 at the final one, before redundancy, a difference of roughly 217 versus 22 dollars per billion tokens at on demand rates. For products with per request economics, such as summarization, extraction, or agent steps, this is the difference between a feature that loses money and one that carries margin.
Two qualifications keep the claim honest. First, these are offline batch numbers on a decode heavy, short context workload; a production service with irregular arrivals, longer prompts, and latency targets will operate below the maximum, which is why Table 1 reports the full throughput and latency frontier rather than one number. The 256 request row, not the 512 row, is the realistic planning figure for interactive services. Second, none of the gain required engineering exotica: one configuration correction and one batching campaign, both reachable by any team willing to profile before tuning.
The economic finding is less that this chip is cheap and more that unmeasured serving configurations routinely forfeit most of the hardware already being paid for. The latency columns carry their own economic content: batch time to first token grew from 84 to 307 milliseconds across the campaign, and per token latency roughly quadrupled between smallest and largest batch. A chat product sells responsiveness and should buy it with a smaller batch; a nightly document pipeline sells volume and should run at the maximum. The contribution here is the measured frontier on which that choice can be made deliberately.
Scope and Limitations
The study’s claims are bounded, and the bounds are part of the result. The workload was an offline batch of short, repeated prompts with 128 token outputs under deterministic selection, a favorable regime for decode throughput: short contexts keep attention cheap, deterministic selection avoids sampling costs, and offline batching avoids the scheduling losses of irregular traffic. Services with long documents, streaming arrivals, or sampling requirements will land below these numbers, though the method for finding their own frontier is identical.
Prompt caching was enabled and, with repeated prompts, makes the prompt processing phase inexpensive; the figures are decode dominated and affected only marginally, but replications should note the setting. The hardware scope is one chip; production TPU deployments run pods of many, an axis this study deliberately did not touch. For a model of this size, replication across chips scales near linearly, so the single chip figures are best read as the unit economics of a fleet.
The software scope is one version of a fast moving stack, and several findings, in particular every quantization result, are version findings a future release may change. Power figures are rated bounds, not measurements, because the environment exposes no telemetry. Output quality was verified by determinism and inspection rather than formal evaluation, sufficient here because no accepted intervention altered the model’s computation.
Conclusions
A single mid tier accelerator chip, running an unmodified open model on an open serving engine, moved from 1,540 to 15,338 tokens per second through measurement driven configuration alone. The recommended operating point delivers roughly 13.5 thousand tokens per second at 17.7 milliseconds per token, at an energy ceiling of 15 millijoules per token and a generation cost near 2.5 cents per million tokens on demand. The chip finished the campaign more than 99 percent occupied, with its dominant memory operations at 84 to 91 percent of theoretical bandwidth: the study did not stop at a convenient number, it stopped at the hardware. The remaining headroom is mapped rather than speculative. The attention kernel operates at roughly a quarter of memory bandwidth in this regime and is the largest cost at scale, with a validated custom kernel harness sitting on exactly that operation; weight quantization would roughly halve the dominant memory traffic and is blocked only by stack maturity; and an identified inefficiency in the output pipeline is worth several percent once fixed at the right layer. Each is a known door with a measured value behind it.
The transferable findings are three. Profile before optimizing, because the largest cost in this system was not where transformer folklore said it would be. Verify that changes took effect, because this stack accepted three configurations it did not apply, and only direct measurement told the difference. And treat every published speedup as a claim about a regime, because the technique measured here as a six times loss is legitimately a win in the regime its advocates measured. None of these require special access or budget; they require the discipline of pairing every number with its verification. The procedures, code, kernel integration techniques, and patch methodology developed during the study are not disclosed in this document.
Acknowledgements
Google ML Developer Programs and Google Developers Program supported this work by providing Google Cloud Credits.
Top comments (0)