{
  "source_file": "0930",
  "source_platform": "Dev.to",
  "author": "Carlos Chinchilla Corbacho",
  "title_listed": "Attention is a learned weighted average, and the cost is in the square",
  "keywords": "",
  "domain": "大模型算法",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://dev.to/cchinchilladev/attention-is-a-learned-weighted-average-and-the-cost-is-in-the-square-2e0o",
  "publish_date": "2026-09-25",
  "corpus_stage": "",
  "notes": "",
  "id": "ec1afcc3e423",
  "collected_at": "2026-10-02T14:05:16",
  "title": "Attention is a learned weighted average, and the cost is in the square - DEV Community",
  "body_markdown": "Originally published at\ncchinchilla.dev\n. Part 4 of From code to weights, a 12-part series on ML fundamentals for engineers.\nPart 3\nended on a square. Two tensors in one decoder block carry\nT\ntwice,\n[B, H, T, T]\n: a\nT × T\nsquare per head,\nH\nheads,\nB\nsequences. At the widget's defaults they were 89% of the block's memory. This post reads what fills them. It's a weighted average. The average is ordinary; the weights are computed from the input, and that is the new part.\nI wrote attention in four lines and checked it against\nF.scaled_dot_product_attention\n, torch's own implementation. Then I read\nLlamaAttention.forward\nwith the shapes in the margin, and priced it for Llama 3.1 8B at the context lengths people promise.\nThis is part of\na series\nfor engineers moving into ML engineering. Each part names one thing that does not behave the way software behaves. This one is cost: you cannot read it from the code.\nA dictionary that matches a little\nA dictionary lookup is exact. The key matches or it doesn't, and you get one value back. Attention is the same operation with the match made soft. Every key matches a little, and you get back the average of all the values, weighted by how well each key matched.\nThose are attention weights, not parameters. They aren't stored anywhere; the model computes them from the input on every call. Until the cost section, \"weights\" means these. Three learned matrices turn each input row into a query, a key, and a value. The query is compared against every key, and the scores become weights:\nq\n,\nk\n,\nv\n=\nx\n@\nWq\n,\nx\n@\nWk\n,\nx\n@\nWv\n# [B, T, d] each\nscores\n=\nq\n@\nk\n.\ntranspose\n(\n-\n1\n,\n-\n2\n)\n/\nsqrt\n(\nd\n)\n# [B, T, T]\nweights\n=\nscores\n.\nsoftmax\n(\n-\n1\n)\n# every row sums to 1\nout\n=\nweights\n@\nv\n# [B, T, d]\nEnter fullscreen mode\nExit fullscreen mode\nRead the shapes.\nx\nis Part 3's\n[B, T, d]\n. The three projections keep it.\nscores\nis\n[B, T, T]\n: one number for every pair of positions, which is the square. The softmax turns each row into weights that sum to one, and the last line is the average. The four lines agree with\nF.scaled_dot_product_attention\nto\n1e-5\n.\nTwo properties follow from \"average\". Each output coordinate sits between the smallest and largest value at that coordinate; attention never leaves that range. And if you replace a softmax row with a one-hot row, a single 1 and zeros elsewhere, the output is exactly the value at that position. The soft dictionary contains the hard one as a limit.\nWhy the square root\nThe division by\nsqrt(d)\nlooks like a detail. It's the difference between a model that trains from the first step and one that starts saturated. The reason is arithmetic you can check in a shell.\nA dot product of two random vectors of width\nd\n, with unit-variance entries, has variance\nd\n. I measured it: 4,104 at\nd = 4096\n. Feed scores that wide into a softmax and it saturates: one weight goes to nearly 1 and the rest to nearly 0. With 1,024 keys, the largest weight in a row averaged 0.976. The row had 0.09 bits of entropy out of a possible 10. That's a hard lookup wearing a soft one's clothes, at initialisation; trained models sharpen rows on purpose, from a start they can learn from. Divide by\nsqrt(d)\nand the variance is 1.0, the largest weight 0.017, the entropy 9.27 bits. Averaged over the scores, the gradient that reaches them is about an order of magnitude larger: 12 to 18 times across seeds, with the loss read through the average. Through a single weight the ratio swings from a few times to many orders of magnitude, depending on which rows that key happens to win. A saturated softmax has almost nothing to learn from: its gradient is under\n1e-6\nfor 99.8% of the scores, and what's left piles onto the few near-ties.\nLlama writes it once:\nself.scaling = self.head_dim**-0.5\n, at\nmodeling_llama.py:226\n. The width that matters is the head's, not the model's.\nHeads are a view, the mask is a triangle\nMulti-head attention is the four lines above run once per head,\nH\ntimes, on slices of\nd\n. The slicing is free.\nq.view(B, T, H, d // H).transpose(1, 2)\ngives\n[B, H, T, d/H]\nand shares memory with\nq\n. The head axis is a stride. Scores are then\n[B, H, T, T]\n, one square per head. Head 0 of the batched call equals attention run on head 0 alone. What the heads buy is\nH\ndifferent sets of weights over the same positions. Each lives in a slice of width\nd/H\n: 128 in Llama 3.1 8B, with 32 heads.\nThe mask is where a decoder stops being an average over everything. To predict the token after position\nt\n, the model may read positions\n0..t\nand nothing after, so the scores above the diagonal are set to\n-inf\nbefore the softmax and their weights come out as zero. I tested it the direct way: attention over eight positions, then rewrite the keys and values at positions 5 to 7. Positions 0 to 4 returned the same output to\n1e-6\n. Without the mask, all eight changed.\nThe original post has a widget for this that doesn't survive being pasted into a feed:\nswitch between the three masks\nand read the live-cell count change. The third chip, the window, is the last lever in the cost section.\nThe triangle composes with Part 3's padding mask by addition. Left-pad a sequence by two and its first two query rows are all\n-inf\n; a row of all\n-inf\nis\n0/0\nin the softmax, NaN. Same trap, one part later, in the four-line form. transformers fills its eager mask with the dtype's minimum instead (\nmasking_utils.py:602-604\n); those rows come out uniform. Quiet, and still meaningless.\nThe eager path computes the exponentials and their sum in fp32 whatever the model's dtype, then casts the row back:\ndtype=torch.float32\nat\nmodeling_llama.py:208\n. A normaliser summed in bf16's eight significant bits is the wrong place to save memory.\nReading\nLlamaAttention.forward\nmodeling_llama.py:243-281\nin transformers 5.17.0 is the four lines with the heads, the mask, and a cache written out. Condensed, with shapes for the 8B model:\nquery_states\n=\nself\n.\nq_proj\n(\nhidden_states\n).\nview\n(\nhidden_shape\n).\ntranspose\n(\n1\n,\n2\n)\n# [B, 32, T, 128]\nkey_states\n=\nself\n.\nk_proj\n(\nhidden_states\n).\nview\n(\nhidden_shape\n).\ntranspose\n(\n1\n,\n2\n)\n# [B, 8, T, 128]\nvalue_states\n=\nself\n.\nv_proj\n(\nhidden_states\n).\nview\n(\nhidden_shape\n).\ntranspose\n(\n1\n,\n2\n)\n# [B, 8, T, 128]\nquery_states\n,\nkey_states\n=\napply_rotary_pos_emb\n(\nquery_states\n,\nkey_states\n,\ncos\n,\nsin\n)\nif\npast_key_values\nis\nnot\nNone\n:\nkey_states\n,\nvalue_states\n=\npast_key_values\n.\nupdate\n(\nkey_states\n,\nvalue_states\n,\nself\n.\nlayer_idx\n)\nattn_output\n,\nattn_weights\n=\nattention_interface\n(\nself\n,\nquery_states\n,\nkey_states\n,\nvalue_states\n,\nattention_mask\n,\nscaling\n=\nself\n.\nscaling\n)\nattn_output\n=\nself\n.\no_proj\n(\nattn_output\n.\nreshape\n(\n*\ninput_shape\n,\n-\n1\n))\nEnter fullscreen mode\nExit fullscreen mode\nThree things in it aren't in the four lines.\npast_key_values.update\nis the cache. At generation time the keys and values of every earlier token are kept, so each new token computes one query row and reads\nT\nstored keys.\nThe key and value have 8 heads where the query has 32:\nk_proj\nmaps 4,096 to 1,024. That's grouped-query attention (GQA), against one key-value head per query head (MHA). It exists for the cache. And\nattention_interface\nis a dispatch whose fallback is\neager_attention_forward\n,\nmodeling_llama.py:191-213\n: the name means the four lines run as written, plus\nrepeat_kv\n, which copies the 8 key heads to 32 at runtime so the matmul lines up. Set\n_attn_implementation\nto\nsdpa\n, the\nscaled_dot_product_attention\nfrom the first section, or to\nflash_attention_2\n, and a different kernel (the routine that runs the arithmetic) computes the same maths. The rotary line is Part 7's.\nThe cost is in the data, not the code\nIn software you read a function's cost from its code. A loop over the input is linear; two nested loops are quadratic; you can tell before it runs. Attention is one function; its code is the four lines. There is no loop in them. The square is inside a single\n@\n, and nothing in the code's shape says quadratic. Its cost is set by\nT\nin the request: positions squared, per head, per layer.\nTake Llama 3.1 8B's shape. The scores for one sequence in one layer, in bf16, are 268 MB at 2,048 tokens, 4.3 GB at 8,192, and 1.1 TB at the 131,072 the config declares; the eager path holds three such squares at its peak, two of them in fp32: about five times those bytes. The weights of the whole model are 16 GB. Nobody materialises that last number. The kernel is why.\nThe second widget on\nthe original post\nis that multiplication, live: drag\nT\nto 16k with\neager\nselected and one layer's scores outgrow the whole model's weights; at 131k they are about seventy times larger. Switch to\nsdpa\nand the number is struck out — never allocated.\nI timed the last three lines on my laptop's CPU with 8 heads of width 64. From 512 tokens to 8,192, sixteen times the input, the eager form went from 1.9 ms to 284 ms: 150 times. A linear layer over the same input went from 0.23 ms to 2.5 ms, eleven times. Doubling\nT\nfrom 4,096 to 8,192 multiplied the linear layer by 1.95 and attention by 4.3: the square, within timing noise.\nThree levers move the number. The first is the kernel.\nF.scaled_dot_product_attention\nwith its flash or memory-efficient backend computes the same numbers in blocks and never holds the\nT × T\nmatrix; its math fallback still does. At 4,096 tokens, in fp32, the profiler charged 545 MB to the eager form's matmuls, the 537 MB square plus its output, and 8.5 MB to sdpa, and it ran 2.4 times faster.\nFlash attention is that idea on a GPU. The arithmetic is still quadratic. What changes is what must exist in memory at once. The line that picks it is one config field,\n_attn_implementation\n. In 5.17.0 the default at load is\nsdpa\n;\neager\nis what you get by asking for it. With padding in the batch, sdpa still holds the mask itself,\n[B, 1, T, T]\nbooleans, 17 MB at 4,096; without padding, transformers passes no mask and lets the kernel apply the triangle (\nis_causal=True\n).\nThe other two change what's computed. Grouped-query attention shrinks the cache. With 8 key-value heads the 8B model stores 131 kB per token across its 32 layers, 17 GB for a 128k context, already a little more than the weights, and with 32 heads it would be 524 kB per token and 69 GB, four times the weights. A sliding window caps how far back a query reads. Llama 3.1 has none; Mistral 7B v0.1 had one, and Gemma 2 and 3 interleave windowed layers with global ones. transformers writes it as\nkv_idx > q_idx - sliding_window\n, at\nmasking_utils.py:100\n. At 32,768 tokens a 4,096 window leaves 126 million live cells where the causal triangle has 537 million. Below the window the cost is a triangle. Above it, a band of width\nW\n. A band grows with\nT\n, not its square.\nBe fair about the square. Count operations rather than bytes, over the prompt pass, and the linear layers of the 8B model do more of them than attention until roughly 30,000 tokens, about twice that if you count only the triangle a causal kernel computes. At 8,192, attention is about a fifth of the block's floating-point operations, its FLOPs, counting the full square, and about an eighth counting the triangle.\nThe square is a memory problem long before it's a compute problem. The number to compute before promising a context length is the memory one.\nWhat the ML engineer watches\nFour things, and the first is a multiplication.\nPrice\n[B, H, T, T]\nand the cache before promising a context.\nHeads times\nT\nsquared times two bytes, per layer. Add the cache per token times the longest sequence times the batch. Do it at the batch you'll serve. The config's\nmax_position_embeddings\nis what the model accepts after its rotary scaling (131,072 here: an 8,192 pre-training window, scaled and then trained out to 128k), not what your card can hold.\nKnow which kernel runs.\neager\nreturns the weights and materialises the square;\nsdpa\nand flash don't. The choice is\n_attn_implementation\n, and it decides whether the 545 MB above exists. A debugging tool that wants the weights can't get them from\nsdpa\n:\noutput_attentions=True\nreturns an empty tuple there, and the fix, loading with\nattn_implementation=\"eager\"\n, materialises the square on every call.\nCompose the masks on purpose.\nCausal, padding, and window are ANDed; added, in the\n-inf\nform. A query row with nothing left to read is NaN in that form and uniform in the library's, and neither is a number you want. Part 3's habit applies twice now: the last-token index comes from the mask, and the mask travels with the batch.\nMeasure latency at the longest prompt you'll accept.\nThe linear layers scale with\nT\n. Attention scales with its square. The cache read at each generated token scales with everything already in it. A p99 measured at 500 tokens says nothing about 8,000. The 1.9 ms to 284 ms above is the shape of the surprise.\nNone of this is modelling work either. A weighted average is four lines, and the weights come from three learned matrices. The square they live in is the first cost in this series that the input sets and the code doesn't.\nPart 4 of twelve. Attention is a weighted average whose weights the model computes from its own input: the average is four lines, the\nT × T\nmatrix of weights is the cost, and the kernel decides whether that square is ever held in memory. Attention is one stage. The next part reads the whole block around it — six tensors, and a residual stream that every stage adds to — and lands on\nthe series index\nwhen it's out. If you arrived here directly,\nPart 1\nreads the tokenizer,\nPart 2\nthe table a token id indexes, and\nPart 3\nthe shape both of them hand to this one.\nWritten with AI assistance. See\nAI policy\n.",
  "metadata": {
    "title": "Attention is a learned weighted average, and the cost is in the square - DEV Community",
    "description": "Originally published at cchinchilla.dev. Part 4 of From code to weights, a 12-part series on ML... Tagged with machinelearning, python, ai, beginners.",
    "keywords": "machinelearning, python, ai, beginners, software, coding, development, engineering, inclusive, community",
    "author": null,
    "og:type": "article",
    "og:url": "https://dev.to/cchinchilladev/attention-is-a-learned-weighted-average-and-the-cost-is-in-the-square-2e0o",
    "og:title": "Attention is a learned weighted average, and the cost is in the square",
    "og:description": "Originally published at cchinchilla.dev. Part 4 of From code to weights, a 12-part series on ML...",
    "og:site_name": "DEV Community",
    "og:image": "https://media2.dev.to/dynamic/image/width=1200,height=627,fit=cover,gravity=auto,format=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2F704ek1ngc179kb1vpwum.png",
    "twitter:site": "@thepracticaldev",
    "twitter:creator": "@",
    "twitter:title": "Attention is a learned weighted average, and the cost is in the square",
    "twitter:description": "Originally published at cchinchilla.dev. Part 4 of From code to weights, a 12-part series on ML...",
    "twitter:card": "summary_large_image",
    "twitter:widgets:new-embed-design": "on",
    "twitter:image:src": "https://media2.dev.to/dynamic/image/width=1200,height=627,fit=cover,gravity=auto,format=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2F704ek1ngc179kb1vpwum.png"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/ec1afcc3e423/code/body-000.txt",
      "chars": 240,
      "language": "text"
    },
    {
      "file": "attachments/ec1afcc3e423/code/body-001.txt",
      "chars": 743,
      "language": "text"
    }
  ],
  "images": [],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/",
      "label": "cchinchilla.dev",
      "type": "a"
    },
    {
      "href": "https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/",
      "label": "Part 3",
      "type": "a"
    },
    {
      "href": "https://cchinchilla.dev/series/",
      "label": "a series",
      "type": "a"
    },
    {
      "href": "https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/",
      "label": "Part 1",
      "type": "a"
    },
    {
      "href": "https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/",
      "label": "Part 2",
      "type": "a"
    },
    {
      "href": "https://cchinchilla.dev/ai-policy/",
      "label": "AI policy",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:49:56Z",
  "source_url": "https://dev.to/cchinchilladev/attention-is-a-learned-weighted-average-and-the-cost-is-in-the-square-2e0o",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 13607,
  "expected_body_images": 0,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/ec1afcc3e423/pages/source.html",
  "local_html": "attachments/ec1afcc3e423/pages/article.html",
  "body_base_path": "attachments/ec1afcc3e423/pages/",
  "body_markdown_path": "attachments/ec1afcc3e423/pages/article.md"
}