Throughput, latency, energy, and cost results for Gemma 2B served with vLLM, with observations on the JAX and Pallas compilation stack
This is the Part 2 of my other article, “From 1,540 to 15,338 Tokens per Second on a Single TPU Chip”, where I approach the “how” I was able to first boost tokens-per-second from 1,540 to 15,338, and then to 19,5111.
The formal paper including both Part 1 and Part 1 articles is available here:
This part is the study’s methods volume. Part I reported the results: a single TPU v5e chip serving the unmodified Gemma 2B model through vLLM moved from 1,540 to 19,511 tokens per second, with the energy ceiling per token falling from roughly 128 to 10 millijoules and the generation cost from roughly 4.6 to 58.5 million tokens per dollar.
This part discloses how, strategy by strategy: the exact configurations, the code that was written and patched, the profiling procedures that motivated each change, the verification that each change had taken effect, and the outcome of each, including the failures. It is written for machine learning engineers and JAX developers; hardware background helps but is not assumed.
Seven strategies were implemented. Three succeeded: correcting the token selection path (2.3), scaling the batch (4.3), and tuning the size of the kernel block (a further 1.3). Four are reported as instructive failures or validations: speculative decoding, a Pallas kernel transplant, three quantization routes, and a numeric precision patch. Every strategy followed the same loop: profile, predict a specific gain, change one variable, verify the change took effect, and measure against the prediction.
Environment, Stack, and Measurement Harness
All work ran in a hosted Colab notebook attached to a single TPU v5e chip (v5e-1, one chip, two cores, 16 GB of high-bandwidth memory at 819 GB per second, 197 teraFLOPs per second peak in bfloat16). The software stack was vllm-tpu 0.23.0 with its tpu-inference plugin, JAX 0.10.1, libtpu 0.0.41, and torch 2.10 at the final sessions; earlier sessions resolved older versions, and the engine version line printed at startup was checked whenever numbers from different sessions were compared.
One stack detail shapes everything that follows: Gemma 2B is not among the plugin’s JAX-native model definitions, so the engine loads the vLLM-native PyTorch definition and executes it through torchax, the bridge that lowers PyTorch operations to JAX and then XLA. Custom kernels on this stack are written in Pallas and lowered through the Mosaic compiler; the stock attention kernel is a Pallas kernel called ragged paged attention, version 3.
Measurement used one benchmark script for every run, holding the protocol fixed while environment variables selected the treatment: SPEC_DECODE, BATCH_MULT, PROFILE, MODEL, QUANT, and KV_DTYPE. The script builds the engine, times engine initialization separately, runs a small warmup generation to trigger first-touch compilation paths, measures batch time to first token by generating exactly one token for the full prompt set, and then runs the main generation of 128 tokens per request under temperature 0.0. The workload is four short prompts repeated BATCH_MULT times. Inter-token latency is derived by subtracting the measured time to first token from the main run and dividing by the tokens per request. Two rules were absolute: headline throughput comes only from runs with the profiler off, because tracing adds overhead, and every claimed effect requires a verification independent of the flag that requested it.
Profiling used the engine’s start_profile and stop_profile calls wrapped around only the main generation, so traces contain the steady state and not initialization. Traces were read two ways. TensorBoard’s XProf plugin provided the HLO operation statistics table, ranking every compiled operation by device time with its full HLO text, tensor shapes, occurrence counts, and source-file attribution; traces were staged into a fresh log directory per analysis to avoid the stale-trace trap, in which TensorBoard silently displays a previous session’s data, and every table was authenticated by checking that its tensor shapes matched the batch under study before any number was read. When the table’s filtering failed, the compressed trace JSON written next to the .xplane file was parsed directly in Python, aggregating operation durations by name and, critically, by the tensor shape extracted from each event’s long name. One environment quirk is worth recording for reproducibility: on this stack the profiler ignored the requested output directory and wrote traces under /content/None/plugins/profile, so traces were located by searching for .xplane files rather than trusting the configured path.
The engine’s own efficiency features were left at their defaults and verified rather than tuned, because the profiles showed they were already doing their job. Asynchronous scheduling overlaps host-side work with device execution; the measured gap between decode steps was 85 microseconds against an 11.9-millisecond step at batch 128, under one percent idle, so no host-side optimization was warranted. Chunked prefill ran with its default budget of 8,192 tokens per pass, visible in the traces as the 4,096-token shape population. Prefix caching was enabled, and with a repeated prompt set it makes the prompt-processing phase inexpensive; this is declared as part of the workload definition since the numbers are decode-dominated. Paged attention managed the key-value cache in 128-token blocks, 3,249 blocks and 415,872 cacheable tokens under the 0.75 memory fraction, capacity that was never the binding constraint (batch 512 used roughly 72,000 tokens). The optimization surface that remained, and that the seven strategies address, was therefore the workload configuration and the compiled operations themselves, not the serving loop.
Strategy 1: The Token Selection Path
The first profile of the default configuration contradicted the expected picture. The top two operations by total device time were fusions attributed to a binary-search source file inside the sampler, together holding roughly 55 percent of the device, while the transformer’s matmuls and the attention custom call divided the remainder. The cause was the request configuration: the baseline used sampling with top- and repetition penalties, and on this stack a top- request routes every step through an implementation that performs an iterative binary search over the 256,000-entry vocabulary distribution. The search executes its full fixed iteration count, 32 iterations, even when top- is 1.0 and the operation is mathematically a no-op; the penalty transforms add further per-step tensor work over the vocabulary. The model itself was the minority of the step.
The change was one line in the request, not in the engine: SamplingParams(max_tokens=128, temperature=0.0), which selects greedy decoding and bypasses the sampling transforms entirely. Verification was the disappearance of the binary-search fusions from the post-change HLO table, with the weight-streaming matmuls promoted to the top ranks: after the fix, the fused gate-and-up projection family led the table at roughly 36 percent of device time, the output projection near 17 percent, and attention near 18 percent, the anatomy of a memory-bound decode. Throughput rose from roughly 1,540 to 3,577 tokens per second at 32 concurrent requests, a factor of 2.3, and the internal engine counter agreed with the wall-clock measurement within a few percent, the paired verification used for every headline number afterward. The observation that a permissive parameter value still pays the full computational cost was documented for the stack’s maintainers. The general procedure transfers: profile with the exact SamplingParams of production, because the selection path is part of the workload, and a team whose product genuinely requires sampling should measure that configuration rather than assume the greedy numbers.
Strategy 2: Speculative Decoding
Speculative decoding was configured through the engine’s native support: method ngram, num_speculative_tokens 5, prompt_lookup_min 2, prompt_lookup_max 4, meaning a lookup mechanism drafts up to five tokens from n-gram matches in the existing text and the model verifies them in one pass. At 32 concurrent requests this produced 607 tokens per second against the 3,555 greedy baseline of the same session, a loss of a factor of six.
The mechanism is arithmetic, not implementation failure. Verification of drafted tokens makes the forward pass process tokens per request per step; at batch 32 that is up to 192 token positions per step through a memory-bound model, and every rejected draft is discarded work. The engine’s internal counters corroborated the wall-clock number, ruling out a measurement artifact, and the configuration itself was confirmed active in the startup log, so this is a genuine regime result, not an inert flag. Speculation converts spare compute into lower per-request latency, which is profitable only when the accelerator has spare capacity, meaning very small batches; the published three-to-four-times claims for the technique come from that regime. It was not pursued further because the study’s objective was throughput at meaningful batch sizes. The one run cost a few minutes and prevented a wrong turn, and its per-token energy figure, 325 millijoules against the 55 of the greedy baseline it replaced, is the study’s clearest illustration that a mismatched optimization is also an energy regression.
Strategy 3: Batch Scaling and Compile Buckets
The post-fix profile showed the step dominated by operations whose cost is set by streaming fixed bytes: the fused gate-and-up projection reading a bfloat16 weight (134 MB per layer, measured at about 91 percent of theoretical bandwidth while executing) and the output projection over the 256,000-entry vocabulary reading roughly 1 GB (about 84 percent). Fixed cost per step means near-free amortization, so the batch was doubled repeatedly: max_num_seqs in the engine configuration raised the concurrency cap, and BATCH_MULT scaled the prompt set to fill it. Each operating point required both, and each new maximum request count triggered fresh ahead-of-time compilation.
That compilation model matters operationally. The engine precompiles a fixed menu of input sizes: token paddings of 16 through 8192 and request paddings of 8 through the configured maximum, plus separate sampling, logprob, and placeholder substitution graphs. Cold initialization ranged from 106 to 285 seconds depending on the JAX compile cache state; the cache persists within a virtual machine but not across rebuilt Colab sessions. Runtime steps then execute in the smallest precompiled bucket that fits, with host overhead measured at 85 microseconds per 11.9-millisecond step, under one percent.
The batch ladder. Energy uses the 197 W rated chip power as an upper bound; cost uses $1.20 per chip-hour on demand. Scaling factors per doubling: 1.81, 1.50, 1.38, 1.13.
Each rung was verified at the operation level by tracing and reading average times per HLO operation. Three families told the whole story, and Table [tab:opladder] gives their ladder. The gate-and-up projection stayed flat while memory-bound, then rose as its arithmetic intensity crossed the chip’s critical 240 FLOPs per byte between batch 256 and 512, where the 197-teraFLOP roofline predicts 349 microseconds against the 375 measured. The output projection behaved identically at larger scale.
The attention custom call scaled linearly, a factor of 7.97 across an 8-times batch increase, because each request reads its own key-value cache and the operation’s intensity is fixed at about 8 FLOPs per byte, the ratio of query heads to key-value heads in this model (Gemma 2B uses multi-query attention: eight query heads of dimension 256 share one key-value head; the cache is paged in 128-token blocks, 3,249 blocks in this configuration).
At batch 512 these families account for the step almost exactly: attention, 18 layers times 714 microseconds, is 12.9 milliseconds (41 percent of the 31.2-millisecond step), the gate-and-up family 6.8 milliseconds (22 percent), the down projection family 4.4 milliseconds (14 percent), and the logits pipeline, matmul plus conversion plus copy plus argmax, near 5.6 milliseconds (18 percent), summing to 30 of 31.2 milliseconds. The batch-linear conversion and copy of the logits, absent from folklore but second and third in the ranked table, became Strategy 6.
One trap in this analysis is worth disclosing in detail because it produced a wrong intermediate conclusion. Aggregating trace events by operation name suggested the decode feed-forward ops ran with a 1024-token shape, implying half of each step was padding at batch 512. Resolving every event to the tensor shape in its own HLO text disproved it: the population with 258 occurrences (the decode steps) carried shape at 375 microseconds, while the 1024-shaped instances occurred 18 times, once per layer, belonging to a single prefill pass; further shape populations at 16, 256, and 4096 tokens belonged to warmup and chunked-prefill buckets. Compiled programs reuse operation names across bucket sizes, so aggregation must key on shape, and occurrence counts must be checked against the expected step count. The corrected analysis confirmed the engine had selected the correct bucket throughout, and the planned max_num_batched_tokens intervention was cancelled before consuming any accelerator time.
What batch scaling did to compilation
Batch scaling is the strategy that interacts most with the ahead-of-time compilation model, so its costs are itemized here. Raising max_num_seqs changes the request padding menu (8, 16, 32, up to the configured maximum) and the attention request padding, so every new operating point paid a fresh precompilation of the backbone across all token buckets, the sampling graphs across all request buckets and flag combinations, the logprob graphs, and the placeholder substitution graphs, roughly 120 compilation tasks in the batch-512 configuration. Within a session the JAX compile cache made repeat initializations several times cheaper (285 seconds cold against 106 warm in the final sessions); across rebuilt virtual machines everything recompiles.
None of this touches steady-state throughput, which is why initialization was timed separately and excluded from every quoted number, but it defines the operational envelope: on this stack, changing the maximum batch is a deployment event, not a runtime knob. The Pallas attention kernel itself compiles through Mosaic once per bucket inside this same pass and was never the visible cost; the multi-minute totals are dominated by the XLA side of the menu, consistent with the compile-time microbenchmark of Section “ Measuring Compile Time: Pallas Versus XLA”
Strategy 4: The Pallas Kernel Transplant
To make the attention kernel modifiable, the stock ragged paged attention v3 source files (kernel.py, kernel_hd64.py, the tuned block-size tables, and util.py) were copied out of the installed tpu_inference package into a project directory, my_rpa, made importable with an empty __init__.py and placed on PYTHONPATH. A small module, imported by the benchmark script before the LLM constructor runs, injects the copies back into the live engine: it enumerates the callables defined in the stock kernel module, replaces each with the copy from my_rpa via setattr on the module, and additionally scans sys.modules for any tpu_inference or vllm module holding a direct reference to a replaced function, rebinding those references as well. Import order is the mechanism: the substitution must complete before the engine builds and precompiles, so the injected functions are what the compiler lowers.
Two wrapping rules mattered. Public entry points are wrapped with a functools.wraps probe that prints once on first invocation, giving runtime proof of execution from inside the engine core process. The inner kernel bodies, whose names begin with an underscore, are installed unwrapped, because pallas_call inspects the kernel function it is handed to build the grid and block specification, and a wrapper interferes with that lowering. Verification was double. At runtime, the probes printed during the engine’s precompilation pass. In hardware traces, the attention custom call, named for its block configuration (RPAd with a one-token query block and 2,048-token key-value blocks in this regime), carried source attribution my_rpa/kernel.py rather than the installed package path, across multiple sessions, batch sizes, and one stack version upgrade. That upgrade is itself a finding about the technique’s robustness: the newer tpu-inference resolved the kernel through the module attribute at call time, so the rebinding scan found zero direct references, yet the probes still fired, confirming the module-level setattr alone carried the substitution.
The measured result of the unmodified transplant was, by design, statistically identical throughput to stock; equality was the success condition, because a substitution that changed results before any intentional modification would be measuring integration error. Because the vendored kernel compiles through Mosaic inside the engine’s own precompilation pass, the substitution added no observable compile time and changed no HLO beyond the source attribution. This validated harness, an editable copy of the production attention kernel proven live, is the platform for the kernel optimization the results volume identifies as the main remaining opportunity: the kernel 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 about 41 percent of the step at batch 512. Strategy 7 reports that harness used for its purpose.
Strategy 5: Three Quantization Routes
With the dominant operations at 84 to 91 percent of bandwidth, halving weight bytes was the remaining structural lever. Three routes were attempted, and each terminated at a different layer of the stack. The verification signal for all three was the engine’s own initialization log: the model memory line (4.67 GiB for bfloat16 weights) and the key-value cache capacity line (415,872 tokens at 16-bit). A working weight quantization must roughly halve the first; a working cache quantization must roughly double the second.
Source inspection explained the pattern. The platform’s supported quantization list advertises several methods (compressed tensors, awq, fp8, and others), but the JAX quantization registry that the loader actually consults implements two, None and fp8, and the PyTorch fallback loader carries scaffolding for quantized layers with a comment stating that only unquantized linear layers are handled for now. The integer route’s checkpoint was produced with llmcompressor’s one-shot flow, scheme W8A16 over Linear targets with lm_head excluded as the conservative first configuration (in this model the output projection shares the embedding weights, the riskiest layer to quantize); the checkpoint remains valid for GPU serving even though this stack rejected it.
One procedural trap deserves a warning: a published quantized checkpoint with a nearly identical name targets Gemma 2, a different architecture from Gemma 2B, and benchmarking it would have silently invalidated every baseline; quantized checkpoints must match the exact base model of the campaign. The silent cache case is the instructive failure: the flag traveled through the whole configuration surface, appeared in the engine’s printed configuration, and changed nothing, which is why the study’s rule is to verify allocations and memory lines rather than flags. One run in the sequence also demonstrated the rule in reverse: a treatment run whose throughput matched baseline was traced to the environment variable never reaching the engine (the script predated the flag), so its number was reclassified as a third baseline replication rather than evidence about quantization. All findings were prepared as reports for the maintainers; they are version findings that a future release may change.
Strategy 6: The Output Precision Patch
The batch-512 trace ranked a 32-bit conversion of the logits tensor (f32, batch by 256000) and a subsequent copy of the same tensor at positions two and three of the whole table, 2.8 milliseconds per step combined, batch-linear, sourced to one line in the runner. Reading the code found the conversion in two places: the runner converts the logits to float32 before calling the sampler, and the sampler converts again internally. The decisive observation came from the sampler’s own ordering: the greedy token is computed by argmax on the original bfloat16 logits before either conversion executes, so under deterministic selection the 32-bit precision feeds only a returned tensor the benchmark never consumes, and since a widening conversion cannot change which element is largest, removing it cannot alter any output.
The patch made both conversion sites conditional on the sampler metadata’s do_sampling flag, a Python-level condition that is static at trace time, so the compiler builds separate graphs per case with no runtime branch: the sampling path stays byte-identical, and only greedy decoding skips the conversions. Application used anchored text replacement into the installed package with a uniqueness assertion on the anchor, the same mechanism as the kernel transplant; both anchors matched exactly once, confirming the code was as read. The prediction was a gain of three to seven percent.
The measurement was the opposite: 11,641 tokens per second on the first patched run and 8,839 on the second, against the 13,470 baseline, with batch time to first token multiplied by five (187 to roughly 995 milliseconds), reproducible on a warm compile cache and unstable between runs, the worst combination. The patch was reverted by the inverse replacement and the baseline reconfirmed. The probable mechanism is that the engine’s precompiled pipeline, including the graphs built during its warmup pass, was specialized around 32-bit logits; changing the returned dtype at one point pushed downstream consumers (output copies, logprob machinery, the runner’s own precompiled sampling graphs) off their precompiled paths on every step. The lesson is stated as a rule for this class of stack: operation-level arithmetic does not compose across just-in-time compilation boundaries, so a cost measured inside one compiled program cannot simply be subtracted by editing that program’s boundary. The inefficiency itself remains real and was documented for the maintainers, where the fix can be made at the pipeline level with all consumers rekeyed together. On a hosted virtual machine, edits to installed packages do not survive an environment rebuild, which makes both the experiment and its revert cheap and fully reversible.
Strategy 7: Attention Kernel Block Size Tuning
The final strategy used the transplant harness of Section 18 for its intended purpose. The target came from the traces: the attention kernel ran at roughly one quarter of memory bandwidth and had grown to the largest cost at scale. The kernel’s compiled name encodes its block configuration, and every trace showed key-value blocks of 2,048 tokens. Reading the vendored source explained why: the default heuristic sizes blocks to carry about 16 megabytes so that each block on its own saturates memory bandwidth, and under this model’s configuration that resolves to 2,048 tokens, the full maximum context. The workload’s sequences average about 140 tokens, so roughly 93 percent of every block was masked padding, loaded and processed for nothing.
The patch added an environment variable override inside the vendored copy’s block-size function for the decode case, leaving stock behavior untouched when the variable is unset. Block sizes are static compile-time arguments, so each override value produces its own compiled kernel, and the change carries two built-in verifications: a print line at every compilation that uses the override, and the compiled kernel’s name itself, which must show the new block size in any trace. Outputs cannot change, because block size affects scheduling, not arithmetic; the generations were confirmed identical to the stock kernel’s.
Measured with the full protocol: in a same-session comparison at 256 concurrent requests, 13,479 tokens per second stock against 16,700 tuned, a gain of 23.9 percent, with latency per token falling from 17.7 to 13.9 milliseconds, replicated across a rebuilt environment within 0.05 percent. A block of 128 tokens measured 16,805, statistically tied with 256, so the optimum is a plateau and only the stock 2,048 was a poor choice for this regime. At 512 concurrent requests the tuned kernel reached 19,511 tokens per second, and the batch-scaling step from 256 to 512 improved from 1.13 to 1.17 times, the expected signature of a cheaper batch-linear term. The campaign total became 12.7 times over the starting configuration, with the energy bound at 10 millijoules per token and 58.5 million tokens per dollar at the throughput maximum.
Two findings about method came with the result. First, profiled runs of the tuned kernel measured near baseline speed and briefly inverted the conclusion: the smaller block multiplies the number of traced kernel invocations roughly eightfold, and per-event tracing overhead becomes visible at that count. The rule of never quoting throughput from a profiled run is what resolved it; the unprofiled comparison and the trace’s kernel names together carry the claim. Second, the tuning is regime-specific. The 16-megabyte heuristic is correct for long contexts that fill the blocks; a long-context deployment should expect no gain and possibly a regression from this change. It is the speculative decoding lesson measured from the winning side: kernels are tuned for a regime, and the regime must be checked against the workload.
Measuring Compile Time: Pallas Versus XLA
The compile-cost comparison in Part I was produced by a standalone script that times the three phases of jax.jit separately: lower (tracing the Python function to an intermediate representation), compile (the backend building the executable), and first execution. The JAX compilation cache was disabled for the measurement, and each trial ran in a fresh process so no cached tracing survived between measurements.
The same attention computation was expressed twice: as high-level operations, two einsums with a float32 softmax between them, compiled by XLA; and as the flash attention Pallas kernel from jax.experimental, lowered through Mosaic, on identical bfloat16 inputs of shape . Over five trials the XLA route compiled in a median of 297.5 milliseconds (range 296.1 to 308.8) and the Pallas route in 50.7 (49.0 to 52.0), a reduction of 83 percent, with lowering costs of 24 and 39 milliseconds respectively and first executions equivalent at 1.1 to 1.8 milliseconds.
The interpretation is bounded: one operation at one shape, measuring developer iteration cost rather than runtime, and consistent with the mechanism that XLA searches a fusion and optimization space over the graph while a Pallas kernel arrives with its implementation decided, leaving only lowering. For kernel development, where the edit, compile, and test loop is traversed hundreds of times, the several-fold difference is the quantity that matters.
Per-Strategy Economics
The business reading is direct. The identical chip-hour, the identical watt budget, and the identical model produce ten times the tokens after two configuration changes, and 12.7 times after one further kernel-level change, which means fleets, energy budgets, and per-token unit costs planned on default configurations are overstated by up to that factor. A service generating one billion tokens per day needs roughly 7.5 chips at the starting configuration and 0.8 at the final one, before redundancy.
Time to first token deserves its own reading, because the strategies moved it in different ways. The selection fix improved it (the prefill’s own logits pass no longer pays the sampler search). Batch scaling raised it steadily, from 84 milliseconds at batch 32 to 307 at 512, because the number is defined here as prefill plus one decode step for the entire batch: more requests mean more prompt tokens through chunked prefill and a heavier first step. The failed precision patch multiplied it by five while claiming to optimize a decode cost, the clearest demonstration in the study that time to first token is a sensitive detector of compilation-path regressions and belongs in every benchmark printout, not only throughput.
Latency prices the choice of operating point: batch 256 holds 17.7 milliseconds per token, suitable for interactive services, while batch 512 trades 79 percent worse per-token latency for 13 percent more throughput and suits machine consumers, agent-to-agent traffic, and offline pipelines. The failed strategies also carry economic content: each was bounded by a single cheap run or a source inspection before it could consume engineering time at scale, and the reverted patch cost one session while producing the study’s most transferable engineering rule.
Reproduction Notes
Reproducing the study requires four artifacts: the notebook’s installation cell (versions resolve at run time and did drift once during the study, from JAX 0.7 to 0.10 with a matching libtpu jump, which the engine’s startup version line detects); the benchmark script with its environment variable switches; the injection module; and the vendored kernel directory created by copying the installed package’s ragged paged attention v3 files. Cell order matters: install, write the project files, copy the kernel sources, then run, with the injection imported before the engine constructor.
The 256-request operating point was replicated in three independent sessions on rebuilt virtual machines at 13,517, 13,470, and 13,484 tokens per second, a spread of one third of one percent; one further session returned 12,991, a 3.6 percent environment dip reported rather than discarded, so replications should expect the headline near 13.5 thousand with occasional session-level variance of a few percent.
Three habits carried the study and are the recommended protocol for any similar effort. First, never quote throughput from a profiled run, and never trust an unprofiled claim about where time goes. Second, pair every treatment with an authenticity check that does not depend on the treatment’s own flag: a memory line, an allocation size, a tensor shape in a trace, a probe print from inside the engine process; this study caught three silently inert configurations and one mislabeled trace by that habit alone.
Third, when aggregating traces, key on tensor shape and check occurrence counts against the expected step count, because operation names are reused across compiled buckets and name-level averages mix programs; that discipline both raised and then correctly killed the padding hypothesis raised in the batch-scaling analysis of Part II. The prompt set, the deterministic selection, the offline batch execution, and the enabled prefix cache define the workload regime; all numbers in this paper are claims about that regime, measured to the hardware’s roofline within it.
Acknowledgments
Google ML Developer Programs and the Google Developers Program supported this work by providing Google Cloud credits.
Top comments (0)