Model architecture
Attention head dimension
Also known as head dimension, head_dim, head dim
In plain English
Attention head dimension is the vector width used by an attention head, which affects its arithmetic, cache footprint, and fit to accelerator matrix tiles.
Technical definition
An attention head dimension is the feature width of a per-head query, key, or value representation; those widths need not all be equal in every architecture.
Engineering details
Matrix engines compute fixed tile shapes. A kernel may pad, pack, or reshape attention dimensions to fit those tiles. Consequently, a head width smaller than a matrix array can waste work in one layout while a different mapping fills more of the array. Head count, sequence length, precision, sharding, and the query-key versus value dimensions all influence the actual utilization.
Why it matters
A dimension favorable for model quality or a GPU kernel may require a different layout on a TPU. Dividing head dimension by array width is not a universal estimate of model throughput or even of every attention matmul’s utilization.
How to read it in InferenceX
The TPU AgentX article discusses Muse Spark Open’s head dimension of 64 and the fit of model shapes to TPU matrix hardware. Read that example together with its kernel and layout context rather than turning the dimensional ratio into a whole-model performance prediction.
See the concept in real benchmarks
All articlesThe Iron(wood) Throne: Google's TPU Stakes Its Claim on Inference
TPUv7 AgentX results across Muse Spark Open, GLM, Qwen3.5, and Kimi K3, with TorchTPU, megakernels, and llm-d disaggregation.
TPU Inference Externalization Full Steam Ahead
InferenceX, Up to 50% Better Performance per Dollar, Rapid Externalization of TPU stack, Growing Customer Base, Ironwood, TPUv8i, Reducing CUDA Moat