{ "cells": [ { "cell_type": "markdown", "source": [ "(pallas_tpu_hardware_reference)=\n", "\n", "# TPU Hardware Reference\n", "\n", "When writing custom Pallas kernels, it is essential to understand the hardware capabilities and constraints of the target TPU generation.\n", "\n", "Note: all specs listed in the table below are **per TensorCore** (a single physical chip may contain multiple TensorCores, e.g., dual-core configurations in v4 or v5p).\n" ], "metadata": { "id": "_tO7Okdrw2db" } }, { "cell_type": "code", "outputs": [], "source": [ "from IPython.display import HTML, display\n", "from jax.experimental.pallas import tpu as pltpu\n", "\n", "headers = [\n", " \"Version\", \"Generation\", \"TensorCores/Chip\", \"VMEM Capacity\", \"CMEM Capacity\",\n", " \"SMEM Capacity\", \"HBM Capacity\", \"HBM BW\", \"BF16 Peak\", \"FP8 Peak\", \"INT8 Peak\", \"INT4 Peak\", \"SparseCore\"\n", "]\n", "\n", "html_lines = []\n", "html_lines.append(\"\")\n", "html_lines.append(\"\")\n", "html_lines.append(\" \")\n", "html_lines.append(\" \")\n", "for h in headers:\n", " html_lines.append(f\" \")\n", "html_lines.append(\" \")\n", "html_lines.append(\" \")\n", "html_lines.append(\" \")\n", "\n", "for cv in pltpu.ChipVersion:\n", " if cv == pltpu.ChipVersion.TPU_7:\n", " continue # Skip TPU 7 as it is redundant with 7x\n", "\n", " # Get per-TensorCore specs (num_cores=1)\n", " info = pltpu.get_tpu_info_for_chip(cv, 1)\n", "\n", " sc = info.sparse_core\n", " sc_str = \"No\"\n", " if sc is not None:\n", " sc_str = f\"Yes ({sc.num_cores} SCs, {sc.num_subcores} subcores, {sc.vmem_capacity_bytes // 1024} KiB VMEM)\"\n", "\n", " row = [\n", " cv.value.upper(),\n", " f\"TPU v{info.generation}\" if info.generation < 7 else f\"TPU {info.generation}\",\n", " str(cv.num_physical_tensor_cores_per_chip),\n", " f\"{info.vmem_capacity_bytes // (1024 * 1024)} MiB\",\n", " f\"{info.cmem_capacity_bytes // 1024} KiB\" if info.cmem_capacity_bytes > 0 else \"N/A\",\n", " f\"{info.smem_capacity_bytes // 1024} KiB\",\n", " f\"{info.hbm_capacity_bytes // 1000000000} GB\",\n", " f\"{info.mem_bw_bytes_per_second / 1e9:.1f} GB/s\" if info.mem_bw_bytes_per_second > 0 else \"N/A\",\n", " f\"{int(round(info.bf16_ops_per_second / 1e12))} TFLOPs/s\" if info.bf16_ops_per_second > 0 else \"N/A\",\n", " f\"{int(round(info.fp8_ops_per_second / 1e12))} TFLOPs/s\" if info.fp8_ops_per_second > 0 else \"N/A\",\n", " f\"{int(round(info.int8_ops_per_second / 1e12))} TOPs/s\" if info.int8_ops_per_second > 0 else \"N/A\",\n", " f\"{int(round(info.int4_ops_per_second / 1e12))} TOPs/s\" if info.int4_ops_per_second > 0 else \"N/A\",\n", " sc_str,\n", " ]\n", "\n", " html_lines.append(\" \")\n", " for cell in row:\n", " html_lines.append(f\" \")\n", " html_lines.append(\" \")\n", "\n", "html_lines.append(\" \")\n", "html_lines.append(\"
{h}
{cell}
\")\n", "\n", "display(HTML(\"\\n\".join(html_lines)))" ], "metadata": { "tags": [ "hide-input" ], "id": "TUMfSJsEw2db" }, "execution_count": null } ], "metadata": { "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat_minor": 0, "nbformat": 4 }