Sequence-on-lane layout
Also known as sequence-on-lane, lane-major KV layout, token-on-lane layout
In plain English
Sequence-on-lane is a KV cache layout that puts a page of tokens along the 128-wide vector lane axis instead of the head dimension, so less memory is wasted on padding.
Technical definition
Sequence-on-lane layout is a TPU KV cache arrangement that maps the tokens of a page onto the 128-lane trailing dimension and the head dimension onto the sublane axis, avoiding the padding that occurs when a small head dimension is placed on the lane axis.
Usable KV pages, FP8, one KV head per device
5,141 to 10,283 pages; +16.5% throughput and -95% median TTFT at concurrency 128 on 8k1k
Engineering details
The TPU vector unit works on tiles whose last dimension is 128 lanes, so any array with a smaller trailing dimension is padded up. The batched attention kernel originally packed keys and values along the head dimension, and for FP8 the packing factor is four. A model with one KV head per device only has two things to pack, so half of every tile was padding. Putting tokens on the lane axis fills tiles regardless of head count, and the head dimension only needs to be a multiple of 32 rather than 128, which lets models with a head dimension of 64 use the kernel.
Why it matters
KV capacity is decode throughput at high concurrency. The layout costs about 3% per-token latency at low concurrency but stops requests from queuing for KV space when the batch is large. It also widens the set of models that run on the shared attention kernel without custom work.
How to read it in InferenceX
In the reported Ironwood configuration usable KV pages doubled from 5,141 to 10,283. At concurrency 128 on 8k1k the extra capacity lifted throughput 16.5% and cut median TTFT by 95% because requests no longer waited for KV space.