Attention Head Calculator

The attention head calculator sizes the multi-head attention block that sits at the heart of every transformer model. Self-attention is the operation that lets a language model decide which earlier tokens matter when predicting the next one, and it is split across several parallel attention heads so the model can attend to different patterns at the same time. This calculator takes the dimensions of that block (hidden size, number of heads, sequence length, batch size, and the attention variant) and returns the projection parameter count, the activation memory footprint, the floating-point operations (FLOPs) per attention pass, and the key-value cache savings you gain from grouped or multi-query attention.

Why does this matter? When you scale a transformer, the attention block is where compute and memory grow quadratically with sequence length. A model that runs comfortably at a 512-token context can run out of GPU memory at 8,192 tokens purely because the attention score matrix is sized as sequence length squared. Knowing these numbers before you launch a training run or deploy an inference server is the difference between a job that fits and a job that crashes with an out-of-memory error. The calculator turns the raw transformer attention formulas into instant, concrete figures.

The tool supports three attention mechanisms used across modern large language models. Multi-Head Attention (MHA) is the original design from the 2017 transformer paper, where every query head has its own dedicated key and value head. Grouped Query Attention (GQA) shares each key-value head across a group of query heads, cutting the key-value cache without the quality loss of going all the way down. Multi-Query Attention (MQA) is the extreme case where a single key-value head serves every query head, giving the largest cache savings. Because the calculator models all three, you can directly compare how a design choice changes parameters, memory, and the KV cache that dominates inference cost.

Everything below explains the exact math the calculator runs, with worked examples that reproduce the on-screen numbers. Use it to plan attention configurations, debug memory limits, or simply understand how transformer attention scales. The attention head calculator weaves together the parameter, memory, FLOPs, and KV-cache views that you would otherwise have to compute by hand across several formulas.

The attention head calculator starts from five inputs: the hidden size d (also called d_model), the number of query heads h, the number of key-value heads kvH (only used for GQA and MQA), the sequence length, and the batch size. From these it derives every result.

The head dimension is simply the hidden size divided by the number of query heads, headDim = d / h. This is the width of the vector each individual head works with, and it is almost always a power-friendly number like 64 or 128. The four learned projection matrices (query, key, value, and output) are each sized from the hidden dimension. In MHA every projection is a square d × d matrix, so each contributes d² parameters and the attention block holds 4d² weights in total. In GQA and MQA the key and value projections shrink because they only produce kvH heads of width headDim, so each becomes d × headDim × kvH parameters instead of d².

The calculator also reports activation memory, which is the temporary tensors created during a forward pass. The attention score matrix is the expensive one: it is shaped batch × heads × sequence × sequence in fp32, so its size scales with the square of the sequence length. The query, key, and value tensors are linear in sequence length. Finally, the FLOPs estimate adds the cost of the query-key product, the softmax, and the weighted value sum. The KV cache savings figure compares the key-value cache an MHA model would need against the smaller cache of your chosen variant, expressed as a percentage reduction.

Inputs and their roles

• Hidden Size (d_model): the embedding width of the model; drives projection parameters.
• Number of Query Heads: how many parallel attention heads the queries use; sets the head dimension.
• Number of KV Heads: how many key-value heads exist (equals query heads for MHA, fewer for GQA, one for MQA).
• Sequence Length: tokens processed at once; the quadratic memory driver.
• Batch Size: sequences processed in parallel; scales all activation memory linearly.

Two costs dominate the attention block: the memory it allocates and the floating-point work it performs. The attention head calculator reports both so you can predict whether a configuration fits on your hardware.

The largest activation is the attention score matrix. After multiplying queries by keys you get a matrix of similarity scores with shape batch × heads × sequence × sequence, and the calculator assumes fp32 (4 bytes). This is why doubling the context length roughly quadruples this tensor. The query tensor is sized batch × heads × sequence × headDim, while the key and value tensors together take twice that. In GQA and MQA the key-value tensors shrink in proportion to the smaller key-value head count, which is exactly where inference savings come from.

The compute side is captured by the GFLOPs figure. The query-key product costs 2 × batch × heads × sequence² × headDim operations, the softmax adds roughly 5 × batch × heads × sequence², and the weighted sum over values costs another 2 × batch × heads × sequence² × headDim. Summing these and dividing by one billion gives GFLOPs. Notice that the two matrix multiplies both scale with the square of the sequence length, which is the well-known quadratic bottleneck of self-attention. Techniques such as FlashAttention reduce the memory traffic of these operations without changing this FLOP count, so the calculator's GFLOPs remain a useful upper-bound estimate for planning.

The single biggest lever in this calculator is the attention type, because it directly controls the key-value cache that dominates the cost of serving large language models. During generation the model must keep the keys and values for every token already produced. That cache grows with sequence length and with the number of key-value heads, so reducing the key-value head count is the cleanest way to make long-context inference affordable.

Multi-Head Attention (MHA) uses one key-value head per query head, giving the richest representation but the largest cache. Grouped Query Attention (GQA) divides the query heads into groups that share a key-value head; with 32 query heads and 8 key-value heads, four queries share each key-value pair, cutting the cache by 75% while keeping most of MHA's quality. Multi-Query Attention (MQA) takes this to the limit with a single key-value head, often saving well over 90% of the cache at the cost of some accuracy. Models like Llama 2 70B and many production chat models adopt GQA precisely because it sits in the sweet spot between cache size and quality.

The calculator quantifies this trade-off with the KV cache savings percentage and the query heads per KV head ratio. Because the key and value projection matrices also shrink under GQA and MQA, you save parameters too, though the bulk of the benefit shows up at inference time as a smaller cache and higher throughput. If you are unsure which variant to pick, start with GQA: it preserves nearly all of MHA's modeling power while delivering most of MQA's memory advantage.

Quick decision guide

• Maximum quality, short context: stay with MHA.
• Long context, production serving: use GQA with 4 to 8 query heads per KV head.
• Extreme memory limits or edge inference: consider MQA and validate the quality drop.

Each output panel in the attention head calculator maps to a real engineering decision. The headline attention parameters figure (in millions) is the weight count of one attention block; multiply it by the number of transformer layers to estimate the attention share of total model size. The head dimension tells you the width each head operates on, and you generally want it to stay at 64 or 128 so the math maps cleanly onto GPU tensor cores.

The GFLOPs value is the compute for a single attention pass at your batch and sequence settings. If you compare two configurations and the GFLOPs jump sharply, it is almost always because you increased the sequence length, which scales the cost quadratically. The parameter breakdown separates the query, key, value, and output projections; under GQA and MQA you will see the key and value rows drop while query and output stay fixed, which visually explains where the savings come from.

The memory usage panel lists the attention score matrix, the query tensor, and the key-value tensors. When the attention score number looks alarmingly large, that is the quadratic term shouting at you, and it is the signal to either shorten the context, reduce the batch, or adopt a memory-efficient attention kernel. Finally, the KV cache savings percentage is your inference cheat sheet: a 75% saving means the serving cache is a quarter of what MHA would need, which roughly translates into supporting four times the batch size or context at the same memory budget. Treat these figures as planning estimates rather than exact production numbers, since real frameworks add overheads and may use lower-precision formats than the fp32 assumed here.