{
  "schema_version": 1,
  "model": "tiny-llama-style-educational",
  "weights": "deterministic random untrained",
  "tokenizer": "hand-written 32-label vocabulary; not a real tokenizer",
  "numpy_version": "2.3.5",
  "seed": 7,
  "dtype": "float64",
  "config": {
    "d_model": 32,
    "n_layers": 2,
    "n_heads": 4,
    "n_kv_heads": 2,
    "head_dim": 8,
    "hidden_dim": 96,
    "vocab_size": 32,
    "norm_eps": 1e-05,
    "rope_theta": 10000.0
  },
  "vocabulary": [
    "<bos>",
    "I",
    "love",
    "learning",
    "about",
    "attention",
    ".",
    "How",
    "does",
    "a",
    "model",
    "use",
    "context",
    "?",
    "the",
    "key",
    "value",
    "query",
    "cache",
    "token",
    "is",
    "shared",
    "across",
    "heads",
    "hello",
    "world",
    "<eos>",
    "and",
    "memory",
    "position",
    "new",
    "old"
  ],
  "cache_bytes_formula": "2 * batch * layers * cached_tokens * kv_heads * head_dim * bytes_per_element",
  "verification": {
    "passed": true,
    "dtype": "float64",
    "tolerance": 1e-10,
    "checks": {
      "cached_single_vs_full_max_abs_error": 2.3314683517128287e-15,
      "cached_chunked_vs_full_max_abs_error": 1.942890293094024e-15,
      "future_change_earlier_logits_max_abs_error": 0.0,
      "gqa_vs_explicit_head_repeat_max_abs_error": 4.440892098500626e-16,
      "masked_attention_weight_max": 0.0,
      "rope_position_0_vs_1_max_abs_difference": 0.6702715711140383,
      "rope_norm_preservation_max_abs_error": 0.0,
      "cache_shape_per_tensor": [
        2,
        2,
        6,
        8
      ],
      "cache_payload_bytes_batch_2_length_6": 6144,
      "expected_cache_payload_bytes": 6144,
      "prefix_cache_unchanged_after_continuation": true
    }
  },
  "passes": [
    {
      "token_ids": [
        [
          1,
          2,
          3,
          4,
          5,
          6
        ]
      ],
      "tokens": [
        [
          "I",
          "love",
          "learning",
          "about",
          "attention",
          "."
        ]
      ],
      "input_shape": [
        1,
        6
      ],
      "embedding_shape": [
        1,
        6,
        32
      ],
      "layers": [
        {
          "positions": [
            0,
            1,
            2,
            3,
            4,
            5
          ],
          "key_positions": [
            0,
            1,
            2,
            3,
            4,
            5
          ],
          "shapes": {
            "input": [
              1,
              6,
              32
            ],
            "normed": [
              1,
              6,
              32
            ],
            "q": [
              1,
              4,
              6,
              8
            ],
            "new_k": [
              1,
              2,
              6,
              8
            ],
            "new_v": [
              1,
              2,
              6,
              8
            ],
            "cached_k": [
              1,
              2,
              6,
              8
            ],
            "cached_v": [
              1,
              2,
              6,
              8
            ],
            "attention_weights": [
              1,
              4,
              6,
              6
            ],
            "joined_heads": [
              1,
              6,
              32
            ],
            "gate": [
              1,
              6,
              96
            ],
            "up": [
              1,
              6,
              96
            ],
            "ffn_output": [
              1,
              6,
              32
            ],
            "output": [
              1,
              6,
              32
            ]
          },
          "cache_payload_bytes": 1536,
          "causal_mask": [
            [
              1,
              0,
              0,
              0,
              0,
              0
            ],
            [
              1,
              1,
              0,
              0,
              0,
              0
            ],
            [
              1,
              1,
              1,
              0,
              0,
              0
            ],
            [
              1,
              1,
              1,
              1,
              0,
              0
            ],
            [
              1,
              1,
              1,
              1,
              1,
              0
            ],
            [
              1,
              1,
              1,
              1,
              1,
              1
            ]
          ],
          "attention_weights": [
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.64418054,
                0.35581946,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.1056534,
                0.09566341,
                0.79868318,
                0.0,
                0.0,
                0.0
              ],
              [
                0.0931381,
                0.0448033,
                0.39086775,
                0.47119085,
                0.0,
                0.0
              ],
              [
                0.31249122,
                0.48498533,
                0.0295792,
                0.10896831,
                0.06397595,
                0.0
              ],
              [
                0.06883099,
                0.02935051,
                0.04723346,
                0.67497671,
                0.14433562,
                0.0352727
              ]
            ],
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.45153516,
                0.54846484,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.20692206,
                0.41218742,
                0.38089052,
                0.0,
                0.0,
                0.0
              ],
              [
                0.07615015,
                0.1024436,
                0.66244626,
                0.15895999,
                0.0,
                0.0
              ],
              [
                0.29692282,
                0.46338444,
                0.10128914,
                0.10220909,
                0.0361945,
                0.0
              ],
              [
                0.23936481,
                0.08105766,
                0.03813087,
                0.16193374,
                0.11798353,
                0.36152938
              ]
            ],
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.39167109,
                0.60832891,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.27367596,
                0.68833033,
                0.03799371,
                0.0,
                0.0,
                0.0
              ],
              [
                0.39563544,
                0.48050171,
                0.0448556,
                0.07900724,
                0.0,
                0.0
              ],
              [
                0.17225398,
                0.07627101,
                0.03302661,
                0.04787584,
                0.67057256,
                0.0
              ],
              [
                0.02164317,
                0.24696145,
                0.23941347,
                0.08052422,
                0.10159209,
                0.30986559
              ]
            ],
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.07365773,
                0.92634227,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.2750586,
                0.16824361,
                0.55669778,
                0.0,
                0.0,
                0.0
              ],
              [
                0.32350946,
                0.2181692,
                0.16545028,
                0.29287106,
                0.0,
                0.0
              ],
              [
                0.26484286,
                0.56199408,
                0.04287384,
                0.05862792,
                0.07166131,
                0.0
              ],
              [
                0.27090091,
                0.19961626,
                0.12937113,
                0.186581,
                0.08832454,
                0.12520616
              ]
            ]
          ],
          "first_query_before_rope": [
            0.35726811,
            1.3522358,
            -0.55286365,
            2.60052177,
            -0.34557496,
            -0.79849068,
            0.94091316,
            0.61754123
          ],
          "first_query_after_rope": [
            0.35726811,
            1.3522358,
            -0.55286365,
            2.60052177,
            -0.34557496,
            -0.79849068,
            0.94091316,
            0.61754123
          ],
          "first_input_values": [
            -0.19570382,
            -0.16176745,
            0.21217972,
            -0.16150694,
            -0.00650434,
            0.17687797,
            -0.11672009,
            -0.02234039
          ],
          "first_output_values": [
            0.84274978,
            1.86694686,
            -0.97891768,
            0.58766354,
            -1.40992571,
            -1.63679098,
            0.06313976,
            -0.82179448
          ],
          "residual_rms": {
            "input": 0.18115600455773284,
            "after_attention": 0.6369966349702493,
            "output": 0.8343706190588777
          },
          "layer": 0
        },
        {
          "positions": [
            0,
            1,
            2,
            3,
            4,
            5
          ],
          "key_positions": [
            0,
            1,
            2,
            3,
            4,
            5
          ],
          "shapes": {
            "input": [
              1,
              6,
              32
            ],
            "normed": [
              1,
              6,
              32
            ],
            "q": [
              1,
              4,
              6,
              8
            ],
            "new_k": [
              1,
              2,
              6,
              8
            ],
            "new_v": [
              1,
              2,
              6,
              8
            ],
            "cached_k": [
              1,
              2,
              6,
              8
            ],
            "cached_v": [
              1,
              2,
              6,
              8
            ],
            "attention_weights": [
              1,
              4,
              6,
              6
            ],
            "joined_heads": [
              1,
              6,
              32
            ],
            "gate": [
              1,
              6,
              96
            ],
            "up": [
              1,
              6,
              96
            ],
            "ffn_output": [
              1,
              6,
              32
            ],
            "output": [
              1,
              6,
              32
            ]
          },
          "cache_payload_bytes": 1536,
          "causal_mask": [
            [
              1,
              0,
              0,
              0,
              0,
              0
            ],
            [
              1,
              1,
              0,
              0,
              0,
              0
            ],
            [
              1,
              1,
              1,
              0,
              0,
              0
            ],
            [
              1,
              1,
              1,
              1,
              0,
              0
            ],
            [
              1,
              1,
              1,
              1,
              1,
              0
            ],
            [
              1,
              1,
              1,
              1,
              1,
              1
            ]
          ],
          "attention_weights": [
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.5812572,
                0.4187428,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.42905994,
                0.29407083,
                0.27686923,
                0.0,
                0.0,
                0.0
              ],
              [
                0.26757383,
                0.204816,
                0.26772976,
                0.25988041,
                0.0,
                0.0
              ],
              [
                0.2173835,
                0.28542046,
                0.26163939,
                0.15682065,
                0.078736,
                0.0
              ],
              [
                0.25022814,
                0.17338616,
                0.17038649,
                0.18927422,
                0.00878578,
                0.20793922
              ]
            ],
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.55120879,
                0.44879121,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.45975944,
                0.26094268,
                0.27929787,
                0.0,
                0.0,
                0.0
              ],
              [
                0.21310178,
                0.34755645,
                0.28924739,
                0.15009438,
                0.0,
                0.0
              ],
              [
                0.2634055,
                0.23296548,
                0.15622913,
                0.2240308,
                0.12336909,
                0.0
              ],
              [
                0.25969032,
                0.20071547,
                0.15442311,
                0.15319395,
                0.06910531,
                0.16287185
              ]
            ],
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.62615305,
                0.37384695,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.09844955,
                0.03886627,
                0.86268418,
                0.0,
                0.0,
                0.0
              ],
              [
                0.06477075,
                0.1028545,
                0.31100237,
                0.52137237,
                0.0,
                0.0
              ],
              [
                0.08351748,
                0.19011723,
                0.10164132,
                0.17327437,
                0.45144961,
                0.0
              ],
              [
                0.03125308,
                0.07584456,
                0.07659871,
                0.10012705,
                0.11454673,
                0.60162988
              ]
            ],
            [
              [
                1.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.86282537,
                0.13717463,
                0.0,
                0.0,
                0.0,
                0.0
              ],
              [
                0.45989351,
                0.20704217,
                0.33306433,
                0.0,
                0.0,
                0.0
              ],
              [
                0.23019404,
                0.11753423,
                0.29777659,
                0.35449514,
                0.0,
                0.0
              ],
              [
                0.14449634,
                0.01681083,
                0.47824842,
                0.33824633,
                0.02219808,
                0.0
              ],
              [
                0.09778509,
                0.10485081,
                0.10265172,
                0.38892376,
                0.16540726,
                0.14038136
              ]
            ]
          ],
          "first_query_before_rope": [
            -0.91716944,
            0.66333279,
            -2.62107042,
            0.96836365,
            0.94219102,
            0.48485138,
            -0.20886093,
            0.21543552
          ],
          "first_query_after_rope": [
            -0.91716944,
            0.66333279,
            -2.62107042,
            0.96836365,
            0.94219102,
            0.48485138,
            -0.20886093,
            0.21543552
          ],
          "first_input_values": [
            0.84274978,
            1.86694686,
            -0.97891768,
            0.58766354,
            -1.40992571,
            -1.63679098,
            0.06313976,
            -0.82179448
          ],
          "first_output_values": [
            0.46084959,
            2.08249409,
            1.08105669,
            0.71356196,
            -2.41331734,
            -2.16342553,
            1.97455485,
            -1.80148333
          ],
          "residual_rms": {
            "input": 0.8343706190588777,
            "after_attention": 1.2283704981484425,
            "output": 1.3669266187673113
          },
          "layer": 1
        }
      ],
      "final_hidden_shape": [
        1,
        6,
        32
      ],
      "logits_shape": [
        1,
        6,
        32
      ],
      "cache_payload_bytes": 3072,
      "name": "prefill",
      "last_position_top5": [
        {
          "token_id": 2,
          "token": "love",
          "logit": 2.0172140453618557,
          "probability": 0.1342117316861278
        },
        {
          "token_id": 20,
          "token": "is",
          "logit": 1.8943120908767246,
          "probability": 0.11869019389011418
        },
        {
          "token_id": 19,
          "token": "token",
          "logit": 1.7742590899940225,
          "probability": 0.10526317940978974
        },
        {
          "token_id": 14,
          "token": "the",
          "logit": 1.639295645065076,
          "probability": 0.09197347680073498
        },
        {
          "token_id": 10,
          "token": "model",
          "logit": 1.141432167948279,
          "probability": 0.05590404633164627
        }
      ]
    },
    {
      "token_ids": [
        [
          2
        ]
      ],
      "tokens": [
        [
          "love"
        ]
      ],
      "input_shape": [
        1,
        1
      ],
      "embedding_shape": [
        1,
        1,
        32
      ],
      "layers": [
        {
          "positions": [
            6
          ],
          "key_positions": [
            0,
            1,
            2,
            3,
            4,
            5,
            6
          ],
          "shapes": {
            "input": [
              1,
              1,
              32
            ],
            "normed": [
              1,
              1,
              32
            ],
            "q": [
              1,
              4,
              1,
              8
            ],
            "new_k": [
              1,
              2,
              1,
              8
            ],
            "new_v": [
              1,
              2,
              1,
              8
            ],
            "cached_k": [
              1,
              2,
              7,
              8
            ],
            "cached_v": [
              1,
              2,
              7,
              8
            ],
            "attention_weights": [
              1,
              4,
              1,
              7
            ],
            "joined_heads": [
              1,
              1,
              32
            ],
            "gate": [
              1,
              1,
              96
            ],
            "up": [
              1,
              1,
              96
            ],
            "ffn_output": [
              1,
              1,
              32
            ],
            "output": [
              1,
              1,
              32
            ]
          },
          "cache_payload_bytes": 1792,
          "causal_mask": [
            [
              1,
              1,
              1,
              1,
              1,
              1,
              1
            ]
          ],
          "attention_weights": [
            [
              [
                0.11141331,
                0.0799265,
                0.17877603,
                0.15937584,
                0.23016037,
                0.1580711,
                0.08227685
              ]
            ],
            [
              [
                0.09310844,
                0.07659029,
                0.10179175,
                0.09441692,
                0.10958535,
                0.46588565,
                0.0586216
              ]
            ],
            [
              [
                0.17715765,
                0.11917223,
                0.19085007,
                0.10775773,
                0.02155416,
                0.13069871,
                0.25280945
              ]
            ],
            [
              [
                0.02047253,
                0.03719491,
                0.33199513,
                0.11635576,
                0.16119715,
                0.11893889,
                0.21384563
              ]
            ]
          ],
          "first_query_before_rope": [
            -0.30174579,
            0.59459629,
            -0.56651155,
            0.07383385,
            0.28104173,
            -0.26200773,
            -0.25825551,
            -0.35589105
          ],
          "first_query_after_rope": [
            -0.21119992,
            0.63868219,
            -0.5500061,
            0.07596786,
            0.35416037,
            0.11949001,
            -0.29176109,
            -0.35544165
          ],
          "first_input_values": [
            -0.11586032,
            -0.03923919,
            0.17975277,
            0.2290444,
            -0.26470556,
            -0.15892847,
            0.12938068,
            -0.39848396
          ],
          "first_output_values": [
            0.58941297,
            -0.19039057,
            -1.17299468,
            0.08758456,
            -0.09384345,
            -2.45370222,
            0.70704482,
            0.15130369
          ],
          "residual_rms": {
            "input": 0.15755413087608236,
            "after_attention": 0.47314858288642947,
            "output": 0.7235052897884617
          },
          "layer": 0
        },
        {
          "positions": [
            6
          ],
          "key_positions": [
            0,
            1,
            2,
            3,
            4,
            5,
            6
          ],
          "shapes": {
            "input": [
              1,
              1,
              32
            ],
            "normed": [
              1,
              1,
              32
            ],
            "q": [
              1,
              4,
              1,
              8
            ],
            "new_k": [
              1,
              2,
              1,
              8
            ],
            "new_v": [
              1,
              2,
              1,
              8
            ],
            "cached_k": [
              1,
              2,
              7,
              8
            ],
            "cached_v": [
              1,
              2,
              7,
              8
            ],
            "attention_weights": [
              1,
              4,
              1,
              7
            ],
            "joined_heads": [
              1,
              1,
              32
            ],
            "gate": [
              1,
              1,
              96
            ],
            "up": [
              1,
              1,
              96
            ],
            "ffn_output": [
              1,
              1,
              32
            ],
            "output": [
              1,
              1,
              32
            ]
          },
          "cache_payload_bytes": 1792,
          "causal_mask": [
            [
              1,
              1,
              1,
              1,
              1,
              1,
              1
            ]
          ],
          "attention_weights": [
            [
              [
                0.33400072,
                0.18022347,
                0.07972657,
                0.10681264,
                0.03104186,
                0.21136959,
                0.05682516
              ]
            ],
            [
              [
                0.1112176,
                0.02416355,
                0.1124483,
                0.52104116,
                0.16312465,
                0.0591612,
                0.00884353
              ]
            ],
            [
              [
                0.10728334,
                0.13745157,
                0.13638266,
                0.03072708,
                0.08154522,
                0.4757343,
                0.03087583
              ]
            ],
            [
              [
                0.0324592,
                0.07380654,
                0.07647249,
                0.17174573,
                0.14634408,
                0.3102951,
                0.18887686
              ]
            ]
          ],
          "first_query_before_rope": [
            -1.84821146,
            -0.43601391,
            0.72214458,
            1.24128476,
            0.64150688,
            -1.49877348,
            0.13485076,
            -0.31303207
          ],
          "first_query_after_rope": [
            -1.59535076,
            0.48641336,
            0.71275892,
            1.2431406,
            1.13237477,
            -1.48318311,
            0.17791078,
            -0.30557877
          ],
          "first_input_values": [
            0.58941297,
            -0.19039057,
            -1.17299468,
            0.08758456,
            -0.09384345,
            -2.45370222,
            0.70704482,
            0.15130369
          ],
          "first_output_values": [
            1.07065608,
            2.02622022,
            -0.61107389,
            0.64968308,
            0.38085165,
            -3.19292459,
            2.37475628,
            -2.02967629
          ],
          "residual_rms": {
            "input": 0.7235052897884617,
            "after_attention": 1.0307965482895942,
            "output": 1.2436349935091706
          },
          "layer": 1
        }
      ],
      "final_hidden_shape": [
        1,
        1,
        32
      ],
      "logits_shape": [
        1,
        1,
        32
      ],
      "cache_payload_bytes": 3584,
      "name": "single_token_decode",
      "last_position_top5": [
        {
          "token_id": 26,
          "token": "<eos>",
          "logit": 1.7028339248043074,
          "probability": 0.1111807705827186
        },
        {
          "token_id": 2,
          "token": "love",
          "logit": 1.6728292894484516,
          "probability": 0.10789438209846562
        },
        {
          "token_id": 14,
          "token": "the",
          "logit": 1.5641475372228273,
          "probability": 0.0967829703691393
        },
        {
          "token_id": 15,
          "token": "key",
          "logit": 1.234501723671116,
          "probability": 0.0696042228975658
        },
        {
          "token_id": 30,
          "token": "new",
          "logit": 1.22127930489295,
          "probability": 0.068689944519895
        }
      ]
    },
    {
      "token_ids": [
        [
          2,
          28
        ]
      ],
      "tokens": [
        [
          "love",
          "memory"
        ]
      ],
      "input_shape": [
        1,
        2
      ],
      "embedding_shape": [
        1,
        2,
        32
      ],
      "layers": [
        {
          "positions": [
            6,
            7
          ],
          "key_positions": [
            0,
            1,
            2,
            3,
            4,
            5,
            6,
            7
          ],
          "shapes": {
            "input": [
              1,
              2,
              32
            ],
            "normed": [
              1,
              2,
              32
            ],
            "q": [
              1,
              4,
              2,
              8
            ],
            "new_k": [
              1,
              2,
              2,
              8
            ],
            "new_v": [
              1,
              2,
              2,
              8
            ],
            "cached_k": [
              1,
              2,
              8,
              8
            ],
            "cached_v": [
              1,
              2,
              8,
              8
            ],
            "attention_weights": [
              1,
              4,
              2,
              8
            ],
            "joined_heads": [
              1,
              2,
              32
            ],
            "gate": [
              1,
              2,
              96
            ],
            "up": [
              1,
              2,
              96
            ],
            "ffn_output": [
              1,
              2,
              32
            ],
            "output": [
              1,
              2,
              32
            ]
          },
          "cache_payload_bytes": 2048,
          "causal_mask": [
            [
              1,
              1,
              1,
              1,
              1,
              1,
              1,
              0
            ],
            [
              1,
              1,
              1,
              1,
              1,
              1,
              1,
              1
            ]
          ],
          "attention_weights": [
            [
              [
                0.11141331,
                0.0799265,
                0.17877603,
                0.15937584,
                0.23016037,
                0.1580711,
                0.08227685,
                0.0
              ],
              [
                0.06206918,
                0.03679566,
                0.06460506,
                0.58964987,
                0.02933755,
                0.00575772,
                0.06768972,
                0.14409523
              ]
            ],
            [
              [
                0.09310844,
                0.07659029,
                0.10179175,
                0.09441692,
                0.10958535,
                0.46588565,
                0.0586216,
                0.0
              ],
              [
                0.06903907,
                0.21112201,
                0.08929936,
                0.01663062,
                0.21920801,
                0.1678675,
                0.19904754,
                0.02778588
              ]
            ],
            [
              [
                0.17715765,
                0.11917223,
                0.19085007,
                0.10775773,
                0.02155416,
                0.13069871,
                0.25280945,
                0.0
              ],
              [
                0.08148264,
                0.05244202,
                0.26111941,
                0.19043549,
                0.09981015,
                0.08749128,
                0.11625603,
                0.11096298
              ]
            ],
            [
              [
                0.02047253,
                0.03719491,
                0.33199513,
                0.11635576,
                0.16119715,
                0.11893889,
                0.21384563,
                0.0
              ],
              [
                0.03440538,
                0.13611444,
                0.05272344,
                0.02675379,
                0.41809136,
                0.09716604,
                0.20744803,
                0.02729753
              ]
            ]
          ],
          "first_query_before_rope": [
            -0.30174579,
            0.59459629,
            -0.56651155,
            0.07383385,
            0.28104173,
            -0.26200773,
            -0.25825551,
            -0.35589105
          ],
          "first_query_after_rope": [
            -0.21119992,
            0.63868219,
            -0.5500061,
            0.07596786,
            0.35416037,
            0.11949001,
            -0.29176109,
            -0.35544165
          ],
          "first_input_values": [
            -0.11586032,
            -0.03923919,
            0.17975277,
            0.2290444,
            -0.26470556,
            -0.15892847,
            0.12938068,
            -0.39848396
          ],
          "first_output_values": [
            0.58941297,
            -0.19039057,
            -1.17299468,
            0.08758456,
            -0.09384345,
            -2.45370222,
            0.70704482,
            0.15130369
          ],
          "residual_rms": {
            "input": 0.17765741112380734,
            "after_attention": 0.4583745147783859,
            "output": 0.6977691843823902
          },
          "layer": 0
        },
        {
          "positions": [
            6,
            7
          ],
          "key_positions": [
            0,
            1,
            2,
            3,
            4,
            5,
            6,
            7
          ],
          "shapes": {
            "input": [
              1,
              2,
              32
            ],
            "normed": [
              1,
              2,
              32
            ],
            "q": [
              1,
              4,
              2,
              8
            ],
            "new_k": [
              1,
              2,
              2,
              8
            ],
            "new_v": [
              1,
              2,
              2,
              8
            ],
            "cached_k": [
              1,
              2,
              8,
              8
            ],
            "cached_v": [
              1,
              2,
              8,
              8
            ],
            "attention_weights": [
              1,
              4,
              2,
              8
            ],
            "joined_heads": [
              1,
              2,
              32
            ],
            "gate": [
              1,
              2,
              96
            ],
            "up": [
              1,
              2,
              96
            ],
            "ffn_output": [
              1,
              2,
              32
            ],
            "output": [
              1,
              2,
              32
            ]
          },
          "cache_payload_bytes": 2048,
          "causal_mask": [
            [
              1,
              1,
              1,
              1,
              1,
              1,
              1,
              0
            ],
            [
              1,
              1,
              1,
              1,
              1,
              1,
              1,
              1
            ]
          ],
          "attention_weights": [
            [
              [
                0.33400072,
                0.18022347,
                0.07972657,
                0.10681264,
                0.03104186,
                0.21136959,
                0.05682516,
                0.0
              ],
              [
                0.11086161,
                0.07489074,
                0.08574444,
                0.13432101,
                0.10854366,
                0.15486515,
                0.14920206,
                0.18157135
              ]
            ],
            [
              [
                0.1112176,
                0.02416355,
                0.1124483,
                0.52104116,
                0.16312465,
                0.0591612,
                0.00884353,
                0.0
              ],
              [
                0.0615137,
                0.03045881,
                0.06464477,
                0.18158749,
                0.52189629,
                0.08091525,
                0.03018207,
                0.02880161
              ]
            ],
            [
              [
                0.10728334,
                0.13745157,
                0.13638266,
                0.03072708,
                0.08154522,
                0.4757343,
                0.03087583,
                0.0
              ],
              [
                0.16825252,
                0.08319738,
                0.24235912,
                0.073545,
                0.07123133,
                0.22536905,
                0.04173278,
                0.09431282
              ]
            ],
            [
              [
                0.0324592,
                0.07380654,
                0.07647249,
                0.17174573,
                0.14634408,
                0.3102951,
                0.18887686,
                0.0
              ],
              [
                0.20367989,
                0.03713137,
                0.1986421,
                0.0754816,
                0.08080152,
                0.2323248,
                0.10097358,
                0.07096515
              ]
            ]
          ],
          "first_query_before_rope": [
            -1.84821146,
            -0.43601391,
            0.72214458,
            1.24128476,
            0.64150688,
            -1.49877348,
            0.13485076,
            -0.31303207
          ],
          "first_query_after_rope": [
            -1.59535076,
            0.48641336,
            0.71275892,
            1.2431406,
            1.13237477,
            -1.48318311,
            0.17791078,
            -0.30557877
          ],
          "first_input_values": [
            0.58941297,
            -0.19039057,
            -1.17299468,
            0.08758456,
            -0.09384345,
            -2.45370222,
            0.70704482,
            0.15130369
          ],
          "first_output_values": [
            1.07065608,
            2.02622022,
            -0.61107389,
            0.64968308,
            0.38085165,
            -3.19292459,
            2.37475628,
            -2.02967629
          ],
          "residual_rms": {
            "input": 0.6977691843823902,
            "after_attention": 1.0206588298222272,
            "output": 1.2532026894853683
          },
          "layer": 1
        }
      ],
      "final_hidden_shape": [
        1,
        2,
        32
      ],
      "logits_shape": [
        1,
        2,
        32
      ],
      "cache_payload_bytes": 4096,
      "name": "chunked_decode",
      "last_position_top5": [
        {
          "token_id": 10,
          "token": "model",
          "logit": 1.8886673141770374,
          "probability": 0.12025894656831196
        },
        {
          "token_id": 20,
          "token": "is",
          "logit": 1.782273804131963,
          "probability": 0.1081213055029847
        },
        {
          "token_id": 14,
          "token": "the",
          "logit": 1.686836894251527,
          "probability": 0.0982796401362856
        },
        {
          "token_id": 25,
          "token": "world",
          "logit": 1.5794357867692832,
          "probability": 0.08827136638735868
        },
        {
          "token_id": 2,
          "token": "love",
          "logit": 1.2729707218783195,
          "probability": 0.0649716300374235
        }
      ]
    }
  ],
  "notes": [
    "All numbers are measured by this NumPy reference; they are not pretrained-model results.",
    "Prefill computes the distribution for the first generated token.",
    "The single-token and chunked demonstrations each start from the same six-token prefix.",
    "Chunked decode uses the greedy next ID followed by a supplied token ID for demonstration.",
    "Cache byte counts include K/V payload, excluding metadata and temporary attention arrays.",
    "Keys are cached after RoPE; only two KV heads are stored for four query heads."
  ]
}
