{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "675a1a18",
   "metadata": {},
   "source": [
    "# Cross Entropy：token PID 与双层扫描\n",
    "\n",
    "运行安全小尺寸的真实融合内核，并展示 token 分块、D 累加和 V 维在线 softmax。\n",
    "\n",
    "> 动画按逻辑 PID 顺序播放，不代表 GPU 的真实执行先后或 SM 调度。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cce8e55",
   "metadata": {},
   "source": [
    "## 1. 环境检查与教程源码加载\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90bcf63b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import importlib.util\n",
    "import math\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "from matplotlib.colors import ListedColormap\n",
    "from matplotlib import font_manager\n",
    "from IPython.display import HTML, display\n",
    "import numpy as np\n",
    "import torch\n",
    "import triton\n",
    "\n",
    "assert torch.cuda.is_available(), \"本 Notebook 必须在支持 CUDA 的环境中运行。\"\n",
    "DEVICE = torch.device(\"cuda\")\n",
    "torch.manual_seed(0)\n",
    "print(\"GPU:\", torch.cuda.get_device_name(0))\n",
    "print(\"PyTorch:\", torch.__version__, \"Triton:\", triton.__version__)\n",
    "\n",
    "def configure_chinese_font():\n",
    "    candidates = [\n",
    "        \"Noto Sans CJK SC\",\n",
    "        \"WenQuanYi Zen Hei\",\n",
    "        \"Microsoft YaHei\",\n",
    "        \"SimHei\",\n",
    "    ]\n",
    "    installed = {font.name for font in font_manager.fontManager.ttflist}\n",
    "    selected = next((name for name in candidates if name in installed), None)\n",
    "    if selected is None:\n",
    "        raise RuntimeError(\n",
    "            \"未找到中文字体。请在 Ubuntu 中安装 fonts-noto-cjk 或 fonts-wqy-zenhei。\"\n",
    "        )\n",
    "    plt.rcParams[\"font.family\"] = \"sans-serif\"\n",
    "    plt.rcParams[\"font.sans-serif\"] = [selected, \"DejaVu Sans\"]\n",
    "    plt.rcParams[\"axes.unicode_minus\"] = False\n",
    "    print(\"Matplotlib 中文字体:\", selected)\n",
    "\n",
    "configure_chinese_font()\n",
    "\n",
    "def load_tutorial(filename, module_name):\n",
    "    path = Path.cwd() / filename\n",
    "    spec = importlib.util.spec_from_file_location(module_name, path)\n",
    "    module = importlib.util.module_from_spec(spec)\n",
    "    spec.loader.exec_module(module)\n",
    "    return module\n",
    "\n",
    "def show_animation(fig, update, frame_count, interval=700):\n",
    "    anim = animation.FuncAnimation(\n",
    "        fig, update, frames=frame_count, interval=interval, repeat=True, blit=False\n",
    "    )\n",
    "    plt.close(fig)\n",
    "    return HTML(anim.to_jshtml())\n",
    "\n",
    "tutorial = load_tutorial('celoss.py', 'celoss_tutorial')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "923f245a",
   "metadata": {},
   "source": [
    "## 2. 参数修改区\n",
    "\n",
    "可修改下列小尺寸参数。请保留注释中说明的内核约束。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3412544",
   "metadata": {},
   "outputs": [],
   "source": [
    "B, N, D, V = 1, 16, 16, 16\n",
    "BS_N = BS_D = BS_V = 16\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e40b563",
   "metadata": {},
   "source": [
    "## 3. 真实 Triton 内核与 PyTorch 数值验证\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6240ae49",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn((B, N, D), device=DEVICE)\n",
    "embedding = torch.randn((D, V), device=DEVICE)\n",
    "targets = torch.randint(0, V, (B, N), device=DEVICE)\n",
    "tri = tutorial.fused_CELoss(x, embedding, targets)\n",
    "ref = torch.nn.functional.cross_entropy((x @ embedding).reshape(-1, V), targets.reshape(-1))\n",
    "torch.testing.assert_close(tri, ref, atol=1e-2, rtol=0)\n",
    "print(\"验证通过，loss =\", float(tri))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b2463ba",
   "metadata": {},
   "source": [
    "## 4. PID 总览静态图\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2a79cde",
   "metadata": {},
   "outputs": [],
   "source": [
    "VIS_TOKENS = 38; num_pids = math.ceil(VIS_TOKENS / BS_N)\n",
    "owner = np.repeat(np.arange(num_pids), BS_N)\n",
    "fig, ax = plt.subplots(figsize=(11, 2)); ax.imshow(owner[None, :], aspect=\"auto\", cmap=\"tab20\")\n",
    "ax.axvline(VIS_TOKENS - 0.5, color=\"red\", linestyle=\"--\")\n",
    "ax.set_yticks([]); ax.set_xlabel(\"展平 token\"); ax.set_title(\"等价追踪：每个 PID 负责 bsN 个 token\"); plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a326c4",
   "metadata": {},
   "source": [
    "## 5. PID 与算法阶段逐帧动画\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb1b4e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "VIS_D, VIS_V = 48, 64; d_steps, v_steps = math.ceil(VIS_D / BS_D), math.ceil(VIS_V / BS_V)\n",
    "phases = [\"矩阵乘法累加\", \"更新在线 softmax\", \"选择目标类别\"]\n",
    "frames = [(pid, vs, ds, phase) for pid in range(num_pids) for vs in range(v_steps) for ds in range(d_steps) for phase in phases]\n",
    "fig, axes = plt.subplots(1, 3, figsize=(14, 4), constrained_layout=True)\n",
    "def update(frame):\n",
    "    pid, vs, ds, phase = frames[frame]\n",
    "    for ax in axes: ax.clear()\n",
    "    token = np.zeros((num_pids, 1)); dm = np.zeros((1, d_steps)); vm = np.zeros((1, v_steps))\n",
    "    token[pid] = dm[0, ds] = vm[0, vs] = pid + 1\n",
    "    for ax, data, title in zip(axes, [token, dm, vm], [\"token PID\", \"D 维块\", \"V 维块\"]):\n",
    "        ax.imshow(data, cmap=\"tab20\", vmin=0, vmax=20, aspect=\"auto\"); ax.set_title(title)\n",
    "    fig.suptitle(f\"PID {pid}：{phase}，D {ds + 1}/{d_steps}，V {vs + 1}/{v_steps}\")\n",
    "display(show_animation(fig, update, len(frames), interval=450))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07bbf975",
   "metadata": {},
   "source": [
    "## 6. 结论\n",
    "\n",
    "外层扫描词表 V，内层扫描隐藏维 D；每个 V 块更新在线 softmax 并选择目标类别。\n",
    "\n",
    "逻辑 PID 映射由 launch grid 和源码索引关系确定；实际 GPU 调度可能并发且顺序不同。\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (triton-tutorial)",
   "language": "python",
   "name": "triton-tutorial"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
