{ "format": "tensorviz.graph", "version": 3, "exportedAt": "2026-09-17T06:01:14.902248+00:00", "status": "static", "warningCount": 0, "document": { "schemaVersion": 2, "semantic": { "graph_id": "snapshot", "name": "HeadSharingComparison", "revision": 0, "nodes": [ { "node_id": "input", "kind": "input", "op": "Input", "params": { "shape": "1,4,16" }, "in_ports": [], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "input", "name": "Input", "displayName": "Input", "description": "Model input node", "icon": "ArrowRight", "category": "io", "parameters": [ { "name": "shape", "type": "string", "defaultValue": "1,3,224,224", "required": true, "description": "Input tensor shape (comma-separated, e.g., 1,3,224,224)" }, { "name": "name", "type": "string", "defaultValue": "input", "required": false, "description": "Name for this input" } ] } } }, { "node_id": "mha", "kind": "module", "op": "Module", "params": { "arg_0": 8 }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "module", "name": "Module", "displayName": "Module", "description": "Custom module group (nn.Module)", "icon": "Package", "category": "module", "parameters": [ { "name": "name", "type": "string", "defaultValue": "CustomModule", "required": true, "description": "Module class name" }, { "name": "description", "type": "string", "defaultValue": "", "required": false, "description": "Module description" } ] }, "moduleInstanceId": "mha", "moduleTemplateId": "model::SharedKVAttention", "isGroup": true, "moduleMembers": [ "python@mha#5", "mha.query", "python@mha#0", "mha.key", "python@mha#1", "python@mha#3", "div@mha#0", "mha.value", "python@mha#2", "python@mha#4", "mha.softmax", "python@mha#6", "mha.project" ], "moduleLayerCount": 13 } }, { "node_id": "gqa", "kind": "module", "op": "Module", "params": { "arg_0": 2 }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "module", "name": "Module", "displayName": "Module", "description": "Custom module group (nn.Module)", "icon": "Package", "category": "module", "parameters": [ { "name": "name", "type": "string", "defaultValue": "CustomModule", "required": true, "description": "Module class name" }, { "name": "description", "type": "string", "defaultValue": "", "required": false, "description": "Module description" } ] }, "moduleInstanceId": "gqa", "moduleTemplateId": "model::SharedKVAttention", "isGroup": true, "moduleMembers": [ "python@gqa#5", "gqa.query", "python@gqa#0", "gqa.key", "python@gqa#1", "python@gqa#3", "div@gqa#0", "gqa.value", "python@gqa#2", "python@gqa#4", "gqa.softmax", "python@gqa#6", "gqa.project" ], "moduleLayerCount": 13 } }, { "node_id": "mqa", "kind": "module", "op": "Module", "params": { "arg_0": 1 }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "module", "name": "Module", "displayName": "Module", "description": "Custom module group (nn.Module)", "icon": "Package", "category": "module", "parameters": [ { "name": "name", "type": "string", "defaultValue": "CustomModule", "required": true, "description": "Module class name" }, { "name": "description", "type": "string", "defaultValue": "", "required": false, "description": "Module description" } ] }, "moduleInstanceId": "mqa", "moduleTemplateId": "model::SharedKVAttention", "isGroup": true, "moduleMembers": [ "python@mqa#5", "mqa.query", "python@mqa#0", "mqa.key", "python@mqa#1", "python@mqa#3", "div@mqa#0", "mqa.value", "python@mqa#2", "python@mqa#4", "mqa.softmax", "python@mqa#6", "mqa.project" ], "moduleLayerCount": 13 } }, { "node_id": "stack@HeadSharingComparison#0", "kind": "layer", "op": "torch.stack", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "stack", "name": "torch.stack", "displayName": "torch.stack", "description": "PyTorch stack operation", "icon": "Activity", "category": "tensor", "parameters": [ { "name": "dim", "type": "string", "defaultValue": "", "required": false, "description": "Dimension(s) to apply the operation on" } ] } } }, { "node_id": "output", "kind": "output", "op": "Output", "params": {}, "in_ports": [ "x" ], "out_ports": [], "metadata": { "layerType": { "id": "output", "name": "Output", "displayName": "Output", "description": "Model output node", "icon": "ArrowLeft", "category": "io", "parameters": [ { "name": "name", "type": "string", "defaultValue": "output", "required": false, "description": "Name for this output" } ] } } }, { "node_id": "python@mha#5", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "mha.query", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "out_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mha" } }, { "node_id": "python@mha#0", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "mha.key", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mha" } }, { "node_id": "python@mha#1", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "python@mha#3", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "div@mha#0", "kind": "layer", "op": "torch.div", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "div", "name": "torch.div", "displayName": "torch.div", "description": "PyTorch div operation", "icon": "Activity", "category": "tensor", "parameters": [ { "name": "dim", "type": "string", "defaultValue": "", "required": false, "description": "Dimension(s) to apply the operation on" } ] }, "parentId": "mha" } }, { "node_id": "mha.value", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mha" } }, { "node_id": "python@mha#2", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "python@mha#4", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "mha.softmax", "kind": "layer", "op": "Softmax", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "softmax", "name": "Softmax", "displayName": "Softmax", "description": "Softmax over a dimension", "icon": "TrendingUp", "category": "activation", "parameters": [ { "name": "dim", "type": "number", "defaultValue": 1, "required": true, "description": "Dimension along which Softmax is computed" } ] }, "parentId": "mha" } }, { "node_id": "python@mha#6", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mha" } }, { "node_id": "mha.project", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "out_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mha" } }, { "node_id": "python@gqa#5", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "gqa.query", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "out_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "gqa" } }, { "node_id": "python@gqa#0", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "gqa.key", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "gqa" } }, { "node_id": "python@gqa#1", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "python@gqa#3", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "div@gqa#0", "kind": "layer", "op": "torch.div", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "div", "name": "torch.div", "displayName": "torch.div", "description": "PyTorch div operation", "icon": "Activity", "category": "tensor", "parameters": [ { "name": "dim", "type": "string", "defaultValue": "", "required": false, "description": "Dimension(s) to apply the operation on" } ] }, "parentId": "gqa" } }, { "node_id": "gqa.value", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "gqa" } }, { "node_id": "python@gqa#2", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "python@gqa#4", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "gqa.softmax", "kind": "layer", "op": "Softmax", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "softmax", "name": "Softmax", "displayName": "Softmax", "description": "Softmax over a dimension", "icon": "TrendingUp", "category": "activation", "parameters": [ { "name": "dim", "type": "number", "defaultValue": 1, "required": true, "description": "Dimension along which Softmax is computed" } ] }, "parentId": "gqa" } }, { "node_id": "python@gqa#6", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "gqa" } }, { "node_id": "gqa.project", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "out_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "gqa" } }, { "node_id": "python@mqa#5", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "mqa.query", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "out_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mqa" } }, { "node_id": "python@mqa#0", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "mqa.key", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mqa" } }, { "node_id": "python@mqa#1", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "python@mqa#3", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "div@mqa#0", "kind": "layer", "op": "torch.div", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "div", "name": "torch.div", "displayName": "torch.div", "description": "PyTorch div operation", "icon": "Activity", "category": "tensor", "parameters": [ { "name": "dim", "type": "string", "defaultValue": "", "required": false, "description": "Dimension(s) to apply the operation on" } ] }, "parentId": "mqa" } }, { "node_id": "mqa.value", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mqa" } }, { "node_id": "python@mqa#2", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "python@mqa#4", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "mqa.softmax", "kind": "layer", "op": "Softmax", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "softmax", "name": "Softmax", "displayName": "Softmax", "description": "Softmax over a dimension", "icon": "TrendingUp", "category": "activation", "parameters": [ { "name": "dim", "type": "number", "defaultValue": 1, "required": true, "description": "Dimension along which Softmax is computed" } ] }, "parentId": "mqa" } }, { "node_id": "python@mqa#6", "kind": "layer", "op": "Python", "params": {}, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "python", "name": "Python", "displayName": "Python", "description": "A statement from the source file that has no graph equivalent. Its code is kept verbatim and its shape is unknown until the model is traced.", "icon": "Hash", "category": "misc", "parameters": [] }, "parentId": "mqa" } }, { "node_id": "mqa.project", "kind": "layer", "op": "Linear", "params": { "in_features": 16, "out_features": 16, "bias": false }, "in_ports": [ "x" ], "out_ports": [ "y" ], "metadata": { "layerType": { "id": "linear", "name": "Linear", "displayName": "Linear", "description": "Fully Connected Layer", "icon": "Minus", "category": "linear", "parameters": [ { "name": "in_features", "type": "number", "defaultValue": 512, "required": true, "description": "Size of input features", "min": 1 }, { "name": "out_features", "type": "number", "defaultValue": 256, "required": true, "description": "Size of output features", "min": 1 }, { "name": "bias", "type": "boolean", "defaultValue": true, "required": false, "description": "If set to False, the layer will not learn an additive bias" } ] }, "parentId": "mqa" } } ], "edges": [ { "edge_id": "edge-1", "source_node_id": "mha.query", "target_node_id": "python@mha#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-2", "source_node_id": "mha.key", "target_node_id": "python@mha#1", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-3", "source_node_id": "mha.value", "target_node_id": "python@mha#2", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-4", "source_node_id": "python@mha#1", "target_node_id": "python@mha#3", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-5", "source_node_id": "python@mha#2", "target_node_id": "python@mha#4", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-6", "source_node_id": "python@mha#0", "target_node_id": "div@mha#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-7", "source_node_id": "python@mha#3", "target_node_id": "div@mha#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-8", "source_node_id": "div@mha#0", "target_node_id": "mha.softmax", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-9", "source_node_id": "python@mha#5", "target_node_id": "mha.softmax", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-10", "source_node_id": "mha.softmax", "target_node_id": "python@mha#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-11", "source_node_id": "python@mha#4", "target_node_id": "python@mha#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-12", "source_node_id": "python@mha#6", "target_node_id": "mha.project", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-13", "source_node_id": "gqa.query", "target_node_id": "python@gqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-14", "source_node_id": "gqa.key", "target_node_id": "python@gqa#1", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-15", "source_node_id": "gqa.value", "target_node_id": "python@gqa#2", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-16", "source_node_id": "python@gqa#1", "target_node_id": "python@gqa#3", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-17", "source_node_id": "python@gqa#2", "target_node_id": "python@gqa#4", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-18", "source_node_id": "python@gqa#0", "target_node_id": "div@gqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-19", "source_node_id": "python@gqa#3", "target_node_id": "div@gqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-20", "source_node_id": "div@gqa#0", "target_node_id": "gqa.softmax", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-21", "source_node_id": "python@gqa#5", "target_node_id": "gqa.softmax", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-22", "source_node_id": "gqa.softmax", "target_node_id": "python@gqa#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-23", "source_node_id": "python@gqa#4", "target_node_id": "python@gqa#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-24", "source_node_id": "python@gqa#6", "target_node_id": "gqa.project", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-25", "source_node_id": "mqa.query", "target_node_id": "python@mqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-26", "source_node_id": "mqa.key", "target_node_id": "python@mqa#1", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-27", "source_node_id": "mqa.value", "target_node_id": "python@mqa#2", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-28", "source_node_id": "python@mqa#1", "target_node_id": "python@mqa#3", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-29", "source_node_id": "python@mqa#2", "target_node_id": "python@mqa#4", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-30", "source_node_id": "python@mqa#0", "target_node_id": "div@mqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-31", "source_node_id": "python@mqa#3", "target_node_id": "div@mqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-32", "source_node_id": "div@mqa#0", "target_node_id": "mqa.softmax", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-33", "source_node_id": "python@mqa#5", "target_node_id": "mqa.softmax", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-34", "source_node_id": "mqa.softmax", "target_node_id": "python@mqa#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-35", "source_node_id": "python@mqa#4", "target_node_id": "python@mqa#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-36", "source_node_id": "python@mqa#6", "target_node_id": "mqa.project", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-37", "source_node_id": "input", "target_node_id": "mha.query", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-38", "source_node_id": "input", "target_node_id": "python@mha#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-39", "source_node_id": "input", "target_node_id": "mha.key", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-40", "source_node_id": "input", "target_node_id": "python@mha#1", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-41", "source_node_id": "input", "target_node_id": "mha.value", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-42", "source_node_id": "input", "target_node_id": "python@mha#2", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-43", "source_node_id": "input", "target_node_id": "python@mha#5", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-44", "source_node_id": "input", "target_node_id": "python@mha#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-45", "source_node_id": "input", "target_node_id": "mha", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-46", "source_node_id": "input", "target_node_id": "gqa.query", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-47", "source_node_id": "input", "target_node_id": "python@gqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-48", "source_node_id": "input", "target_node_id": "gqa.key", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-49", "source_node_id": "input", "target_node_id": "python@gqa#1", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-50", "source_node_id": "input", "target_node_id": "gqa.value", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-51", "source_node_id": "input", "target_node_id": "python@gqa#2", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-52", "source_node_id": "input", "target_node_id": "python@gqa#5", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-53", "source_node_id": "input", "target_node_id": "python@gqa#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-54", "source_node_id": "input", "target_node_id": "gqa", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-55", "source_node_id": "input", "target_node_id": "mqa.query", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-56", "source_node_id": "input", "target_node_id": "python@mqa#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-57", "source_node_id": "input", "target_node_id": "mqa.key", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-58", "source_node_id": "input", "target_node_id": "python@mqa#1", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-59", "source_node_id": "input", "target_node_id": "mqa.value", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-60", "source_node_id": "input", "target_node_id": "python@mqa#2", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-61", "source_node_id": "input", "target_node_id": "python@mqa#5", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-62", "source_node_id": "input", "target_node_id": "python@mqa#6", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-63", "source_node_id": "input", "target_node_id": "mqa", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-64", "source_node_id": "mha.project", "target_node_id": "stack@HeadSharingComparison#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-65", "source_node_id": "mha", "target_node_id": "stack@HeadSharingComparison#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-66", "source_node_id": "gqa.project", "target_node_id": "stack@HeadSharingComparison#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-67", "source_node_id": "gqa", "target_node_id": "stack@HeadSharingComparison#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-68", "source_node_id": "mqa.project", "target_node_id": "stack@HeadSharingComparison#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-69", "source_node_id": "mqa", "target_node_id": "stack@HeadSharingComparison#0", "source_port": "y", "target_port": "x", "metadata": {} }, { "edge_id": "edge-70", "source_node_id": "stack@HeadSharingComparison#0", "target_node_id": "output", "source_port": "y", "target_port": "x", "metadata": {} } ], "inputs": [], "modules": [], "operations": [], "metadata": {} }, "view": { "graph_id": "snapshot", "node_positions": { "input": { "x": 0, "y": 160 }, "mha": { "x": 340, "y": 160 }, "gqa": { "x": 680, "y": 160 }, "mqa": { "x": 1020, "y": 160 }, "stack@HeadSharingComparison#0": { "x": 1360, "y": 160 }, "output": { "x": 1700, "y": 160 }, "python@mha#5": { "x": 24, "y": 214 }, "mha.query": { "x": 24, "y": 364 }, "python@mha#0": { "x": 24, "y": 514 }, "mha.key": { "x": 24, "y": 664 }, "python@mha#1": { "x": 24, "y": 814 }, "python@mha#3": { "x": 24, "y": 964 }, "div@mha#0": { "x": 24, "y": 1114 }, "mha.value": { "x": 24, "y": 1264 }, "python@mha#2": { "x": 24, "y": 1414 }, "python@mha#4": { "x": 24, "y": 1564 }, "mha.softmax": { "x": 24, "y": 1714 }, "python@mha#6": { "x": 24, "y": 1864 }, "mha.project": { "x": 24, "y": 2014 }, "python@gqa#5": { "x": 24, "y": 214 }, "gqa.query": { "x": 24, "y": 364 }, "python@gqa#0": { "x": 24, "y": 514 }, "gqa.key": { "x": 24, "y": 664 }, "python@gqa#1": { "x": 24, "y": 814 }, "python@gqa#3": { "x": 24, "y": 964 }, "div@gqa#0": { "x": 24, "y": 1114 }, "gqa.value": { "x": 24, "y": 1264 }, "python@gqa#2": { "x": 24, "y": 1414 }, "python@gqa#4": { "x": 24, "y": 1564 }, "gqa.softmax": { "x": 24, "y": 1714 }, "python@gqa#6": { "x": 24, "y": 1864 }, "gqa.project": { "x": 24, "y": 2014 }, "python@mqa#5": { "x": 24, "y": 214 }, "mqa.query": { "x": 24, "y": 364 }, "python@mqa#0": { "x": 24, "y": 514 }, "mqa.key": { "x": 24, "y": 664 }, "python@mqa#1": { "x": 24, "y": 814 }, "python@mqa#3": { "x": 24, "y": 964 }, "div@mqa#0": { "x": 24, "y": 1114 }, "mqa.value": { "x": 24, "y": 1264 }, "python@mqa#2": { "x": 24, "y": 1414 }, "python@mqa#4": { "x": 24, "y": 1564 }, "mqa.softmax": { "x": 24, "y": 1714 }, "python@mqa#6": { "x": 24, "y": 1864 }, "mqa.project": { "x": 24, "y": 2014 } }, "collapsed_modules": [ "mha", "gqa", "mqa" ], "selection": [], "layout": {} } }, "snapshotId": "snapshot-4f8b112e717399b8", "collaboration": { "walkthrough": { "title": "Eight queries, fewer key/value heads", "steps": [ { "id": "heads", "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.", "view": { "camera": { "x": 1567.8846153846152, "y": 81.7902097902097, "z": 0.26469227209625174 }, "nodePositions": { "input": { "x": 0, "y": 160 }, "mha": { "x": 340, "y": 160 }, "gqa": { "x": 680, "y": 160 }, "mqa": { "x": 1020, "y": 160 }, "stack@HeadSharingComparison#0": { "x": 1360, "y": 160 }, "output": { "x": 1700, "y": 160 }, "python@mha#5": { "x": 24, "y": 214 }, "mha.query": { "x": 24, "y": 364 }, "python@mha#0": { "x": 24, "y": 514 }, "mha.key": { "x": 24, "y": 664 }, "python@mha#1": { "x": 24, "y": 814 }, "python@mha#3": { "x": 24, "y": 964 }, "div@mha#0": { "x": 24, "y": 1114 }, "mha.value": { "x": 24, "y": 1264 }, "python@mha#2": { "x": 24, "y": 1414 }, "python@mha#4": { "x": 24, "y": 1564 }, "mha.softmax": { "x": 24, "y": 1714 }, "python@mha#6": { "x": 24, "y": 1864 }, "mha.project": { "x": 24, "y": 2014 }, "python@gqa#5": { "x": 24, "y": 214 }, "gqa.query": { "x": 24, "y": 364 }, "python@gqa#0": { "x": 24, "y": 514 }, "gqa.key": { "x": 24, "y": 664 }, "python@gqa#1": { "x": 24, "y": 814 }, "python@gqa#3": { "x": 24, "y": 964 }, "div@gqa#0": { "x": 24, "y": 1114 }, "gqa.value": { "x": 24, "y": 1264 }, "python@gqa#2": { "x": 24, "y": 1414 }, "python@gqa#4": { "x": 24, "y": 1564 }, "gqa.softmax": { "x": 24, "y": 1714 }, "python@gqa#6": { "x": 24, "y": 1864 }, "gqa.project": { "x": 24, "y": 2014 }, "python@mqa#5": { "x": 24, "y": 214 }, "mqa.query": { "x": 24, "y": 364 }, "python@mqa#0": { "x": 24, "y": 514 }, "mqa.key": { "x": 24, "y": 664 }, "python@mqa#1": { "x": 24, "y": 814 }, "python@mqa#3": { "x": 24, "y": 964 }, "div@mqa#0": { "x": 24, "y": 1114 }, "mqa.value": { "x": 24, "y": 1264 }, "python@mqa#2": { "x": 24, "y": 1414 }, "python@mqa#4": { "x": 24, "y": 1564 }, "mqa.softmax": { "x": 24, "y": 1714 }, "python@mqa#6": { "x": 24, "y": 1864 }, "mqa.project": { "x": 24, "y": 2014 } }, "expandedModuleIds": [ "mha" ], "selectedNodeIds": [ "mha" ], "viewport": { "width": 1100, "height": 700 } }, "anchors": [ { "target": { "kind": "node", "nodeId": "mha" }, "snapshotId": "snapshot-4f8b112e717399b8", "exportedAt": "2026-09-17T06:01:14.902248+00:00", "modelName": "HeadSharingComparison", "context": "{\"connections\":[\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"div@mha#0\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha.softmax\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha.key\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha.query\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha.value\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#1\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#2\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#5\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#6\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mha\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"stack@HeadSharingComparison#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mha.key\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#1\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mha.project\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"stack@HeadSharingComparison#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mha.query\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mha.softmax\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#6\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mha.value\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#2\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#0\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"div@mha#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#1\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#3\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#2\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#4\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#3\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"div@mha#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#4\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mha#6\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#5\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha.softmax\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@mha#6\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mha.project\\\",\\\"targetPort\\\":\\\"x\\\"}\"],\"contents\":[\"{\\\"id\\\":\\\"div@mha#0\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"torch.div\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"mha.key\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Linear\\\",\\\"params\\\":{\\\"bias\\\":false,\\\"in_features\\\":16},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"mha.project\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Linear\\\",\\\"params\\\":{\\\"bias\\\":false,\\\"in_features\\\":16,\\\"out_features\\\":16},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"mha.query\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Linear\\\",\\\"params\\\":{\\\"bias\\\":false,\\\"in_features\\\":16,\\\"out_features\\\":16},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"mha.softmax\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Softmax\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"mha.value\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Linear\\\",\\\"params\\\":{\\\"bias\\\":false,\\\"in_features\\\":16},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#0\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#1\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#2\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#3\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#4\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#5\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\",\"{\\\"id\\\":\\\"python@mha#6\\\",\\\"kind\\\":\\\"layer\\\",\\\"op\\\":\\\"Python\\\",\\\"params\\\":{},\\\"parent\\\":\\\"mha\\\"}\"],\"node\":{\"id\":\"mha\",\"kind\":\"module\",\"op\":\"Module\",\"params\":{\"arg_0\":8}}}" } ] }, { "id": "groups", "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.", "view": { "camera": { "x": 308.5, "y": -444, "z": 1.25 }, "nodePositions": { "input": { "x": 0, "y": 160 }, "mha": { "x": 340, "y": 160 }, "gqa": { "x": 680, "y": 160 }, "mqa": { "x": 1020, "y": 160 }, "stack@HeadSharingComparison#0": { "x": 1360, "y": 160 }, "output": { "x": 1700, "y": 160 }, "python@mha#5": { "x": 24, "y": 214 }, "mha.query": { "x": 24, "y": 364 }, "python@mha#0": { "x": 24, "y": 514 }, "mha.key": { "x": 24, "y": 664 }, "python@mha#1": { "x": 24, "y": 814 }, "python@mha#3": { "x": 24, "y": 964 }, "div@mha#0": { "x": 24, "y": 1114 }, "mha.value": { "x": 24, "y": 1264 }, "python@mha#2": { "x": 24, "y": 1414 }, "python@mha#4": { "x": 24, "y": 1564 }, "mha.softmax": { "x": 24, "y": 1714 }, "python@mha#6": { "x": 24, "y": 1864 }, "mha.project": { "x": 24, "y": 2014 }, "python@gqa#5": { "x": 24, "y": 214 }, "gqa.query": { "x": 24, "y": 364 }, "python@gqa#0": { "x": 24, "y": 514 }, "gqa.key": { "x": 24, "y": 664 }, "python@gqa#1": { "x": 24, "y": 814 }, "python@gqa#3": { "x": 24, "y": 964 }, "div@gqa#0": { "x": 24, "y": 1114 }, "gqa.value": { "x": 24, "y": 1264 }, "python@gqa#2": { "x": 24, "y": 1414 }, "python@gqa#4": { "x": 24, "y": 1564 }, "gqa.softmax": { "x": 24, "y": 1714 }, "python@gqa#6": { "x": 24, "y": 1864 }, "gqa.project": { "x": 24, "y": 2014 }, "python@mqa#5": { "x": 24, "y": 214 }, "mqa.query": { "x": 24, "y": 364 }, "python@mqa#0": { "x": 24, "y": 514 }, "mqa.key": { "x": 24, "y": 664 }, "python@mqa#1": { "x": 24, "y": 814 }, "python@mqa#3": { "x": 24, "y": 964 }, "div@mqa#0": { "x": 24, "y": 1114 }, "mqa.value": { "x": 24, "y": 1264 }, "python@mqa#2": { "x": 24, "y": 1414 }, "python@mqa#4": { "x": 24, "y": 1564 }, "mqa.softmax": { "x": 24, "y": 1714 }, "python@mqa#6": { "x": 24, "y": 1864 }, "mqa.project": { "x": 24, "y": 2014 } }, "expandedModuleIds": [ "gqa" ], "selectedNodeIds": [ "gqa.key" ], "viewport": { "width": 1100, "height": 700 } }, "anchors": [ { "target": { "kind": "node", "nodeId": "gqa.key" }, "snapshotId": "snapshot-4f8b112e717399b8", "exportedAt": "2026-09-17T06:01:14.902248+00:00", "modelName": "HeadSharingComparison", "context": "{\"connections\":[\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"gqa.key\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@gqa#1\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"gqa.key\\\",\\\"targetPort\\\":\\\"x\\\"}\"],\"contents\":[],\"node\":{\"id\":\"gqa.key\",\"kind\":\"layer\",\"op\":\"Linear\",\"params\":{\"bias\":false,\"in_features\":16},\"parent\":\"gqa\"}}" } ] }, { "id": "single", "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.", "view": { "camera": { "x": 308.5, "y": -1044, "z": 1.25 }, "nodePositions": { "input": { "x": 0, "y": 160 }, "mha": { "x": 340, "y": 160 }, "gqa": { "x": 680, "y": 160 }, "mqa": { "x": 1020, "y": 160 }, "stack@HeadSharingComparison#0": { "x": 1360, "y": 160 }, "output": { "x": 1700, "y": 160 }, "python@mha#5": { "x": 24, "y": 214 }, "mha.query": { "x": 24, "y": 364 }, "python@mha#0": { "x": 24, "y": 514 }, "mha.key": { "x": 24, "y": 664 }, "python@mha#1": { "x": 24, "y": 814 }, "python@mha#3": { "x": 24, "y": 964 }, "div@mha#0": { "x": 24, "y": 1114 }, "mha.value": { "x": 24, "y": 1264 }, "python@mha#2": { "x": 24, "y": 1414 }, "python@mha#4": { "x": 24, "y": 1564 }, "mha.softmax": { "x": 24, "y": 1714 }, "python@mha#6": { "x": 24, "y": 1864 }, "mha.project": { "x": 24, "y": 2014 }, "python@gqa#5": { "x": 24, "y": 214 }, "gqa.query": { "x": 24, "y": 364 }, "python@gqa#0": { "x": 24, "y": 514 }, "gqa.key": { "x": 24, "y": 664 }, "python@gqa#1": { "x": 24, "y": 814 }, "python@gqa#3": { "x": 24, "y": 964 }, "div@gqa#0": { "x": 24, "y": 1114 }, "gqa.value": { "x": 24, "y": 1264 }, "python@gqa#2": { "x": 24, "y": 1414 }, "python@gqa#4": { "x": 24, "y": 1564 }, "gqa.softmax": { "x": 24, "y": 1714 }, "python@gqa#6": { "x": 24, "y": 1864 }, "gqa.project": { "x": 24, "y": 2014 }, "python@mqa#5": { "x": 24, "y": 214 }, "mqa.query": { "x": 24, "y": 364 }, "python@mqa#0": { "x": 24, "y": 514 }, "mqa.key": { "x": 24, "y": 664 }, "python@mqa#1": { "x": 24, "y": 814 }, "python@mqa#3": { "x": 24, "y": 964 }, "div@mqa#0": { "x": 24, "y": 1114 }, "mqa.value": { "x": 24, "y": 1264 }, "python@mqa#2": { "x": 24, "y": 1414 }, "python@mqa#4": { "x": 24, "y": 1564 }, "mqa.softmax": { "x": 24, "y": 1714 }, "python@mqa#6": { "x": 24, "y": 1864 }, "mqa.project": { "x": 24, "y": 2014 } }, "expandedModuleIds": [ "mqa" ], "selectedNodeIds": [ "mqa.value" ], "viewport": { "width": 1100, "height": 700 } }, "anchors": [ { "target": { "kind": "node", "nodeId": "mqa.value" }, "snapshotId": "snapshot-4f8b112e717399b8", "exportedAt": "2026-09-17T06:01:14.902248+00:00", "modelName": "HeadSharingComparison", "context": "{\"connections\":[\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"mqa.value\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"mqa.value\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@mqa#2\\\",\\\"targetPort\\\":\\\"x\\\"}\"],\"contents\":[],\"node\":{\"id\":\"mqa.value\",\"kind\":\"layer\",\"op\":\"Linear\",\"params\":{\"bias\":false,\"in_features\":16},\"parent\":\"mqa\"}}" } ] }, { "id": "output", "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.", "view": { "camera": { "x": 308.5, "y": -1797.5, "z": 1.25 }, "nodePositions": { "input": { "x": 0, "y": 160 }, "mha": { "x": 340, "y": 160 }, "gqa": { "x": 680, "y": 160 }, "mqa": { "x": 1020, "y": 160 }, "stack@HeadSharingComparison#0": { "x": 1360, "y": 160 }, "output": { "x": 1700, "y": 160 }, "python@mha#5": { "x": 24, "y": 214 }, "mha.query": { "x": 24, "y": 364 }, "python@mha#0": { "x": 24, "y": 514 }, "mha.key": { "x": 24, "y": 664 }, "python@mha#1": { "x": 24, "y": 814 }, "python@mha#3": { "x": 24, "y": 964 }, "div@mha#0": { "x": 24, "y": 1114 }, "mha.value": { "x": 24, "y": 1264 }, "python@mha#2": { "x": 24, "y": 1414 }, "python@mha#4": { "x": 24, "y": 1564 }, "mha.softmax": { "x": 24, "y": 1714 }, "python@mha#6": { "x": 24, "y": 1864 }, "mha.project": { "x": 24, "y": 2014 }, "python@gqa#5": { "x": 24, "y": 214 }, "gqa.query": { "x": 24, "y": 364 }, "python@gqa#0": { "x": 24, "y": 514 }, "gqa.key": { "x": 24, "y": 664 }, "python@gqa#1": { "x": 24, "y": 814 }, "python@gqa#3": { "x": 24, "y": 964 }, "div@gqa#0": { "x": 24, "y": 1114 }, "gqa.value": { "x": 24, "y": 1264 }, "python@gqa#2": { "x": 24, "y": 1414 }, "python@gqa#4": { "x": 24, "y": 1564 }, "gqa.softmax": { "x": 24, "y": 1714 }, "python@gqa#6": { "x": 24, "y": 1864 }, "gqa.project": { "x": 24, "y": 2014 }, "python@mqa#5": { "x": 24, "y": 214 }, "mqa.query": { "x": 24, "y": 364 }, "python@mqa#0": { "x": 24, "y": 514 }, "mqa.key": { "x": 24, "y": 664 }, "python@mqa#1": { "x": 24, "y": 814 }, "python@mqa#3": { "x": 24, "y": 964 }, "div@mqa#0": { "x": 24, "y": 1114 }, "mqa.value": { "x": 24, "y": 1264 }, "python@mqa#2": { "x": 24, "y": 1414 }, "python@mqa#4": { "x": 24, "y": 1564 }, "mqa.softmax": { "x": 24, "y": 1714 }, "python@mqa#6": { "x": 24, "y": 1864 }, "mqa.project": { "x": 24, "y": 2014 } }, "expandedModuleIds": [ "gqa" ], "selectedNodeIds": [ "gqa.project" ], "viewport": { "width": 1100, "height": 700 } }, "anchors": [ { "target": { "kind": "node", "nodeId": "gqa.project" }, "snapshotId": "snapshot-4f8b112e717399b8", "exportedAt": "2026-09-17T06:01:14.902248+00:00", "modelName": "HeadSharingComparison", "context": "{\"connections\":[\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"gqa.project\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"stack@HeadSharingComparison#0\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"python@gqa#6\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"gqa.project\\\",\\\"targetPort\\\":\\\"x\\\"}\"],\"contents\":[],\"node\":{\"id\":\"gqa.project\",\"kind\":\"layer\",\"op\":\"Linear\",\"params\":{\"bias\":false,\"in_features\":16,\"out_features\":16},\"parent\":\"gqa\"}}" } ] } ] }, "threads": [], "savedViews": [ { "id": "groups", "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.", "view": { "camera": { "x": 308.5, "y": -444, "z": 1.25 }, "nodePositions": { "input": { "x": 0, "y": 160 }, "mha": { "x": 340, "y": 160 }, "gqa": { "x": 680, "y": 160 }, "mqa": { "x": 1020, "y": 160 }, "stack@HeadSharingComparison#0": { "x": 1360, "y": 160 }, "output": { "x": 1700, "y": 160 }, "python@mha#5": { "x": 24, "y": 214 }, "mha.query": { "x": 24, "y": 364 }, "python@mha#0": { "x": 24, "y": 514 }, "mha.key": { "x": 24, "y": 664 }, "python@mha#1": { "x": 24, "y": 814 }, "python@mha#3": { "x": 24, "y": 964 }, "div@mha#0": { "x": 24, "y": 1114 }, "mha.value": { "x": 24, "y": 1264 }, "python@mha#2": { "x": 24, "y": 1414 }, "python@mha#4": { "x": 24, "y": 1564 }, "mha.softmax": { "x": 24, "y": 1714 }, "python@mha#6": { "x": 24, "y": 1864 }, "mha.project": { "x": 24, "y": 2014 }, "python@gqa#5": { "x": 24, "y": 214 }, "gqa.query": { "x": 24, "y": 364 }, "python@gqa#0": { "x": 24, "y": 514 }, "gqa.key": { "x": 24, "y": 664 }, "python@gqa#1": { "x": 24, "y": 814 }, "python@gqa#3": { "x": 24, "y": 964 }, "div@gqa#0": { "x": 24, "y": 1114 }, "gqa.value": { "x": 24, "y": 1264 }, "python@gqa#2": { "x": 24, "y": 1414 }, "python@gqa#4": { "x": 24, "y": 1564 }, "gqa.softmax": { "x": 24, "y": 1714 }, "python@gqa#6": { "x": 24, "y": 1864 }, "gqa.project": { "x": 24, "y": 2014 }, "python@mqa#5": { "x": 24, "y": 214 }, "mqa.query": { "x": 24, "y": 364 }, "python@mqa#0": { "x": 24, "y": 514 }, "mqa.key": { "x": 24, "y": 664 }, "python@mqa#1": { "x": 24, "y": 814 }, "python@mqa#3": { "x": 24, "y": 964 }, "div@mqa#0": { "x": 24, "y": 1114 }, "mqa.value": { "x": 24, "y": 1264 }, "python@mqa#2": { "x": 24, "y": 1414 }, "python@mqa#4": { "x": 24, "y": 1564 }, "mqa.softmax": { "x": 24, "y": 1714 }, "python@mqa#6": { "x": 24, "y": 1864 }, "mqa.project": { "x": 24, "y": 2014 } }, "expandedModuleIds": [ "gqa" ], "selectedNodeIds": [ "gqa.key" ], "viewport": { "width": 1100, "height": 700 } }, "anchors": [ { "target": { "kind": "node", "nodeId": "gqa.key" }, "snapshotId": "snapshot-4f8b112e717399b8", "exportedAt": "2026-09-17T06:01:14.902248+00:00", "modelName": "HeadSharingComparison", "context": "{\"connections\":[\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"gqa.key\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"python@gqa#1\\\",\\\"targetPort\\\":\\\"x\\\"}\",\"{\\\"kind\\\":\\\"edge\\\",\\\"source\\\":\\\"input\\\",\\\"sourcePort\\\":\\\"y\\\",\\\"target\\\":\\\"gqa.key\\\",\\\"targetPort\\\":\\\"x\\\"}\"],\"contents\":[],\"node\":{\"id\":\"gqa.key\",\"kind\":\"layer\",\"op\":\"Linear\",\"params\":{\"bias\":false,\"in_features\":16},\"parent\":\"gqa\"}}" } ] } ], "openingViewId": "groups" } }