This repository contains the measurement harness used to produce the study "The Sliding Window Premium: Throughput, Goodput, and Per Layer Attention Costs of Hybrid Attention Models at Realistic Context Lengths on TPU v6e" (DOI 10.5281/zenodo.21518443).
The object under measurement is the Ragged Paged Attention (RPA) v3 kernel of the vLLM TPU backend, as exercised by vLLM serving Gemma 4 31B on a TPU v6e slice. Gemma 4 31B is a hybrid attention model: of its sixty attention layers, fifty attend only within a sliding window of 1,024 tokens and ten attend over the full history. The harness measures what the serving stack charges for each of the two layer types at decode time, as a function of history length.
The harness performs four functions.
- It substitutes the vendored RPA v3 kernel for the installed one at import time, so the decode block size can be treated as a controlled variable rather than a value the shipped heuristic selects.
- It runs a campaign of vLLM generations across a history ladder, a batch ladder, and two block size treatments, writing one JSON record per run.
- It reads XProf traces and rolls decode kernel device time up by layer type, separating windowed layers from global layers by compiled kernel name.
- It parses the run records into a table, aggregates repeats, and reports the coefficient of variation across repeats.
The scope of the campaign is one slice size, one model, tensor parallelism 4, greedy decoding, and 128 generated tokens per request.
| Component | Value |
|---|---|
| Instance | ct6e-standard-4t (TPU v6e, four Trillium chips, 2x2 mesh, 32 GB HBM per chip) |
| Provisioning | Flex start, with a Hyperdisk Balanced persistent data disk |
| Zone | us-east5-b (default in provision/create_vm.sh) |
| Image | ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e |
| Python | 3.12, in a uv managed virtual environment |
| Serving stack | vllm-tpu with tpu-inference, JAX, Pallas, Mosaic |
| Model | google/gemma-4-31B, bf16 weights, tensor parallelism 4 |
| KV page size | 128 tokens |
| Layer mix | 50 windowed layers (window 1,024) plus 10 global layers |
| Observability | XProf traces, tensorboard-plugin-profile, tpu-info sampling |
The OS image ships Python 3.10, while current vllm-tpu and tpu-inference pin
JAX versions that publish wheels only for newer interpreters. setup.sh
therefore builds a Python 3.12 environment with uv and places it on
/mnt/data when the persistent disk is mounted, so the environment, the
Hugging Face cache, the results, and the traces survive flex start VM
recreation. On a second and later VM the setup step completes in seconds.
| File | Function |
|---|---|
provision/create_vm.sh |
Creates the flex start v6e-4 VM and attaches the data disk. Run from a local machine. |
provision/startup.sh |
Runs as root at every boot. Formats the data disk on first boot, mounts it at /mnt/data, creates the durable directories. |
setup.sh |
Installs system packages, builds the Python 3.12 environment, installs vllm-tpu, verifies that JAX sees four chips, vendors the RPA v3 kernel into my_rpa/, and applies the block size override patch. |
patch_bkv.py |
Applies the RPA_BKV decode block size override to the vendored kernel. Anchored, idempotent, and asserts anchor uniqueness. |
patch_sw.py |
Applies the diagnostic patch to the sliding window skip logic of the vendored kernel. Provides --find and --check modes for anchor verification before application. |
custom_pallas.py |
Substitutes the vendored kernel functions for the installed ones and rebinds references already captured by other modules. Prints proof of first invocation on the hot path. |
sitecustomize.py |
Applies the substitution in every Python process on PYTHONPATH, including worker processes. Gated on KERNEL_TRANSPLANT=1. |
run_vllm.py |
One benchmark run. Builds the engine, probes KV cache capacity, warms up, measures a batch time to first token proxy, runs the main generation, computes throughput, latency, goodput, and an interchip communication estimate, hashes the generated token streams, and writes one JSON record. |
bench_sw.sh |
The campaign driver. Runs the sequence of configurations under a wall clock budget, tees each log, and synchronizes results to Cloud Storage when GCS_BUCKET is set. |
rpa_stats.py |
Reads an XProf trace and rolls decode kernel time up by layer type, reporting average microseconds per layer per step, share of decode attention, the compiled block sizes present, and the windowed premium. |
parse_results.py |
Parses the JSON records, falls back to log parsing, writes metrics.csv, prints a comparison table, aggregates repeats into mean, standard deviation, and coefficient of variation, and plots throughput against goodput. |
my_rpa/ is created by setup.sh and is not committed. It holds the vendored
copy of the kernel and is the only copy the patches touch. The installed
package is never modified.
Create the persistent data disk:
gcloud compute disks create my-data-disk --size=200GB \
--zone=us-east5-b --type=hyperdisk-balancedFrom a local machine:
bash provision/create_vm.sh
gcloud compute scp --recurse . tpu-v6e-vm:~/tpu_project --zone=us-east5-bVM and ZONE are environment variables of create_vm.sh.
bash ~/tpu_project/setup.shThe step asserts that JAX sees four chips and exits if it does not. It also
locates the installed RPA v3 kernel directory, copies it into my_rpa/, and
applies patch_bkv.py.
export HF_TOKEN=hf_...
export GCS_BUCKET=...
tmux new -s run
bash ~/tpu_project/bench_sw.shThe campaign runs under a wall clock budget of BUDGET_MIN minutes, 140 by
default. When the budget is exhausted the driver skips the remaining runs
rather than truncating one. Each run writes $RESULTS/<tag>.log and
$RESULTS/<tag>.json, where $RESULTS is /mnt/data/results when the data
disk is mounted.
The campaign has three phases.
Phase 1, verification. sw_anchor_bkv128 establishes the session anchor at
batch 256 and block size 128. sw_anchor_profiled and sw_heur_profiled
supply the profiled baselines at batch 64.
Phase 2a, the profiled crossover. Histories of approximately 140, 880, 1,300, 4,000, and 8,000 tokens, at batch 64, 64, 64, 32, and 16 respectively. Batch is reduced at the longest histories to keep the KV cache resident, so absolute microseconds are not comparable across the ladder while the windowed to global ratio within each run is.
Phase 2b, the throughput sweep. Histories of 512, 880, and 1,300 tokens at the batch 256 operating point, unprofiled, plus a heuristic against override contrast at 1,300 tokens.
Phase 3, the intervention. Executed only when my_rpa/kernel.py carries the
SW_PATCH marker. When the marker is absent the driver spends the freed budget
on repeats of the two runs that carry the paper claims.
Phase 3 requires patch_sw.py to have been applied to the vendored kernel:
python3 patch_sw.py my_rpa/kernel.py --find # locate the window logic
python3 patch_sw.py my_rpa/kernel.py --check # verify the anchor matches
python3 patch_sw.py my_rpa/kernel.py # applyThe patch refuses to apply unless the anchor appears exactly once, so upstream
drift produces a failure rather than a silent corruption. The inserted code is
evaluated at trace time against the SW_PATCH environment variable, so it
compiles into a distinct kernel and control runs with the variable unset
execute stock behavior in the same session.
Per layer type rollups from the traces:
python3 rpa_stats.py <profiles>/sw_*_profiled/plugins/profile/*/tpu-v6e-vm.trace.json.gz
python3 rpa_stats.py <trace> --full # add the per kernel tableRun records into a table:
python3 parse_results.py /mnt/data/results --filter 'sw_*'parse_results.py writes metrics.csv and metrics.png next to the records,
flags profiled runs so their throughput is not quoted, flags runs whose goodput
came from the batch level proxy, and marks configurations whose repeats exceed
a coefficient of variation of 1 percent.
Set by bench_sw.sh and readable by run_vllm.py.
| Variable | Default | Function |
|---|---|---|
MODEL |
google/gemma-4-31B |
Model identifier |
TP |
4 | Tensor parallelism |
MAX_MODEL_LEN |
2048 | Context length |
MAX_SEQS |
256 | Concurrency ceiling, the target batch |
BATCH_MULT |
8 | Prompt list multiplier, batch is four times this value |
GPU_MEM_UTIL |
0.85 in the driver, 0.92 at the engine | HBM fraction |
RPA_BKV |
unset | Decode block size request. Unset means the shipped heuristic decides |
SW_PATCH |
unset | Compiles the diagnostic patch into the windowed kernel variant |
PROMPT_LEN |
0 | Pads each prompt to approximately N tokens. 0 selects the short prompts |
PROFILE |
0 | Traces the main generation |
KERNEL_TRANSPLANT |
1 | 0 selects the installed kernel instead of the vendored copy |
SLO_TTFT_MS |
1000 | Goodput objective, time to first token |
SLO_TPOT_MS |
50 | Goodput objective, time per output token |
RATED_W_PER_CHIP |
200 | Third party power rating for the efficiency line. Google publishes no official v6e figure |
BUDGET_MIN |
140 | Campaign wall clock budget |
GCS_BUCKET |
unset | Destination for result synchronization |
The protocol is inherited from the preceding studies in the series and is enforced by the code.
Throughput is never quoted from a profiled run. Tracing overhead scales with
the kernel invocation count and therefore penalizes configurations that invoke
the kernel more often. Profiled runs attribute device time; separate unprofiled
runs supply every throughput number. bench_sw.sh and parse_results.py both
print the warning.
Three levels of configuration are recorded separately, because they disagree.
The requested block size is the value of RPA_BKV. The effective block size is
the value predicted after alignment to the KV page size, recorded as
rpa_bkv_effective. The compiled block size is the value encoded in the kernel
name in the hardware trace. Only the third is ground truth. The JSON records
the first two and rpa_stats.py reads the third.
KV cache capacity is probed at engine initialization and compared against batch times context length. A run whose capacity falls below its demand may preempt and recompute, and is not steady state.
The full generated token stream of every request is hashed into
output_sha256, sorted by request identifier so scheduler completion order
cannot change the digest. This is the correctness gate for the intervention
runs and the instrument behind the determinism result below.
The kernel substitution prints proof of first invocation on the hot path, and the block size override prints the surviving request at compile time.
The numbers below are from the study this harness produced. All per layer values are average device time per layer per decode step at block size 128.
At approximately 140 tokens of history, where the 1,024 token window clips nothing and both layer types read identical data, the fifty windowed layers cost 52.3 microseconds per layer against 34.2 for the ten global layers, a premium of 53.1 percent for identical work. The premium is a property of the compiled sliding window kernel variant rather than of the campaign configuration: it is present under the shipped heuristic at 49.1 percent (144.2 against 96.7 microseconds at the heuristic's 2,048 token blocks), it survived a change of KV cache precision, and it survived the targeted intervention at 51.8 percent. At this history the windowed layers carry 88.5 percent of decode attention time.
The premium declines monotonically with history and then inverts.
| History (tokens) | Batch | Windowed (µs) | Global (µs) | Premium |
|---|---|---|---|---|
| 140 | 64 | 52.3 | 34.2 | +53.1% |
| 880 | 64 | 346.0 | 252.2 | +37.2% |
| 1,300 | 64 | 412.4 | 366.8 | +12.4% |
| 4,000 | 32 | 144.7 | 322.7 | 55.2% |
| 8,000 | 16 | 88.5 | 440.8 | 79.9% |
The ordering inverts between 1,300 and 4,000 tokens. Past the window the windowed layers' cost stops tracking history, while the global layers' cost continues to climb. At 8,000 tokens the ten global layers, one sixth of the model, account for 50.1 percent of decode attention, as much as the fifty windowed layers combined. The 4,000 token row compiled an effective block size of 256 and is caveated as such; the inversion conclusion rests on the 1,300 and 8,000 token points, which compiled block size 128.
Dividing each layer type's time by the number of KV blocks it processes, with the windowed block count capped by the window, places the premium between 37 and 58 percent at every history measured. The sign never changes. The raw inversion is block count arithmetic: past the window a windowed layer performs up to eight times less work at 8,000 tokens, so it wins despite remaining about fifty percent less efficient per unit of work.
The diagnostic patch of patch_sw.py removes the skip index computation from
the windowed variant's per grid step path, the mechanism the source code most
directly suggests. It left the premium intact at both ends of the history
ladder: 51.8 percent against 53.1 in the control at short history, and 79.9
percent in both at 8,000 tokens. At the throughput level the control pair
measured 8,001.5 tokens per second at a coefficient of variation of 0.09
percent and the intervention pair 8,017.2 at 0.16 percent, a difference of 0.19
percent against a noise floor of 0.16. The candidate mechanism is therefore
excluded by experiment. What remains is a per grid step cost within the
compiled sliding window variant, which is a compiler and kernel codegen
quantity rather than a serving configuration.
At 8,000 tokens the two layer types have opposite block size optima. The windowed layers cost 88.5 microseconds at block size 128 against 148.1 under the heuristic's larger blocks, because the window caps their useful data and the rest of a large block is waste. The global layers cost 157.4 microseconds under the heuristic against 440.8 at block size 128, because their data fills large blocks and small blocks multiply per block cost.
| Policy at 8,000 tokens | Windowed (50 layers) | Global (10 layers) | Total |
|---|---|---|---|
| Uniform block size 128 | 50 x 88.5 | 10 x 440.8 | 8,833 µs |
| Uniform heuristic | 50 x 148.1 | 10 x 157.4 | 8,979 µs |
| Mixed | 50 x 88.5 | 10 x 157.4 | 5,999 µs |
Assigning each layer type its measured optimum spends about 32 percent less time in decode attention per step than either uniform configuration. The policy is implementable because the stack already compiles the two layer types as separate kernels, and because the heuristic itself already emits heterogeneous block sizes at 8,000 tokens. The change required is in the block size selection function, not in hardware and not in the model.
The two uniform totals sit within 2 percent of each other. An operator sweeping uniform block sizes at this length would conclude that block size had stopped mattering, while a third of the cost sits between the two policies, visible only per layer type.
At batch 256, unprofiled, with goodput defined by a time to first token objective of 1,000 ms and a time per output token objective of 50 ms, applied as a batch level all or nothing proxy.
| History (tokens) | Treatment | Tokens per second | Median TPOT (ms) | Goodput |
|---|---|---|---|---|
| 140 | Override 128 | 8,001.5 | 28.9 | 8,001.5 |
| 512 | Override 128 | 3,194.4 | 76.0 | 0 |
| 880 | Override 128 | 1,864.2 | 125.0 | 0 |
| 1,300 | Override 128 | 1,860.8 | 131.5 | 0 |
| 1,300 | Heuristic | 2,729.7 | 87.4 | 0 |
Throughput falls steeply to 880 tokens and is then flat to 1,300, which is the end to end signature of window clipping. The goodput cliff moves with history at fixed batch: a 512 token history already exceeds the 50 ms objective at batch 256. A fleet sized at batch 256 from short prompt benchmarks violates its latency objective as soon as transcripts reach a few hundred tokens. At 1,300 tokens the shipped heuristic outperforms the uniform small block by 46.7 percent, an advantage carried almost entirely by the ten global layers.
Two findings revise the preceding studies rather than the hardware.
The kernel's block size selection aligns any request up to the KV page size of 128. Every run labeled block size 16, 32, or 64 in the earlier campaigns therefore compiled a block size of 128. The plateau reported from 16 to 128 was one kernel wearing four labels. No throughput conclusion changes; the description does. This campaign requests 128 explicitly. Discipline proved insufficient once: the 4,000 token run requested 128, its own JSON predicted an effective 128, and the trace shows a compiled 256. Only the trace is ground truth.
Greedy decoding on this platform is deterministic within a compiled program but not bit deterministic across runs. Two stock runs on identical configuration produced token streams that diverge after about forty tokens. Treatment verification logic, which compares a treatment against its control within one program, is unaffected. Claims that depend on exact cross run reproducibility are not: byte identical caching of generations and hash based deduplication of outputs should key on semantic rather than byte identity.
Under a third party rating of 200 W per chip, the short history anchor runs at approximately 10.0 tokens per second per rated watt, an upper bound near 100 mJ per token. At 1,300 tokens of history the same slice under the uniform small block delivers 2.33 tokens per second per rated watt, an upper bound near 430 mJ per token. History multiplies the energy of a token by more than four across this ladder. At an illustrative on demand price of 3.22 dollars per chip hour, the mixed policy's reduction in decode attention corresponds to an estimated 15 to 25 percent end to end reduction in the long history regime, on the order of 0.30 to 0.50 dollars per million long context tokens. The percentage claims are ratios of measured values and are independent of the assumed price and power figures.
The campaign was planned at memory utilization 0.85 and executed at 0.92, because the driver version on the machine did not export the setting. The per run JSON caught the divergence. The correct comparator is the preceding campaign's block size 128 measurement at 0.92, against which the anchors agree within 0.7 percent across campaigns and builds.
The per layer costs at 4,000 and 8,000 tokens come from runs at batch 32 and 16, because the KV cache of 256 concurrent long histories does not fit the slice. The premium ratio is robust within a run and the mixed policy arithmetic uses within run values, but absolute microseconds across the ladder span three batch sizes and are not directly comparable.
The 4,000 token point compiled an effective block size of 256 despite a request and a prediction of 128.
The mixed policy is computed from measured per layer costs. It has not been measured end to end as a deployed configuration. The arithmetic is exact for decode attention; the end to end economic figure is labeled an estimate.
Block counts in the per block normalization are estimated from mid generation history. Goodput uses the batch level proxy, all or nothing per run. The long history rows exceed the time to first token objective by construction, so they are throughput comparisons rather than service level results.
The long prompts are synthetic, a repeated filler sentence after a distinct stem. Several points are single runs. The intervention null bounds the premium's mechanism but does not identify the instruction level cause, which requires compiler level work.
Scope is one slice size, one model, tensor parallelism 4, greedy decoding, and 128 generated tokens per request. The layer mix of five to one with a 1,024 token window is one architecture's choice, and models with different mixes will move every boundary reported here.
Google ML Developer Programs and the Google Developers Program supported this work by providing Google Cloud credits.
R. Zimbres, "From 1,540 to 19,511 Tokens per Second on a Single TPU v5e Chip: A Measurement Study of Large Language Model Inference Optimization", Jul. 06, 2026, Zenodo. doi: 10.5281/zenodo.21221952. https://zenodo.org/records/21221952
R. Zimbres, "Token Velocity on a Single TPU v5e Chip: Per Bucket Prefill and Decode Rates for Gemma 2B and Their Dependence on Kernel Configuration", Jul. 06, 2026, Zenodo. doi: 10.5281/zenodo.21227936. https://zenodo.org/records/21227936
R. Zimbres, "The Decode Block Size Heuristic in TPU Ragged Paged Attention Reduces LLM Inference Throughput by 28 to 69 Percent", Jul. 17, 2026, Zenodo. doi: 10.5281/zenodo.21404155. https://zenodo.org/records/21404155
R. Zimbres, "Batch Scaling and Goodput of a Tuned Attention Kernel on TPU v6e: Throughput, Latency, Energy, and Cost Measurements of vLLM Serving Gemma 4 31B", Zenodo. doi: 10.5281/zenodo.21462837. https://zenodo.org/records/21462837
R. Zimbres, "The Sliding Window Premium: Throughput, Goodput, and Per Layer Attention Costs of Hybrid Attention Models at Realistic Context Lengths on TPU v6e", Zenodo. doi: 10.5281/zenodo.21518443. https://zenodo.org/records/21518443