{
  "schemaVersion": 1,
  "slug": "grouped-query",
  "title": "Eight queries, fewer key/value heads",
  "className": "HeadSharingComparison",
  "question": "Which attention heads can share a key/value cache?",
  "description": "Compare MHA, GQA and MQA while keeping eight query heads. Follow the head mapping and calculate compact KV-cache storage.",
  "category": "Modern LLMs",
  "tags": [
    "GQA",
    "MQA",
    "Multi-head attention",
    "KV cache"
  ],
  "curator": "TensorViz",
  "attribution": "Original TensorViz teaching example. PyTorch provides the underlying operators.",
  "license": "No separate redistribution license has been declared for these project examples.",
  "limitations": "Untrained width-16 comparison. Averaging K/V weights is an illustrative conversion without the paper's uptraining. Cache estimates exclude temporary repeated tensors, attention activations and allocator overhead; no quality or latency claim.",
  "openingStep": "groups",
  "provenance": {
    "capturedAt": "2026-09-17T06:01:14.902248+00:00",
    "productRevision": "cf0ab1c9c58c1dc640235b742fa032cade00b43c",
    "sourceSha256": "2d7dbb8bead241d951d07146b0a15f05cd19c7ae0d0bb5a90bf1ad1dd974aa16",
    "python": "3.13.13",
    "torch": "2.7.1",
    "execution": {
      "inputs": {
        "args": [
          {
            "shape": [
              1,
              4,
              16
            ],
            "dtype": "float32",
            "fill": "randn"
          }
        ]
      },
      "device": "cpu",
      "mode": "eval",
      "seed": 0,
      "deterministic": true
    },
    "outcome": "passed",
    "outputShape": [
      3,
      1,
      4,
      16
    ],
    "method": "TensorViz multifile projection and isolated execution worker",
    "scope": "Reduced causal MHA/GQA/MQA CPU forwards, independent per-head cache calculation, and duplicated-weight equivalence. Analytical cache estimates are not runtime measurements.",
    "numericalChecks": [
      "MHA with duplicated K/V weights matches grouped attention for 8, 2 and 1 KV heads",
      "Future token changes cannot affect earlier causal outputs",
      "Compact per-group K/V calculation reproduces the final query output",
      "Analytical KV bytes scale with the number of KV heads"
    ]
  },
  "layers": [
    {
      "id": "input",
      "label": "input",
      "op": "Input",
      "kind": "input",
      "parent": null,
      "params": {
        "shape": "1,4,16"
      }
    },
    {
      "id": "mha",
      "label": "mha",
      "op": "Module",
      "kind": "module",
      "parent": null,
      "params": {
        "arg_0": 8
      },
      "source": {
        "line": 39,
        "endLine": 39
      }
    },
    {
      "id": "gqa",
      "label": "gqa",
      "op": "Module",
      "kind": "module",
      "parent": null,
      "params": {
        "arg_0": 2
      },
      "source": {
        "line": 40,
        "endLine": 40
      }
    },
    {
      "id": "mqa",
      "label": "mqa",
      "op": "Module",
      "kind": "module",
      "parent": null,
      "params": {
        "arg_0": 1
      },
      "source": {
        "line": 41,
        "endLine": 41
      }
    },
    {
      "id": "stack@HeadSharingComparison#0",
      "label": "stack@HeadSharingComparison#0",
      "op": "torch.stack",
      "kind": "layer",
      "parent": null,
      "params": {},
      "source": {
        "line": 44,
        "endLine": 44
      }
    },
    {
      "id": "output",
      "label": "output",
      "op": "Output",
      "kind": "output",
      "parent": null,
      "params": {}
    },
    {
      "id": "python@mha#5",
      "label": "python@mha#5",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 29,
        "endLine": 29
      }
    },
    {
      "id": "mha.query",
      "label": "mha.query",
      "op": "Linear",
      "kind": "layer",
      "parent": "mha",
      "params": {
        "in_features": 16,
        "out_features": 16,
        "bias": false
      },
      "source": {
        "line": 15,
        "endLine": 15
      }
    },
    {
      "id": "python@mha#0",
      "label": "python@mha#0",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 23,
        "endLine": 23
      }
    },
    {
      "id": "mha.key",
      "label": "mha.key",
      "op": "Linear",
      "kind": "layer",
      "parent": "mha",
      "params": {
        "in_features": 16,
        "bias": false
      },
      "source": {
        "line": 16,
        "endLine": 16
      }
    },
    {
      "id": "python@mha#1",
      "label": "python@mha#1",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 24,
        "endLine": 24
      }
    },
    {
      "id": "python@mha#3",
      "label": "python@mha#3",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 26,
        "endLine": 26
      }
    },
    {
      "id": "div@mha#0",
      "label": "div@mha#0",
      "op": "torch.div",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 28,
        "endLine": 28
      }
    },
    {
      "id": "mha.value",
      "label": "mha.value",
      "op": "Linear",
      "kind": "layer",
      "parent": "mha",
      "params": {
        "in_features": 16,
        "bias": false
      },
      "source": {
        "line": 17,
        "endLine": 17
      }
    },
    {
      "id": "python@mha#2",
      "label": "python@mha#2",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 25,
        "endLine": 25
      }
    },
    {
      "id": "python@mha#4",
      "label": "python@mha#4",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 27,
        "endLine": 27
      }
    },
    {
      "id": "mha.softmax",
      "label": "mha.softmax",
      "op": "Softmax",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 18,
        "endLine": 18
      }
    },
    {
      "id": "python@mha#6",
      "label": "python@mha#6",
      "op": "Python",
      "kind": "layer",
      "parent": "mha",
      "params": {},
      "source": {
        "line": 31,
        "endLine": 31
      }
    },
    {
      "id": "mha.project",
      "label": "mha.project",
      "op": "Linear",
      "kind": "layer",
      "parent": "mha",
      "params": {
        "in_features": 16,
        "out_features": 16,
        "bias": false
      },
      "source": {
        "line": 19,
        "endLine": 19
      }
    },
    {
      "id": "python@gqa#5",
      "label": "python@gqa#5",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 29,
        "endLine": 29
      }
    },
    {
      "id": "gqa.query",
      "label": "gqa.query",
      "op": "Linear",
      "kind": "layer",
      "parent": "gqa",
      "params": {
        "in_features": 16,
        "out_features": 16,
        "bias": false
      },
      "source": {
        "line": 15,
        "endLine": 15
      }
    },
    {
      "id": "python@gqa#0",
      "label": "python@gqa#0",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 23,
        "endLine": 23
      }
    },
    {
      "id": "gqa.key",
      "label": "gqa.key",
      "op": "Linear",
      "kind": "layer",
      "parent": "gqa",
      "params": {
        "in_features": 16,
        "bias": false
      },
      "source": {
        "line": 16,
        "endLine": 16
      }
    },
    {
      "id": "python@gqa#1",
      "label": "python@gqa#1",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 24,
        "endLine": 24
      }
    },
    {
      "id": "python@gqa#3",
      "label": "python@gqa#3",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 26,
        "endLine": 26
      }
    },
    {
      "id": "div@gqa#0",
      "label": "div@gqa#0",
      "op": "torch.div",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 28,
        "endLine": 28
      }
    },
    {
      "id": "gqa.value",
      "label": "gqa.value",
      "op": "Linear",
      "kind": "layer",
      "parent": "gqa",
      "params": {
        "in_features": 16,
        "bias": false
      },
      "source": {
        "line": 17,
        "endLine": 17
      }
    },
    {
      "id": "python@gqa#2",
      "label": "python@gqa#2",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 25,
        "endLine": 25
      }
    },
    {
      "id": "python@gqa#4",
      "label": "python@gqa#4",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 27,
        "endLine": 27
      }
    },
    {
      "id": "gqa.softmax",
      "label": "gqa.softmax",
      "op": "Softmax",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 18,
        "endLine": 18
      }
    },
    {
      "id": "python@gqa#6",
      "label": "python@gqa#6",
      "op": "Python",
      "kind": "layer",
      "parent": "gqa",
      "params": {},
      "source": {
        "line": 31,
        "endLine": 31
      }
    },
    {
      "id": "gqa.project",
      "label": "gqa.project",
      "op": "Linear",
      "kind": "layer",
      "parent": "gqa",
      "params": {
        "in_features": 16,
        "out_features": 16,
        "bias": false
      },
      "source": {
        "line": 19,
        "endLine": 19
      }
    },
    {
      "id": "python@mqa#5",
      "label": "python@mqa#5",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 29,
        "endLine": 29
      }
    },
    {
      "id": "mqa.query",
      "label": "mqa.query",
      "op": "Linear",
      "kind": "layer",
      "parent": "mqa",
      "params": {
        "in_features": 16,
        "out_features": 16,
        "bias": false
      },
      "source": {
        "line": 15,
        "endLine": 15
      }
    },
    {
      "id": "python@mqa#0",
      "label": "python@mqa#0",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 23,
        "endLine": 23
      }
    },
    {
      "id": "mqa.key",
      "label": "mqa.key",
      "op": "Linear",
      "kind": "layer",
      "parent": "mqa",
      "params": {
        "in_features": 16,
        "bias": false
      },
      "source": {
        "line": 16,
        "endLine": 16
      }
    },
    {
      "id": "python@mqa#1",
      "label": "python@mqa#1",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 24,
        "endLine": 24
      }
    },
    {
      "id": "python@mqa#3",
      "label": "python@mqa#3",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 26,
        "endLine": 26
      }
    },
    {
      "id": "div@mqa#0",
      "label": "div@mqa#0",
      "op": "torch.div",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 28,
        "endLine": 28
      }
    },
    {
      "id": "mqa.value",
      "label": "mqa.value",
      "op": "Linear",
      "kind": "layer",
      "parent": "mqa",
      "params": {
        "in_features": 16,
        "bias": false
      },
      "source": {
        "line": 17,
        "endLine": 17
      }
    },
    {
      "id": "python@mqa#2",
      "label": "python@mqa#2",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 25,
        "endLine": 25
      }
    },
    {
      "id": "python@mqa#4",
      "label": "python@mqa#4",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 27,
        "endLine": 27
      }
    },
    {
      "id": "mqa.softmax",
      "label": "mqa.softmax",
      "op": "Softmax",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 18,
        "endLine": 18
      }
    },
    {
      "id": "python@mqa#6",
      "label": "python@mqa#6",
      "op": "Python",
      "kind": "layer",
      "parent": "mqa",
      "params": {},
      "source": {
        "line": 31,
        "endLine": 31
      }
    },
    {
      "id": "mqa.project",
      "label": "mqa.project",
      "op": "Linear",
      "kind": "layer",
      "parent": "mqa",
      "params": {
        "in_features": 16,
        "out_features": 16,
        "bias": false
      },
      "source": {
        "line": 19,
        "endLine": 19
      }
    }
  ],
  "connections": [
    {
      "id": "edge-1",
      "source": "mha.query",
      "target": "python@mha#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-2",
      "source": "mha.key",
      "target": "python@mha#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-3",
      "source": "mha.value",
      "target": "python@mha#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-4",
      "source": "python@mha#1",
      "target": "python@mha#3",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-5",
      "source": "python@mha#2",
      "target": "python@mha#4",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-6",
      "source": "python@mha#0",
      "target": "div@mha#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-7",
      "source": "python@mha#3",
      "target": "div@mha#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-8",
      "source": "div@mha#0",
      "target": "mha.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-9",
      "source": "python@mha#5",
      "target": "mha.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-10",
      "source": "mha.softmax",
      "target": "python@mha#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-11",
      "source": "python@mha#4",
      "target": "python@mha#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-12",
      "source": "python@mha#6",
      "target": "mha.project",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-13",
      "source": "gqa.query",
      "target": "python@gqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-14",
      "source": "gqa.key",
      "target": "python@gqa#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-15",
      "source": "gqa.value",
      "target": "python@gqa#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-16",
      "source": "python@gqa#1",
      "target": "python@gqa#3",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-17",
      "source": "python@gqa#2",
      "target": "python@gqa#4",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-18",
      "source": "python@gqa#0",
      "target": "div@gqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-19",
      "source": "python@gqa#3",
      "target": "div@gqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-20",
      "source": "div@gqa#0",
      "target": "gqa.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-21",
      "source": "python@gqa#5",
      "target": "gqa.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-22",
      "source": "gqa.softmax",
      "target": "python@gqa#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-23",
      "source": "python@gqa#4",
      "target": "python@gqa#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-24",
      "source": "python@gqa#6",
      "target": "gqa.project",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-25",
      "source": "mqa.query",
      "target": "python@mqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-26",
      "source": "mqa.key",
      "target": "python@mqa#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-27",
      "source": "mqa.value",
      "target": "python@mqa#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-28",
      "source": "python@mqa#1",
      "target": "python@mqa#3",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-29",
      "source": "python@mqa#2",
      "target": "python@mqa#4",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-30",
      "source": "python@mqa#0",
      "target": "div@mqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-31",
      "source": "python@mqa#3",
      "target": "div@mqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-32",
      "source": "div@mqa#0",
      "target": "mqa.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-33",
      "source": "python@mqa#5",
      "target": "mqa.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-34",
      "source": "mqa.softmax",
      "target": "python@mqa#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-35",
      "source": "python@mqa#4",
      "target": "python@mqa#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-36",
      "source": "python@mqa#6",
      "target": "mqa.project",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-37",
      "source": "input",
      "target": "mha.query",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-38",
      "source": "input",
      "target": "python@mha#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-39",
      "source": "input",
      "target": "mha.key",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-40",
      "source": "input",
      "target": "python@mha#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-41",
      "source": "input",
      "target": "mha.value",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-42",
      "source": "input",
      "target": "python@mha#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-43",
      "source": "input",
      "target": "python@mha#5",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-44",
      "source": "input",
      "target": "python@mha#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-45",
      "source": "input",
      "target": "mha",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-46",
      "source": "input",
      "target": "gqa.query",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-47",
      "source": "input",
      "target": "python@gqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-48",
      "source": "input",
      "target": "gqa.key",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-49",
      "source": "input",
      "target": "python@gqa#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-50",
      "source": "input",
      "target": "gqa.value",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-51",
      "source": "input",
      "target": "python@gqa#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-52",
      "source": "input",
      "target": "python@gqa#5",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-53",
      "source": "input",
      "target": "python@gqa#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-54",
      "source": "input",
      "target": "gqa",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-55",
      "source": "input",
      "target": "mqa.query",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-56",
      "source": "input",
      "target": "python@mqa#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-57",
      "source": "input",
      "target": "mqa.key",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-58",
      "source": "input",
      "target": "python@mqa#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-59",
      "source": "input",
      "target": "mqa.value",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-60",
      "source": "input",
      "target": "python@mqa#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-61",
      "source": "input",
      "target": "python@mqa#5",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-62",
      "source": "input",
      "target": "python@mqa#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-63",
      "source": "input",
      "target": "mqa",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-64",
      "source": "mha.project",
      "target": "stack@HeadSharingComparison#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-65",
      "source": "mha",
      "target": "stack@HeadSharingComparison#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-66",
      "source": "gqa.project",
      "target": "stack@HeadSharingComparison#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-67",
      "source": "gqa",
      "target": "stack@HeadSharingComparison#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-68",
      "source": "mqa.project",
      "target": "stack@HeadSharingComparison#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-69",
      "source": "mqa",
      "target": "stack@HeadSharingComparison#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-70",
      "source": "stack@HeadSharingComparison#0",
      "target": "output",
      "sourcePort": null,
      "targetPort": null
    }
  ],
  "steps": [
    {
      "id": "heads",
      "target": "mha",
      "title": "Give every query its own K/V head",
      "note": "MHA projects eight distinct query, key and value heads. At four input tokens, each causal head has a four-by-four attention map."
    },
    {
      "id": "groups",
      "target": "gqa.key",
      "title": "Share within two groups",
      "note": "GQA projects only two key heads and two value heads. Queries 0–3 consult KV 0; queries 4–7 consult KV 1. The query projection still produces all eight heads."
    },
    {
      "id": "single",
      "target": "mqa.value",
      "title": "Use one shared key/value head",
      "note": "MQA is the single-KV-head endpoint. Every query has its own attention distribution, even though the key/value contents are shared."
    },
    {
      "id": "output",
      "target": "gqa.project",
      "title": "Rejoin all eight query outputs",
      "note": "Each query head returns two features. Concatenating eight outputs recovers width 16 before the output projection. The comparison below separates compact cache bytes from this reference implementation's repeated temporary tensors."
    }
  ],
  "milestone": {
    "name": "MHA / MQA / GQA",
    "year": 2023,
    "kind": "Component",
    "fidelity": "Eight query heads of width two with 8, 2 or 1 KV heads. Transparent repeat_interleave CPU implementation; no optimized attention kernel.",
    "papers": [
      {
        "title": "Attention Is All You Need (2017)",
        "url": "https://arxiv.org/abs/1706.03762"
      },
      {
        "title": "Fast Transformer Decoding: One Write-Head is All You Need (2019)",
        "url": "https://arxiv.org/abs/1911.02150"
      },
      {
        "title": "GQA: Training Generalized Multi-Query Transformer Models (2023)",
        "url": "https://arxiv.org/abs/2305.13245"
      }
    ]
  },
  "experiment": {
    "kind": "routing",
    "title": "Keep queries, share keys and values.",
    "description": "Head connections from the recorded configuration; cache bytes calculated as 2 × batch × tokens × KV heads × head width × bytes per value. This CPU reference does not measure decoding speed.",
    "controlLabel": "Head and cache configuration",
    "cases": [
      {
        "id": "mha-128",
        "label": "MHA · 8 KV heads · 128 cached tokens",
        "note": "All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.",
        "target": "mha",
        "routing": {
          "sourceLabel": "Query heads",
          "destinationLabel": "Shared key/value heads",
          "sources": [
            "Q 0",
            "Q 1",
            "Q 2",
            "Q 3",
            "Q 4",
            "Q 5",
            "Q 6",
            "Q 7"
          ],
          "destinations": [
            "KV 0",
            "KV 1",
            "KV 2",
            "KV 3",
            "KV 4",
            "KV 5",
            "KV 6",
            "KV 7"
          ],
          "weights": [
            [
              1,
              0,
              0,
              0,
              0,
              0,
              0,
              0
            ],
            [
              0,
              1,
              0,
              0,
              0,
              0,
              0,
              0
            ],
            [
              0,
              0,
              1,
              0,
              0,
              0,
              0,
              0
            ],
            [
              0,
              0,
              0,
              1,
              0,
              0,
              0,
              0
            ],
            [
              0,
              0,
              0,
              0,
              1,
              0,
              0,
              0
            ],
            [
              0,
              0,
              0,
              0,
              0,
              1,
              0,
              0
            ],
            [
              0,
              0,
              0,
              0,
              0,
              0,
              1,
              0
            ],
            [
              0,
              0,
              0,
              0,
              0,
              0,
              0,
              1
            ]
          ]
        },
        "vectors": [
          {
            "label": "Last token output · four-token forward",
            "values": [
              -0.21111758544348944,
              0.11683543780294618,
              0.5017139662611414,
              -0.1425208747346436,
              0.18604252478128183,
              -0.21919487167003918,
              -0.3739623315636151,
              -0.005229620135255769,
              -0.00484164277273208,
              0.06731108802342632,
              -0.06666299333829089,
              -0.01772044374467424,
              0.11593483275768837,
              -0.06246878466753765,
              0.12209215737486764,
              -0.13132224575309784
            ]
          }
        ],
        "metrics": [
          {
            "label": "Analytical K + V cache · bytes",
            "value": 16384
          },
          {
            "label": "KV heads",
            "value": 8
          },
          {
            "label": "Query heads per KV head",
            "value": 1
          }
        ]
      },
      {
        "id": "mha-2048",
        "label": "MHA · 8 KV heads · 2048 cached tokens",
        "note": "All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.",
        "target": "mha",
        "routing": {
          "sourceLabel": "Query heads",
          "destinationLabel": "Shared key/value heads",
          "sources": [
            "Q 0",
            "Q 1",
            "Q 2",
            "Q 3",
            "Q 4",
            "Q 5",
            "Q 6",
            "Q 7"
          ],
          "destinations": [
            "KV 0",
            "KV 1",
            "KV 2",
            "KV 3",
            "KV 4",
            "KV 5",
            "KV 6",
            "KV 7"
          ],
          "weights": [
            [
              1,
              0,
              0,
              0,
              0,
              0,
              0,
              0
            ],
            [
              0,
              1,
              0,
              0,
              0,
              0,
              0,
              0
            ],
            [
              0,
              0,
              1,
              0,
              0,
              0,
              0,
              0
            ],
            [
              0,
              0,
              0,
              1,
              0,
              0,
              0,
              0
            ],
            [
              0,
              0,
              0,
              0,
              1,
              0,
              0,
              0
            ],
            [
              0,
              0,
              0,
              0,
              0,
              1,
              0,
              0
            ],
            [
              0,
              0,
              0,
              0,
              0,
              0,
              1,
              0
            ],
            [
              0,
              0,
              0,
              0,
              0,
              0,
              0,
              1
            ]
          ]
        },
        "vectors": [
          {
            "label": "Last token output · four-token forward",
            "values": [
              -0.21111758544348944,
              0.11683543780294618,
              0.5017139662611414,
              -0.1425208747346436,
              0.18604252478128183,
              -0.21919487167003918,
              -0.3739623315636151,
              -0.005229620135255769,
              -0.00484164277273208,
              0.06731108802342632,
              -0.06666299333829089,
              -0.01772044374467424,
              0.11593483275768837,
              -0.06246878466753765,
              0.12209215737486764,
              -0.13132224575309784
            ]
          }
        ],
        "metrics": [
          {
            "label": "Analytical K + V cache · bytes",
            "value": 262144
          },
          {
            "label": "KV heads",
            "value": 8
          },
          {
            "label": "Query heads per KV head",
            "value": 1
          }
        ]
      },
      {
        "id": "gqa-128",
        "label": "GQA · 2 KV heads · 128 cached tokens",
        "note": "All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.",
        "target": "gqa",
        "routing": {
          "sourceLabel": "Query heads",
          "destinationLabel": "Shared key/value heads",
          "sources": [
            "Q 0",
            "Q 1",
            "Q 2",
            "Q 3",
            "Q 4",
            "Q 5",
            "Q 6",
            "Q 7"
          ],
          "destinations": [
            "KV 0",
            "KV 1"
          ],
          "weights": [
            [
              1,
              0
            ],
            [
              1,
              0
            ],
            [
              1,
              0
            ],
            [
              1,
              0
            ],
            [
              0,
              1
            ],
            [
              0,
              1
            ],
            [
              0,
              1
            ],
            [
              0,
              1
            ]
          ]
        },
        "vectors": [
          {
            "label": "Last token output · four-token forward",
            "values": [
              0.05293419789594468,
              0.08292123901724062,
              0.16872712593069916,
              -0.00382339625361104,
              0.0269688312978711,
              -0.06342084995709538,
              -0.16084319610580047,
              0.15609071989360707,
              -0.03662947042752858,
              0.19799422200720415,
              -0.05458768818051347,
              0.049211804757439345,
              -0.0025121055467179354,
              -0.15444945950968647,
              0.03860389321179173,
              0.09010099360193347
            ]
          }
        ],
        "metrics": [
          {
            "label": "Analytical K + V cache · bytes",
            "value": 4096
          },
          {
            "label": "KV heads",
            "value": 2
          },
          {
            "label": "Query heads per KV head",
            "value": 4
          }
        ]
      },
      {
        "id": "gqa-2048",
        "label": "GQA · 2 KV heads · 2048 cached tokens",
        "note": "All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.",
        "target": "gqa",
        "routing": {
          "sourceLabel": "Query heads",
          "destinationLabel": "Shared key/value heads",
          "sources": [
            "Q 0",
            "Q 1",
            "Q 2",
            "Q 3",
            "Q 4",
            "Q 5",
            "Q 6",
            "Q 7"
          ],
          "destinations": [
            "KV 0",
            "KV 1"
          ],
          "weights": [
            [
              1,
              0
            ],
            [
              1,
              0
            ],
            [
              1,
              0
            ],
            [
              1,
              0
            ],
            [
              0,
              1
            ],
            [
              0,
              1
            ],
            [
              0,
              1
            ],
            [
              0,
              1
            ]
          ]
        },
        "vectors": [
          {
            "label": "Last token output · four-token forward",
            "values": [
              0.05293419789594468,
              0.08292123901724062,
              0.16872712593069916,
              -0.00382339625361104,
              0.0269688312978711,
              -0.06342084995709538,
              -0.16084319610580047,
              0.15609071989360707,
              -0.03662947042752858,
              0.19799422200720415,
              -0.05458768818051347,
              0.049211804757439345,
              -0.0025121055467179354,
              -0.15444945950968647,
              0.03860389321179173,
              0.09010099360193347
            ]
          }
        ],
        "metrics": [
          {
            "label": "Analytical K + V cache · bytes",
            "value": 65536
          },
          {
            "label": "KV heads",
            "value": 2
          },
          {
            "label": "Query heads per KV head",
            "value": 4
          }
        ]
      },
      {
        "id": "mqa-128",
        "label": "MQA · 1 KV heads · 128 cached tokens",
        "note": "All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.",
        "target": "mqa",
        "routing": {
          "sourceLabel": "Query heads",
          "destinationLabel": "Shared key/value heads",
          "sources": [
            "Q 0",
            "Q 1",
            "Q 2",
            "Q 3",
            "Q 4",
            "Q 5",
            "Q 6",
            "Q 7"
          ],
          "destinations": [
            "KV 0"
          ],
          "weights": [
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ]
          ]
        },
        "vectors": [
          {
            "label": "Last token output · four-token forward",
            "values": [
              0.05607162997985175,
              -0.0007644707496270195,
              0.047266194824350176,
              -0.02635026405502313,
              -0.00046292182700524796,
              -0.004557621623896115,
              -0.02339663203969375,
              0.024000452665990494,
              -0.003866001274110578,
              0.027019821004131375,
              -0.05387535332853814,
              -0.04623341221607427,
              -0.03834209820525859,
              -0.04486158178561829,
              0.026035407161435705,
              0.06696209451278268
            ]
          }
        ],
        "metrics": [
          {
            "label": "Analytical K + V cache · bytes",
            "value": 2048
          },
          {
            "label": "KV heads",
            "value": 1
          },
          {
            "label": "Query heads per KV head",
            "value": 8
          }
        ]
      },
      {
        "id": "mqa-2048",
        "label": "MQA · 1 KV heads · 2048 cached tokens",
        "note": "All eight query heads remain distinct. Lines show which K/V head each query consults. K/V weights are group means of the same MHA weights, so outputs need not match. Cache size is an analytical estimate for one layer, batch 1, head width 2 and float32, excluding temporary expansion and framework overhead.",
        "target": "mqa",
        "routing": {
          "sourceLabel": "Query heads",
          "destinationLabel": "Shared key/value heads",
          "sources": [
            "Q 0",
            "Q 1",
            "Q 2",
            "Q 3",
            "Q 4",
            "Q 5",
            "Q 6",
            "Q 7"
          ],
          "destinations": [
            "KV 0"
          ],
          "weights": [
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ],
            [
              1
            ]
          ]
        },
        "vectors": [
          {
            "label": "Last token output · four-token forward",
            "values": [
              0.05607162997985175,
              -0.0007644707496270195,
              0.047266194824350176,
              -0.02635026405502313,
              -0.00046292182700524796,
              -0.004557621623896115,
              -0.02339663203969375,
              0.024000452665990494,
              -0.003866001274110578,
              0.027019821004131375,
              -0.05387535332853814,
              -0.04623341221607427,
              -0.03834209820525859,
              -0.04486158178561829,
              0.026035407161435705,
              0.06696209451278268
            ]
          }
        ],
        "metrics": [
          {
            "label": "Analytical K + V cache · bytes",
            "value": 32768
          },
          {
            "label": "KV heads",
            "value": 1
          },
          {
            "label": "Query heads per KV head",
            "value": 8
          }
        ]
      }
    ]
  },
  "revision": "803295b72f9f4cbd",
  "files": {
    "graph.tensorviz.json": {
      "url": "/models/grouped-query/803295b72f9f4cbd/graph.tensorviz.json",
      "sha256": "ab31f160b871ec47ababcbd63e680cb1537d91c3344f83dca6fe234bc96fcb2c",
      "bytes": 101230
    },
    "model.py": {
      "url": "/models/grouped-query/803295b72f9f4cbd/model.py",
      "sha256": "2d7dbb8bead241d951d07146b0a15f05cd19c7ae0d0bb5a90bf1ad1dd974aa16",
      "bytes": 5884
    },
    "run.py": {
      "url": "/models/grouped-query/803295b72f9f4cbd/run.py",
      "sha256": "2725ca1925ae07c729d6c5df663b073aabbfdaa8ebbd11549fa7e085949a9844",
      "bytes": 670
    },
    "requirements.txt": {
      "url": "/models/grouped-query/803295b72f9f4cbd/requirements.txt",
      "sha256": "edb86016b42cb3cc6cdc0c490aab929d3bb098ced4515394f0b55766033e0daf",
      "bytes": 91
    },
    "preview.svg": {
      "url": "/models/grouped-query/803295b72f9f4cbd/preview.svg",
      "sha256": "a1ad79c03c85fd65fed6ae5f336816c6b76c2ea7087139a8e76742134e04be96",
      "bytes": 1915
    }
  }
}