{
  "schemaVersion": 1,
  "lessonId": "tpu-03",
  "contentHash": "e6dd2b245a8dfadb2626f5ecca9e89c0a274b8b8c32b51ec9ce69dffaefa00f9",
  "rendererHash": "a1bd5b648f4731ae5d5f0a24811c2540d1a0a132149ce589d17c4113509cbedf",
  "kind": "analytic",
  "generatedAt": "2026-10-20T10:58:50.515621+00:00",
  "provenance": "analytic_calculation_not_device_execution",
  "environment": {
    "backend": "cpu",
    "deviceCount": 1,
    "jax": "0.9.2",
    "numpy": "2.4.4",
    "matplotlib": "3.10.8"
  },
  "stdout": "TPU ridge points (FLOP/byte): {'v5e': 240.5, 'v5p': 166.0, 'v6e': 573.8}\n7B memory budget (GiB): {'fp32_train_state_gib': 104.308, 'mixed_bf16_train_state_gib': 117.347, 'kv_fp32_gib': 2.0, 'kv_bf16_gib': 1.0, 'kv_int8_gib': 0.5}\n1.5B memory budget (GiB): {'fp32_train_state_gib': 22.352, 'mixed_bf16_train_state_gib': 25.146, 'kv_fp32_gib': 1.0, 'kv_bf16_gib': 0.5, 'kv_int8_gib': 0.25}\n",
  "data": {
    "kind": "bar",
    "panels": [
      {
        "kind": "bar",
        "labels": [
          "v5e",
          "v5p",
          "v6e"
        ],
        "xlabel": "Cloud TPU generation",
        "ylabel": "HBM per chip (GiB)",
        "series": [
          {
            "label": "HBM capacity",
            "y": [
              16.0,
              95.0,
              32.0
            ]
          }
        ]
      },
      {
        "kind": "bar",
        "labels": [
          "v5e",
          "v5p",
          "v6e"
        ],
        "xlabel": "Cloud TPU generation",
        "ylabel": "theoretical BF16 ridge point (FLOP/byte)",
        "series": [
          {
            "label": "peak BF16 FLOP/s \u00f7 HBM bytes/s",
            "y": [
              240.5,
              166.0,
              573.8
            ]
          }
        ]
      }
    ]
  },
  "diagram": {
    "title": "TPU HBM-to-MXU dataflow and mixed-precision state budget",
    "prediction": "Which tensors stay in `float32` in HBM and which operands stream as `bfloat16` into the MXU?",
    "reading": "Read top to bottom: HBM stores master `fp32` parameters and Adam moments (`m`, `v`) alongside `bf16` activations and KV cache; operands are cast to `bf16` (or `int8`) to stream across the HBM bandwidth bottleneck into the MXU, which accumulates dot products in `fp32` before writing updated state back.",
    "alt": "Read top to bottom: HBM holds fp32 master weights and Adam moments plus bf16/int8 caches, streams bf16 operands into the MXU, accumulates dot products in fp32, and updates sharded HBM state.",
    "scope": "Architecture and dataflow mechanism diagram.",
    "height": 6,
    "panel": {
      "kind": "graph",
      "nodes": [
        {
          "id": "0",
          "x": 5,
          "y": 1,
          "label": "TPU HBM capacity: fp32 master weights + Adam (m, v) + bf16/int8 KV cache",
          "tone": "data",
          "w": 8.6,
          "h": 1.05
        },
        {
          "id": "1",
          "x": 5,
          "y": 2.7,
          "label": "HBM bandwidth pipe (819-2765 GB/s): stream bf16 / int8 operands",
          "tone": "state",
          "w": 8.6,
          "h": 1.05
        },
        {
          "id": "2",
          "x": 5,
          "y": 4.4,
          "label": "TPU MXU systolic array: bf16/int8 multiply + fp32 accumulation",
          "tone": "result",
          "w": 8.6,
          "h": 1.05
        },
        {
          "id": "3",
          "x": 5,
          "y": 6.1,
          "label": "XProf *.xplane.pb trace: verify MXU busy % vs HBM bandwidth stall",
          "tone": "warning",
          "w": 8.6,
          "h": 1.05
        }
      ],
      "edges": [
        {
          "from": "0",
          "to": "1"
        },
        {
          "from": "1",
          "to": "2"
        },
        {
          "from": "2",
          "to": "3"
        }
      ],
      "height": 7,
      "width": 10
    }
  }
}