{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LayerNorm：PID 内存映射与数学公式\n",
    "\n",
    "选择内核、步骤和 PID，查看符号地址映射。公式通过 IPython.display.Math 输出，由 Jupyter 内置 MathJax 渲染，不依赖外部 KaTeX 或具体张量数值。\n",
    "\n",
    "<style>\n",
    ".ln-map { background: #ffffff; color: #172033; padding: 12px; border: 1px solid #b8c2d1; }\n",
    ".ln-label { color: #172033; font-weight: 600; }\n",
    ".ln-cell { background: #f8fafc; color: #172033; }\n",
    ".ln-cell.current { background: #c7e2ff; color: #0b2a4a; border-color: #1464b4; font-weight: 700; }\n",
    ".ln-cell.mask { background: #e4e7ec; color: #4b5563; }\n",
    ".ln-code { color: #172033; }\n",
    "</style>"
   ],
   "id": "34f1f157a0b1a8db"
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-08-18T08:24:31.131141400Z",
     "start_time": "2026-08-18T08:24:31.086103700Z"
    }
   },
   "source": [
    "from IPython.display import HTML, display\n",
    "\n",
    "layernorm_visualization_html = '<style>\\n#layernorm-pid-visualization{background:#fff;color:#172033;color-scheme:light;border:1px solid #cbd5e1;padding:16px;font-family:system-ui,\\'Microsoft YaHei\\',sans-serif;line-height:1.45}\\n#layernorm-pid-visualization *{box-sizing:border-box}.ln-controls{display:flex;flex-wrap:wrap;gap:8px;align-items:end;margin-bottom:12px}.ln-controls label{display:flex;flex-direction:column;gap:3px;min-width:170px;font-size:13px}.ln-controls select,.ln-controls input,.ln-controls button{font:inherit}.ln-controls button{border:1px solid #94a3b8;background:#fff;color:#172033;padding:6px 10px;cursor:pointer}.ln-controls button.primary{background:#2563eb;color:#fff}.ln-status{padding:8px;background:#f8fafc;border:1px solid #cbd5e1;margin-bottom:10px}.ln-layout{display:grid;grid-template-columns:minmax(0,1.2fr) minmax(300px,.8fr);gap:14px}.ln-memory,.ln-panel{background:#fff;border:1px solid #cbd5e1;padding:10px;min-width:0}.ln-memory{overflow:auto}.ln-row{display:grid;grid-template-columns:120px repeat(8,minmax(48px,1fr));gap:3px;margin:3px 0;min-width:520px}.ln-label,.ln-cell{padding:5px 3px;border:1px solid #cbd5e1;text-align:center;font-size:12px}.ln-label{text-align:left;border-color:transparent;font-weight:600}.ln-cell.inactive{opacity:.35}.ln-cell.read{background:#bfdbfe;color:#12345b;border-color:#2563eb}.ln-cell.write{background:#bbf7d0;color:#14532d;border-color:#16a34a}.ln-cell.shared{background:#fde68a;color:#713f12;border-color:#d97706}.ln-cell.partial{background:#ddd6fe;color:#4c1d95;border-color:#7c3aed}.ln-cell.lock{background:#fecaca;color:#7f1d1d;border-color:#dc2626}.ln-cell.mask{background:repeating-linear-gradient(135deg,#e5e7eb 0,#e5e7eb 5px,#fff 5px,#fff 10px);color:#4b5563}.ln-panel h3{margin:0 0 8px;font-size:16px}.ln-panel h4{margin:12px 0 4px;font-size:13px;color:#334155}.ln-formula,.ln-address{background:#f8fafc;border-left:3px solid #2563eb;padding:8px;white-space:normal;overflow-x:auto;overflow-y:hidden;color:#172033}.ln-formula .MathJax,.ln-address .MathJax{max-width:100%;overflow-x:auto}.ln-code{margin:0;background:#f1f5f9;border-left:3px solid #64748b;padding:8px;white-space:pre;overflow-x:auto;overflow-y:auto;color:#172033;font:12px ui-monospace,SFMono-Regular,Consolas,monospace}.ln-legend{display:flex;flex-wrap:wrap;gap:5px;margin:8px 0;font-size:11px}.ln-legend span{padding:3px 6px;border:1px solid #cbd5e1}.ln-legend .read{background:#bfdbfe}.ln-legend .write{background:#bbf7d0}.ln-legend .shared{background:#fde68a}.ln-legend .partial{background:#ddd6fe}.ln-legend .lock{background:#fecaca}.ln-legend .mask{background:#e5e7eb}.ln-no-mathjax{font:12px ui-monospace,monospace;white-space:pre-wrap;color:#7f1d1d}@media(max-width:760px){.ln-layout{grid-template-columns:1fr}.ln-row{grid-template-columns:90px repeat(8,minmax(34px,1fr))}}\\n</style>\\n<div id=\"layernorm-pid-visualization\">\\n<div class=\"ln-controls\"><label>内核<select data-role=\"kernel\"><option value=\"forward\">前向内核</option><option value=\"backward_input\">反向：输入与部分和</option><option value=\"backward_params\">反向：参数归约</option></select></label><label>PID <output data-role=\"pid-value\">0</output><input data-role=\"pid\" type=\"range\" min=\"0\" max=\"3\" value=\"0\"></label><label>步骤 <output data-role=\"step-value\">0</output><input data-role=\"step\" type=\"range\" min=\"0\" max=\"3\" value=\"0\"></label><button data-action=\"previous\">上一步</button><button class=\"primary\" data-action=\"play\">播放</button><button data-action=\"next\">下一步</button></div>\\n<div class=\"ln-status\" data-role=\"status\"></div><div class=\"ln-legend\"><span class=\"read\">读取</span><span class=\"write\">写入</span><span class=\"shared\">共享参数</span><span class=\"partial\">部分和</span><span class=\"lock\">锁/计数</span><span class=\"mask\">mask lane</span></div>\\n<div class=\"ln-layout\"><div class=\"ln-memory\" data-role=\"memory\"></div><div class=\"ln-panel\"><h3 data-role=\"title\"></h3><div data-role=\"detail\"></div><h4>数学公式</h4><div class=\"ln-formula\" data-role=\"formula\"></div><h4>地址映射</h4><div class=\"ln-address\" data-role=\"address\"></div><h4>对应 Triton 代码</h4><pre class=\"ln-code\" data-role=\"code\"></pre></div></div></div>\\n<script>\\n(() => {\\n  const root=document.getElementById(\\'layernorm-pid-visualization\\'); if(!root)return;\\n  const $=role=>root.querySelector(`[data-role=\"${role}\"],[data-action=\"${role}\"]`);\\n  const M=4,N=6,G=2,BLOCK=3;\\n  const math=lines=>String.raw`\\\\[\\\\begin{aligned}${lines.join(String.raw`\\\\\\\\`)}\\\\end{aligned}\\\\]`;\\n  const stage=(title,detail,formula,address,code)=>({title,detail,formulaLines:formula,addressLines:address,formula:math(formula),address:math(address),codeLines:code});\\n  const stages={\\n    forward:[\\n      stage(\\'PID 行映射\\',\\'grid=(M,)，一个 PID 负责一整行。\\',[String.raw`r=\\\\mathrm{program\\\\_id}(0)`],[String.raw`x_{row}=x_{base}+r\\\\cdot row\\\\_stride`,String.raw`y_{row}=y_{base}+r\\\\cdot row\\\\_stride`],[\\'row = tl.program_id(0)\\',\\'x_ptr += row * row_stride\\',\\'y_ptr += row * row_stride\\']),\\n      stage(\\'均值、列偏移与 mask\\',\\'有效 lane 参与均值归约，越界 lane 由 mask 屏蔽。\\',[String.raw`\\\\mu_r=\\\\frac{1}{N}\\\\sum_{j=0}^{N-1}x_{rj}`],[String.raw`columns_j=offset+j`,String.raw`addr(x_{rj})=x_{base}+r\\\\cdot row\\\\_stride+columns_j`,String.raw`addr(mean_r)=mean_{base}+r`,String.raw`mask_j=(columns_j<N)`],[\\'columns = offset + tl.arange(0, block_size)\\',\\'x = tl.load(x_ptr + columns, mask=mask, other=0.0)\\',\\'mean = tl.sum(x, axis=0) / n_cols\\']),\\n      stage(\\'方差与 rstd\\',\\'再次读取 x 和 mean，并写入当前行的 rstd。\\',[String.raw`\\\\sigma_r^2=\\\\frac{1}{N}\\\\sum_j(x_{rj}-\\\\mu_r)^2`,String.raw`rstd_r=(\\\\sigma_r^2+\\\\varepsilon)^{-1/2}`],[String.raw`addr(x_{rj})=x_{base}+r\\\\cdot row\\\\_stride+columns_j`,String.raw`addr(mean_r)=mean_{base}+r`,String.raw`addr(rstd_r)=rstd_{base}+r`,String.raw`difference_j=0\\\\text{ if }columns_j\\\\ge N`],[\\'difference = tl.where(mask, x - mean, 0.0)\\',\\'variance = tl.sum(difference * difference, axis=0) / n_cols\\',\\'rstd = 1 / tl.sqrt(variance + eps)\\']),\\n      stage(\\'weight、bias 与输出写回\\',\\'共享参数被所有 PID 读取，完成仿射变换并写回 y。\\',[String.raw`\\\\gamma_j=weight[j],\\\\quad \\\\beta_j=bias[j]`,String.raw`\\\\hat{x}_{rj}=(x_{rj}-\\\\mu_r)rstd_r`,String.raw`y_{rj}=\\\\gamma_j\\\\hat{x}_{rj}+\\\\beta_j`],[String.raw`addr(weight_j)=weight_{base}+j`,String.raw`addr(bias_j)=bias_{base}+j`,String.raw`addr(y_{rj})=y_{base}+r\\\\cdot row\\\\_stride+j`,String.raw`mask_j=(j<N)`],[\\'weight = tl.load(weight_ptr + columns, mask=mask)\\',\\'bias = tl.load(bias_ptr + columns, mask=mask)\\',\\'y = normalized_x * weight + bias\\',\\'tl.store(y_ptr + columns, y, mask=mask)\\'])\\n    ],\\n    backward_input:[\\n      stage(\\'反向行映射\\',\\'同一 PID 对齐访问 x、grad_y、grad_x 及保存的统计量。\\',[String.raw`r=\\\\mathrm{program\\\\_id}(0)`,String.raw`\\\\hat{x}_{rj}=(x_{rj}-\\\\mu_r)rstd_r`],[String.raw`addr(x_{rj})=x_{base}+r\\\\cdot row\\\\_stride+j`,String.raw`addr(grad\\\\_y_{rj})=grad\\\\_y_{base}+r\\\\cdot row\\\\_stride+j`,String.raw`addr(grad\\\\_x_{rj})=grad\\\\_x_{base}+r\\\\cdot row\\\\_stride+j`,String.raw`addr(weight_j)=weight_{base}+j`,String.raw`addr(mean_r)=mean_{base}+r`,String.raw`addr(rstd_r)=rstd_{base}+r`,String.raw`mask_j=(j<N)`],[\\'program_id = tl.program_id(0)\\',\\'x_ptr += program_id * row_stride\\',\\'grad_y_ptr += program_id * row_stride\\',\\'grad_x_ptr += program_id * row_stride\\']),\\n      stage(\\'输入梯度行内归约\\',\\'projection 和 grad_mean 都沿特征列归约。\\',[String.raw`g_{rj}=\\\\frac{\\\\partial L}{\\\\partial y_{rj}}\\\\gamma_j`,String.raw`\\\\frac{\\\\partial L}{\\\\partial x_{rj}}=rstd_r[g_{rj}-mean(g_r)-\\\\hat{x}_{rj}mean(\\\\hat{x}_rg_r)]`],[String.raw`read: x_{rj},\\\\ grad\\\\_y_{rj},\\\\ weight_j,\\\\ mean_r,\\\\ rstd_r`,String.raw`write: grad\\\\_x_{base}+r\\\\cdot row\\\\_stride+j`,String.raw`j\\\\in[0,N),\\\\quad mask_j=(j<N)`],[\\'grad_normalized_x = weight * grad_y\\',\\'projection = tl.sum(normalized_x * grad_normalized_x, axis=0) / n_cols\\',\\'grad_mean = tl.sum(grad_normalized_x, axis=0) / n_cols\\',\\'grad_x = (grad_normalized_x - normalized_x * projection - grad_mean) * rstd\\']),\\n      stage(\\'参数梯度部分和与锁\\',\\'lock_id=r mod G；同组 PID 串行累加部分和。\\',[String.raw`\\\\Delta\\\\gamma_{rj}=\\\\frac{\\\\partial L}{\\\\partial y_{rj}}\\\\hat{x}_{rj}`,String.raw`\\\\Delta\\\\beta_{rj}=\\\\frac{\\\\partial L}{\\\\partial y_{rj}}`,String.raw`lock\\\\_id=r\\\\bmod G`],[String.raw`lock\\\\_addr=locks_{base}+lock\\\\_id`,String.raw`count\\\\_addr=locks_{base}+G+lock\\\\_id`,String.raw`addr(P^\\\\gamma_{gj})=grad\\\\_weight\\\\_partial_{base}+gN+j`,String.raw`addr(P^\\\\beta_{gj})=grad\\\\_bias\\\\_partial_{base}+gN+j`],[\\'lock_id = program_id % group_size\\',\\'locks_ptr += lock_id\\',\\'count_ptr = locks_ptr + group_size\\',\\'tl.atomic_cas(locks_ptr, 0, 1)\\'])\\n    ],\\n    backward_params:[\\n      stage(\\'参数归约列块\\',\\'PID q 负责一个列块；教学图 BLOCK_SIZE_N=3，源码为 128。\\',[String.raw`Q=\\\\left\\\\lceil\\\\frac{N}{BLOCK\\\\_SIZE\\\\_N}\\\\right\\\\rceil`,String.raw`q=\\\\mathrm{program\\\\_id}(0)`],[String.raw`columns_j=q\\\\cdot BLOCK\\\\_SIZE\\\\_N+j`,String.raw`0\\\\le columns_j<N`,String.raw`output\\\\_mask_j=(columns_j<N)`],[\\'q = tl.program_id(0)\\',\\'columns = q * block_size_n + tl.arange(0, block_size_n)\\']),\\n      stage(\\'二维广播读取部分和\\',\\'rows[:,None] 与 columns[None,:] 形成二维地址。\\',[String.raw`offset_{g,j}=gN+columns_j`,String.raw`g\\\\in[0,G)`],[String.raw`addr(P^\\\\gamma_{g,j})=grad\\\\_weight\\\\_partial_{base}+gN+columns_j`,String.raw`addr(P^\\\\beta_{g,j})=grad\\\\_bias\\\\_partial_{base}+gN+columns_j`,String.raw`mask_{g,j}=(g<G)\\\\land(columns_j<N)`],[\\'rows = tl.arange(0, group_size)\\',\\'columns = q * block_size_n + tl.arange(0, block_size_n)\\',\\'offsets = rows[:, None] * n_cols + columns[None, :]\\']),\\n      stage(\\'归约并写回最终参数梯度\\',\\'沿 group 维归约并写回当前列块。\\',[String.raw`\\\\frac{\\\\partial L}{\\\\partial\\\\gamma_j}=\\\\sum_gP^\\\\gamma_{gj}`,String.raw`\\\\frac{\\\\partial L}{\\\\partial\\\\beta_j}=\\\\sum_gP^\\\\beta_{gj}`],[String.raw`addr(grad\\\\_weight_j)=grad\\\\_weight_{base}+columns_j`,String.raw`addr(grad\\\\_bias_j)=grad\\\\_bias_{base}+columns_j`,String.raw`output\\\\_mask_j=(columns_j<N)`],[\\'grad_weight = tl.sum(partial_weight, axis=0)\\',\\'grad_bias = tl.sum(partial_bias, axis=0)\\',\\'tl.store(grad_weight_ptr + columns, grad_weight, mask=output_mask)\\',\\'tl.store(grad_bias_ptr + columns, grad_bias, mask=output_mask\\'])\\n    ]\\n  };\\n  const state={kernel:\\'forward\\',step:0,pid:0,playing:false,timer:null};\\n  const active=(len,start,end,kind)=>Array.from({length:len},(_,i)=>i>=start&&i<end?kind:\\'inactive\\');\\n  const row=(label,values,states)=>`<div class=\"ln-row\"><div class=\"ln-label\">${label}</div>${values.map((v,i)=>`<div class=\"ln-cell ${states[i]||\\'inactive\\'}\">${v}</div>`).join(\\'\\')}</div>`;\\n  const matrix=(name,r,kind)=>row(`${name}[${r},:]`,Array.from({length:N},(_,j)=>`${name}[${r},${j}]`).concat([\\'mask\\',\\'mask\\']),active(N+2,0,N,kind).map((x,i)=>i>=N?\\'mask\\':x));\\n  function renderMemory(){const k=state.kernel,s=state.step,p=state.pid,blocks=[],lock=p%G;if(k===\\'forward\\'){for(let r=0;r<M;r++)blocks.push(matrix(\\'x\\',r,r===p&&s>=0?\\'read\\':\\'inactive\\'));blocks.push(row(\\'mean[M]\\',Array.from({length:M},(_,i)=>`mu[${i}]`),active(M,p,p+1,s===1?\\'write\\':s>=2?\\'read\\':\\'inactive\\')));blocks.push(row(\\'rstd[M]\\',Array.from({length:M},(_,i)=>`rstd[${i}]`),active(M,p,p+1,s===2?\\'write\\':s===3?\\'read\\':\\'inactive\\')));blocks.push(row(\\'weight[N]\\',Array.from({length:N},(_,i)=>`gamma[${i}]`),active(N,0,N,s===3?\\'shared\\':\\'inactive\\')));blocks.push(row(\\'bias[N]\\',Array.from({length:N},(_,i)=>`beta[${i}]`),active(N,0,N,s===3?\\'shared\\':\\'inactive\\')));for(let r=0;r<M;r++)blocks.push(matrix(\\'y\\',r,s===3&&r===p?\\'write\\':\\'inactive\\'))}else if(k===\\'backward_input\\'){for(const n of [\\'x\\',\\'grad_y\\',\\'grad_x\\'])for(let r=0;r<M;r++)blocks.push(matrix(n,r,r===p?(n===\\'grad_x\\'&&s>=1?\\'write\\':\\'read\\'):\\'inactive\\'));blocks.push(row(\\'weight[N]\\',Array.from({length:N},(_,i)=>`gamma[${i}]`),active(N,0,N,s>=1?\\'shared\\':\\'inactive\\')));for(const n of [\\'mean\\',\\'rstd\\'])blocks.push(row(`${n}[M]`,Array.from({length:M},(_,i)=>`${n===\\'mean\\'?\\'mu\\':\\'rstd\\'}[${i}]`),active(M,p,p+1,\\'read\\')));for(let g=0;g<G;g++){blocks.push(row(`Pgamma[${g},:]`,Array.from({length:N},(_,i)=>`Pgamma[${g},${i}]`),active(N,0,N,s===2&&g===lock?\\'partial\\':\\'inactive\\')));blocks.push(row(`Pbeta[${g},:]`,Array.from({length:N},(_,i)=>`Pbeta[${g},${i}]`),active(N,0,N,s===2&&g===lock?\\'partial\\':\\'inactive\\')))}blocks.push(row(\\'locks[0:G]\\',Array.from({length:G},(_,i)=>`lock[${i}]`),active(G,lock,lock+1,s===2?\\'lock\\':\\'inactive\\')));blocks.push(row(\\'locks[G:2G]\\',Array.from({length:G},(_,i)=>`count[${i}]`),active(G,lock,lock+1,s===2?\\'lock\\':\\'inactive\\')))}else{const start=p*BLOCK,end=Math.min(N,start+BLOCK);for(let g=0;g<G;g++){blocks.push(row(`Pgamma[${g},:]`,Array.from({length:N},(_,i)=>`Pgamma[${g},${i}]`),active(N,start,end,s>=1?\\'partial\\':\\'inactive\\')));blocks.push(row(`Pbeta[${g},:]`,Array.from({length:N},(_,i)=>`Pbeta[${g},${i}]`),active(N,start,end,s>=1?\\'partial\\':\\'inactive\\')))}blocks.push(row(\\'grad_weight[N]\\',Array.from({length:N},(_,i)=>`dgamma[${i}]`),active(N,start,end,s===2?\\'write\\':\\'inactive\\')));blocks.push(row(\\'grad_bias[N]\\',Array.from({length:N},(_,i)=>`dbeta[${i}]`),active(N,start,end,s===2?\\'write\\':\\'inactive\\')))}$(\\'memory\\').innerHTML=blocks.join(\\'\\')}\\n  let mathJaxLoader=null;\\n  function ensureMathJax(){\\n    if(window.MathJax?.typesetPromise||window.MathJax?.Hub?.Queue)return Promise.resolve();\\n    if(mathJaxLoader)return mathJaxLoader;\\n    window.MathJax={tex:{displayMath:[[\\'\\\\\\\\[\\',\\'\\\\\\\\]\\']]},startup:{typeset:false}};\\n    mathJaxLoader=new Promise((resolve,reject)=>{const script=document.createElement(\\'script\\');script.src=\\'https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-svg.js\\';script.async=true;script.onload=resolve;script.onerror=reject;document.head.appendChild(script)});\\n    return mathJaxLoader;\\n  }\\n  async function setMath(el,markup,plain){\\n    el.textContent=\\'公式渲染中...\\';el.classList.remove(\\'ln-no-mathjax\\');\\n    try{await ensureMathJax();el.innerHTML=markup;if(window.MathJax?.typesetPromise){await window.MathJax.typesetPromise([el])}else if(window.MathJax?.Hub?.Queue){window.MathJax.Hub.Queue([\\'Typeset\\',window.MathJax.Hub,el])}else{throw new Error(\\'MathJax unavailable\\')}}catch(error){el.textContent=\\'公式渲染失败：\\'+plain;el.classList.add(\\'ln-no-mathjax\\')}\\n  }\\n  function render(){const item=stages[state.kernel][state.step];renderMemory();$(\\'title\\').textContent=item.title;$(\\'detail\\').textContent=item.detail;setMath($(\\'formula\\'),item.formula,item.formulaLines.join(\\'\\\\n\\'));setMath($(\\'address\\'),item.address,item.addressLines.join(\\'\\\\n\\'));$(\\'code\\').textContent=item.codeLines.join(\\'\\\\n\\');$(\\'pid\\').max=state.kernel===\\'backward_params\\'?String(Math.ceil(N/BLOCK)-1):String(M-1);$(\\'pid\\').value=state.pid;$(\\'pid-value\\').textContent=state.pid;$(\\'step\\').max=String(stages[state.kernel].length-1);$(\\'step\\').value=state.step;$(\\'step-value\\').textContent=state.step;$(\\'status\\').textContent=`${item.title} · PID ${state.pid} · ${state.kernel}`}\\n  function stepBy(delta){state.step=(state.step+delta+stages[state.kernel].length)%stages[state.kernel].length;render()}$(\\'kernel\\').onchange=e=>{state.kernel=e.target.value;state.step=0;state.pid=0;render()};$(\\'pid\\').oninput=e=>{state.pid=+e.target.value;render()};$(\\'step\\').oninput=e=>{state.step=+e.target.value;render()};$(\\'previous\\').onclick=()=>stepBy(-1);$(\\'next\\').onclick=()=>stepBy(1);$(\\'play\\').onclick=()=>{if(state.playing){clearInterval(state.timer);state.playing=false;$(\\'play\\').textContent=\\'播放\\'}else{state.playing=true;$(\\'play\\').textContent=\\'暂停\\';state.timer=setInterval(()=>stepBy(1),900)}};render();\\n})();\\n</script>'\n",
    "\n",
    "# Python 字符串会保留 JavaScript 模板中的双反斜杠；还原为 MathJax 需要的单反斜杠。\n",
    "layernorm_visualization_html = layernorm_visualization_html.replace('\\\\\\\\', '\\\\').replace('String.raw`\\\\`', 'String.raw`\\\\\\\\`')\n",
    "_slash = chr(92)\n",
    "_math_start = layernorm_visualization_html.index('const math=')\n",
    "_math_end = layernorm_visualization_html.index(';', _math_start) + 1\n",
    "_math = \"const math=lines=>lines.map(line=>String.raw`\" + _slash + \"[${line}\" + _slash + \"]`).join('');\"\n",
    "layernorm_visualization_html = layernorm_visualization_html[:_math_start] + _math + layernorm_visualization_html[_math_end:]\n",
    "layernorm_visualization_html = layernorm_visualization_html.replace('el.innerHTML=markup;', \"el.innerHTML=markup.replace(/</g, '&lt;');\")\n",
    "display(HTML(layernorm_visualization_html))\n"
   ],
   "id": "4d51d8a88dd1f682",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "<style>\n",
       "#layernorm-pid-visualization{background:#fff;color:#172033;color-scheme:light;border:1px solid #cbd5e1;padding:16px;font-family:system-ui,'Microsoft YaHei',sans-serif;line-height:1.45}\n",
       "#layernorm-pid-visualization *{box-sizing:border-box}.ln-controls{display:flex;flex-wrap:wrap;gap:8px;align-items:end;margin-bottom:12px}.ln-controls label{display:flex;flex-direction:column;gap:3px;min-width:170px;font-size:13px}.ln-controls select,.ln-controls input,.ln-controls button{font:inherit}.ln-controls button{border:1px solid #94a3b8;background:#fff;color:#172033;padding:6px 10px;cursor:pointer}.ln-controls button.primary{background:#2563eb;color:#fff}.ln-status{padding:8px;background:#f8fafc;border:1px solid #cbd5e1;margin-bottom:10px}.ln-layout{display:grid;grid-template-columns:minmax(0,1.2fr) minmax(300px,.8fr);gap:14px}.ln-memory,.ln-panel{background:#fff;border:1px solid #cbd5e1;padding:10px;min-width:0}.ln-memory{overflow:auto}.ln-row{display:grid;grid-template-columns:120px repeat(8,minmax(48px,1fr));gap:3px;margin:3px 0;min-width:520px}.ln-label,.ln-cell{padding:5px 3px;border:1px solid #cbd5e1;text-align:center;font-size:12px}.ln-label{text-align:left;border-color:transparent;font-weight:600}.ln-cell.inactive{opacity:.35}.ln-cell.read{background:#bfdbfe;color:#12345b;border-color:#2563eb}.ln-cell.write{background:#bbf7d0;color:#14532d;border-color:#16a34a}.ln-cell.shared{background:#fde68a;color:#713f12;border-color:#d97706}.ln-cell.partial{background:#ddd6fe;color:#4c1d95;border-color:#7c3aed}.ln-cell.lock{background:#fecaca;color:#7f1d1d;border-color:#dc2626}.ln-cell.mask{background:repeating-linear-gradient(135deg,#e5e7eb 0,#e5e7eb 5px,#fff 5px,#fff 10px);color:#4b5563}.ln-panel h3{margin:0 0 8px;font-size:16px}.ln-panel h4{margin:12px 0 4px;font-size:13px;color:#334155}.ln-formula,.ln-address{background:#f8fafc;border-left:3px solid #2563eb;padding:8px;white-space:normal;overflow-x:auto;overflow-y:hidden;color:#172033}.ln-formula .MathJax,.ln-address .MathJax{max-width:100%;overflow-x:auto}.ln-code{margin:0;background:#f1f5f9;border-left:3px solid #64748b;padding:8px;white-space:pre;overflow-x:auto;overflow-y:auto;color:#172033;font:12px ui-monospace,SFMono-Regular,Consolas,monospace}.ln-legend{display:flex;flex-wrap:wrap;gap:5px;margin:8px 0;font-size:11px}.ln-legend span{padding:3px 6px;border:1px solid #cbd5e1}.ln-legend .read{background:#bfdbfe}.ln-legend .write{background:#bbf7d0}.ln-legend .shared{background:#fde68a}.ln-legend .partial{background:#ddd6fe}.ln-legend .lock{background:#fecaca}.ln-legend .mask{background:#e5e7eb}.ln-no-mathjax{font:12px ui-monospace,monospace;white-space:pre-wrap;color:#7f1d1d}@media(max-width:760px){.ln-layout{grid-template-columns:1fr}.ln-row{grid-template-columns:90px repeat(8,minmax(34px,1fr))}}\n",
       "</style>\n",
       "<div id=\"layernorm-pid-visualization\">\n",
       "<div class=\"ln-controls\"><label>内核<select data-role=\"kernel\"><option value=\"forward\">前向内核</option><option value=\"backward_input\">反向：输入与部分和</option><option value=\"backward_params\">反向：参数归约</option></select></label><label>PID <output data-role=\"pid-value\">0</output><input data-role=\"pid\" type=\"range\" min=\"0\" max=\"3\" value=\"0\"></label><label>步骤 <output data-role=\"step-value\">0</output><input data-role=\"step\" type=\"range\" min=\"0\" max=\"3\" value=\"0\"></label><button data-action=\"previous\">上一步</button><button class=\"primary\" data-action=\"play\">播放</button><button data-action=\"next\">下一步</button></div>\n",
       "<div class=\"ln-status\" data-role=\"status\"></div><div class=\"ln-legend\"><span class=\"read\">读取</span><span class=\"write\">写入</span><span class=\"shared\">共享参数</span><span class=\"partial\">部分和</span><span class=\"lock\">锁/计数</span><span class=\"mask\">mask lane</span></div>\n",
       "<div class=\"ln-layout\"><div class=\"ln-memory\" data-role=\"memory\"></div><div class=\"ln-panel\"><h3 data-role=\"title\"></h3><div data-role=\"detail\"></div><h4>数学公式</h4><div class=\"ln-formula\" data-role=\"formula\"></div><h4>地址映射</h4><div class=\"ln-address\" data-role=\"address\"></div><h4>对应 Triton 代码</h4><pre class=\"ln-code\" data-role=\"code\"></pre></div></div></div>\n",
       "<script>\n",
       "(() => {\n",
       "  const root=document.getElementById('layernorm-pid-visualization'); if(!root)return;\n",
       "  const $=role=>root.querySelector(`[data-role=\"${role}\"],[data-action=\"${role}\"]`);\n",
       "  const M=4,N=6,G=2,BLOCK=3;\n",
       "  const math=lines=>lines.map(line=>String.raw`\\[${line}\\]`).join('');\n",
       "  const stage=(title,detail,formula,address,code)=>({title,detail,formulaLines:formula,addressLines:address,formula:math(formula),address:math(address),codeLines:code});\n",
       "  const stages={\n",
       "    forward:[\n",
       "      stage('PID 行映射','grid=(M,)，一个 PID 负责一整行。',[String.raw`r=\\mathrm{program\\_id}(0)`],[String.raw`x_{row}=x_{base}+r\\cdot row\\_stride`,String.raw`y_{row}=y_{base}+r\\cdot row\\_stride`],['row = tl.program_id(0)','x_ptr += row * row_stride','y_ptr += row * row_stride']),\n",
       "      stage('均值、列偏移与 mask','有效 lane 参与均值归约，越界 lane 由 mask 屏蔽。',[String.raw`\\mu_r=\\frac{1}{N}\\sum_{j=0}^{N-1}x_{rj}`],[String.raw`columns_j=offset+j`,String.raw`addr(x_{rj})=x_{base}+r\\cdot row\\_stride+columns_j`,String.raw`addr(mean_r)=mean_{base}+r`,String.raw`mask_j=(columns_j<N)`],['columns = offset + tl.arange(0, block_size)','x = tl.load(x_ptr + columns, mask=mask, other=0.0)','mean = tl.sum(x, axis=0) / n_cols']),\n",
       "      stage('方差与 rstd','再次读取 x 和 mean，并写入当前行的 rstd。',[String.raw`\\sigma_r^2=\\frac{1}{N}\\sum_j(x_{rj}-\\mu_r)^2`,String.raw`rstd_r=(\\sigma_r^2+\\varepsilon)^{-1/2}`],[String.raw`addr(x_{rj})=x_{base}+r\\cdot row\\_stride+columns_j`,String.raw`addr(mean_r)=mean_{base}+r`,String.raw`addr(rstd_r)=rstd_{base}+r`,String.raw`difference_j=0\\text{ if }columns_j\\ge N`],['difference = tl.where(mask, x - mean, 0.0)','variance = tl.sum(difference * difference, axis=0) / n_cols','rstd = 1 / tl.sqrt(variance + eps)']),\n",
       "      stage('weight、bias 与输出写回','共享参数被所有 PID 读取，完成仿射变换并写回 y。',[String.raw`\\gamma_j=weight[j],\\quad \\beta_j=bias[j]`,String.raw`\\hat{x}_{rj}=(x_{rj}-\\mu_r)rstd_r`,String.raw`y_{rj}=\\gamma_j\\hat{x}_{rj}+\\beta_j`],[String.raw`addr(weight_j)=weight_{base}+j`,String.raw`addr(bias_j)=bias_{base}+j`,String.raw`addr(y_{rj})=y_{base}+r\\cdot row\\_stride+j`,String.raw`mask_j=(j<N)`],['weight = tl.load(weight_ptr + columns, mask=mask)','bias = tl.load(bias_ptr + columns, mask=mask)','y = normalized_x * weight + bias','tl.store(y_ptr + columns, y, mask=mask)'])\n",
       "    ],\n",
       "    backward_input:[\n",
       "      stage('反向行映射','同一 PID 对齐访问 x、grad_y、grad_x 及保存的统计量。',[String.raw`r=\\mathrm{program\\_id}(0)`,String.raw`\\hat{x}_{rj}=(x_{rj}-\\mu_r)rstd_r`],[String.raw`addr(x_{rj})=x_{base}+r\\cdot row\\_stride+j`,String.raw`addr(grad\\_y_{rj})=grad\\_y_{base}+r\\cdot row\\_stride+j`,String.raw`addr(grad\\_x_{rj})=grad\\_x_{base}+r\\cdot row\\_stride+j`,String.raw`addr(weight_j)=weight_{base}+j`,String.raw`addr(mean_r)=mean_{base}+r`,String.raw`addr(rstd_r)=rstd_{base}+r`,String.raw`mask_j=(j<N)`],['program_id = tl.program_id(0)','x_ptr += program_id * row_stride','grad_y_ptr += program_id * row_stride','grad_x_ptr += program_id * row_stride']),\n",
       "      stage('输入梯度行内归约','projection 和 grad_mean 都沿特征列归约。',[String.raw`g_{rj}=\\frac{\\partial L}{\\partial y_{rj}}\\gamma_j`,String.raw`\\frac{\\partial L}{\\partial x_{rj}}=rstd_r[g_{rj}-mean(g_r)-\\hat{x}_{rj}mean(\\hat{x}_rg_r)]`],[String.raw`read: x_{rj},\\ grad\\_y_{rj},\\ weight_j,\\ mean_r,\\ rstd_r`,String.raw`write: grad\\_x_{base}+r\\cdot row\\_stride+j`,String.raw`j\\in[0,N),\\quad mask_j=(j<N)`],['grad_normalized_x = weight * grad_y','projection = tl.sum(normalized_x * grad_normalized_x, axis=0) / n_cols','grad_mean = tl.sum(grad_normalized_x, axis=0) / n_cols','grad_x = (grad_normalized_x - normalized_x * projection - grad_mean) * rstd']),\n",
       "      stage('参数梯度部分和与锁','lock_id=r mod G；同组 PID 串行累加部分和。',[String.raw`\\Delta\\gamma_{rj}=\\frac{\\partial L}{\\partial y_{rj}}\\hat{x}_{rj}`,String.raw`\\Delta\\beta_{rj}=\\frac{\\partial L}{\\partial y_{rj}}`,String.raw`lock\\_id=r\\bmod G`],[String.raw`lock\\_addr=locks_{base}+lock\\_id`,String.raw`count\\_addr=locks_{base}+G+lock\\_id`,String.raw`addr(P^\\gamma_{gj})=grad\\_weight\\_partial_{base}+gN+j`,String.raw`addr(P^\\beta_{gj})=grad\\_bias\\_partial_{base}+gN+j`],['lock_id = program_id % group_size','locks_ptr += lock_id','count_ptr = locks_ptr + group_size','tl.atomic_cas(locks_ptr, 0, 1)'])\n",
       "    ],\n",
       "    backward_params:[\n",
       "      stage('参数归约列块','PID q 负责一个列块；教学图 BLOCK_SIZE_N=3，源码为 128。',[String.raw`Q=\\left\\lceil\\frac{N}{BLOCK\\_SIZE\\_N}\\right\\rceil`,String.raw`q=\\mathrm{program\\_id}(0)`],[String.raw`columns_j=q\\cdot BLOCK\\_SIZE\\_N+j`,String.raw`0\\le columns_j<N`,String.raw`output\\_mask_j=(columns_j<N)`],['q = tl.program_id(0)','columns = q * block_size_n + tl.arange(0, block_size_n)']),\n",
       "      stage('二维广播读取部分和','rows[:,None] 与 columns[None,:] 形成二维地址。',[String.raw`offset_{g,j}=gN+columns_j`,String.raw`g\\in[0,G)`],[String.raw`addr(P^\\gamma_{g,j})=grad\\_weight\\_partial_{base}+gN+columns_j`,String.raw`addr(P^\\beta_{g,j})=grad\\_bias\\_partial_{base}+gN+columns_j`,String.raw`mask_{g,j}=(g<G)\\land(columns_j<N)`],['rows = tl.arange(0, group_size)','columns = q * block_size_n + tl.arange(0, block_size_n)','offsets = rows[:, None] * n_cols + columns[None, :]']),\n",
       "      stage('归约并写回最终参数梯度','沿 group 维归约并写回当前列块。',[String.raw`\\frac{\\partial L}{\\partial\\gamma_j}=\\sum_gP^\\gamma_{gj}`,String.raw`\\frac{\\partial L}{\\partial\\beta_j}=\\sum_gP^\\beta_{gj}`],[String.raw`addr(grad\\_weight_j)=grad\\_weight_{base}+columns_j`,String.raw`addr(grad\\_bias_j)=grad\\_bias_{base}+columns_j`,String.raw`output\\_mask_j=(columns_j<N)`],['grad_weight = tl.sum(partial_weight, axis=0)','grad_bias = tl.sum(partial_bias, axis=0)','tl.store(grad_weight_ptr + columns, grad_weight, mask=output_mask)','tl.store(grad_bias_ptr + columns, grad_bias, mask=output_mask'])\n",
       "    ]\n",
       "  };\n",
       "  const state={kernel:'forward',step:0,pid:0,playing:false,timer:null};\n",
       "  const active=(len,start,end,kind)=>Array.from({length:len},(_,i)=>i>=start&&i<end?kind:'inactive');\n",
       "  const row=(label,values,states)=>`<div class=\"ln-row\"><div class=\"ln-label\">${label}</div>${values.map((v,i)=>`<div class=\"ln-cell ${states[i]||'inactive'}\">${v}</div>`).join('')}</div>`;\n",
       "  const matrix=(name,r,kind)=>row(`${name}[${r},:]`,Array.from({length:N},(_,j)=>`${name}[${r},${j}]`).concat(['mask','mask']),active(N+2,0,N,kind).map((x,i)=>i>=N?'mask':x));\n",
       "  function renderMemory(){const k=state.kernel,s=state.step,p=state.pid,blocks=[],lock=p%G;if(k==='forward'){for(let r=0;r<M;r++)blocks.push(matrix('x',r,r===p&&s>=0?'read':'inactive'));blocks.push(row('mean[M]',Array.from({length:M},(_,i)=>`mu[${i}]`),active(M,p,p+1,s===1?'write':s>=2?'read':'inactive')));blocks.push(row('rstd[M]',Array.from({length:M},(_,i)=>`rstd[${i}]`),active(M,p,p+1,s===2?'write':s===3?'read':'inactive')));blocks.push(row('weight[N]',Array.from({length:N},(_,i)=>`gamma[${i}]`),active(N,0,N,s===3?'shared':'inactive')));blocks.push(row('bias[N]',Array.from({length:N},(_,i)=>`beta[${i}]`),active(N,0,N,s===3?'shared':'inactive')));for(let r=0;r<M;r++)blocks.push(matrix('y',r,s===3&&r===p?'write':'inactive'))}else if(k==='backward_input'){for(const n of ['x','grad_y','grad_x'])for(let r=0;r<M;r++)blocks.push(matrix(n,r,r===p?(n==='grad_x'&&s>=1?'write':'read'):'inactive'));blocks.push(row('weight[N]',Array.from({length:N},(_,i)=>`gamma[${i}]`),active(N,0,N,s>=1?'shared':'inactive')));for(const n of ['mean','rstd'])blocks.push(row(`${n}[M]`,Array.from({length:M},(_,i)=>`${n==='mean'?'mu':'rstd'}[${i}]`),active(M,p,p+1,'read')));for(let g=0;g<G;g++){blocks.push(row(`Pgamma[${g},:]`,Array.from({length:N},(_,i)=>`Pgamma[${g},${i}]`),active(N,0,N,s===2&&g===lock?'partial':'inactive')));blocks.push(row(`Pbeta[${g},:]`,Array.from({length:N},(_,i)=>`Pbeta[${g},${i}]`),active(N,0,N,s===2&&g===lock?'partial':'inactive')))}blocks.push(row('locks[0:G]',Array.from({length:G},(_,i)=>`lock[${i}]`),active(G,lock,lock+1,s===2?'lock':'inactive')));blocks.push(row('locks[G:2G]',Array.from({length:G},(_,i)=>`count[${i}]`),active(G,lock,lock+1,s===2?'lock':'inactive')))}else{const start=p*BLOCK,end=Math.min(N,start+BLOCK);for(let g=0;g<G;g++){blocks.push(row(`Pgamma[${g},:]`,Array.from({length:N},(_,i)=>`Pgamma[${g},${i}]`),active(N,start,end,s>=1?'partial':'inactive')));blocks.push(row(`Pbeta[${g},:]`,Array.from({length:N},(_,i)=>`Pbeta[${g},${i}]`),active(N,start,end,s>=1?'partial':'inactive')))}blocks.push(row('grad_weight[N]',Array.from({length:N},(_,i)=>`dgamma[${i}]`),active(N,start,end,s===2?'write':'inactive')));blocks.push(row('grad_bias[N]',Array.from({length:N},(_,i)=>`dbeta[${i}]`),active(N,start,end,s===2?'write':'inactive')))}$('memory').innerHTML=blocks.join('')}\n",
       "  let mathJaxLoader=null;\n",
       "  function ensureMathJax(){\n",
       "    if(window.MathJax?.typesetPromise||window.MathJax?.Hub?.Queue)return Promise.resolve();\n",
       "    if(mathJaxLoader)return mathJaxLoader;\n",
       "    window.MathJax={tex:{displayMath:[['\\[','\\]']]},startup:{typeset:false}};\n",
       "    mathJaxLoader=new Promise((resolve,reject)=>{const script=document.createElement('script');script.src='https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-svg.js';script.async=true;script.onload=resolve;script.onerror=reject;document.head.appendChild(script)});\n",
       "    return mathJaxLoader;\n",
       "  }\n",
       "  async function setMath(el,markup,plain){\n",
       "    el.textContent='公式渲染中...';el.classList.remove('ln-no-mathjax');\n",
       "    try{await ensureMathJax();el.innerHTML=markup.replace(/</g, '&lt;');if(window.MathJax?.typesetPromise){await window.MathJax.typesetPromise([el])}else if(window.MathJax?.Hub?.Queue){window.MathJax.Hub.Queue(['Typeset',window.MathJax.Hub,el])}else{throw new Error('MathJax unavailable')}}catch(error){el.textContent='公式渲染失败：'+plain;el.classList.add('ln-no-mathjax')}\n",
       "  }\n",
       "  function render(){const item=stages[state.kernel][state.step];renderMemory();$('title').textContent=item.title;$('detail').textContent=item.detail;setMath($('formula'),item.formula,item.formulaLines.join('\\n'));setMath($('address'),item.address,item.addressLines.join('\\n'));$('code').textContent=item.codeLines.join('\\n');$('pid').max=state.kernel==='backward_params'?String(Math.ceil(N/BLOCK)-1):String(M-1);$('pid').value=state.pid;$('pid-value').textContent=state.pid;$('step').max=String(stages[state.kernel].length-1);$('step').value=state.step;$('step-value').textContent=state.step;$('status').textContent=`${item.title} · PID ${state.pid} · ${state.kernel}`}\n",
       "  function stepBy(delta){state.step=(state.step+delta+stages[state.kernel].length)%stages[state.kernel].length;render()}$('kernel').onchange=e=>{state.kernel=e.target.value;state.step=0;state.pid=0;render()};$('pid').oninput=e=>{state.pid=+e.target.value;render()};$('step').oninput=e=>{state.step=+e.target.value;render()};$('previous').onclick=()=>stepBy(-1);$('next').onclick=()=>stepBy(1);$('play').onclick=()=>{if(state.playing){clearInterval(state.timer);state.playing=false;$('play').textContent='播放'}else{state.playing=true;$('play').textContent='暂停';state.timer=setInterval(()=>stepBy(1),900)}};render();\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "14a66fb1f65f624e"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
