diff --git a/_extra/viz-base.js b/_extra/viz-base.js index f3339bc2..ce3e1a3f 100644 --- a/_extra/viz-base.js +++ b/_extra/viz-base.js @@ -3,6 +3,17 @@ document.addEventListener('DOMContentLoaded', function() { var p = new URLSearchParams(location.search); if (p.has('notitle')) document.body.classList.add('notitle'); + // Directly opened demo pages are copied outside Sphinx templates, so they do + // not receive html_js_files. Load the same language switch used by book pages. + if (window.parent === window && !p.has('notitle')) { + var langSwitchPath = window.location.pathname.indexOf('/zh/demo/') >= 0 + ? '../../_static/lang-switch.js' + : '../_static/lang-switch.js'; + var s = document.createElement('script'); + s.src = new URL(langSwitchPath, window.location.href).href; + document.head.appendChild(s); + } + // Forward arrow keys to parent (reveal.js) when embedded if (window.parent !== window) { document.addEventListener('keydown', function(e) { diff --git a/_extra_zh/zh/_static/tirx-layout-demo/index.html b/_extra_zh/zh/_static/tirx-layout-demo/index.html new file mode 100644 index 00000000..ae2ec85e --- /dev/null +++ b/_extra_zh/zh/_static/tirx-layout-demo/index.html @@ -0,0 +1,167 @@ + + + + + + +TIRx Tensor Layout — Interactive Demo + + + + + + +

TIRx Tensor Layout

+
how a TileLayout (shard / replica / offset) maps logical tensor elements to physical threads
+ +
+ load an example + + ↓ fills the two fields below +
+ +
+
+ + +
+
+ +
+ + +
+
+
+ + +
+
+
+ +
+
+

Logical tensor

+
+
+
+
+

Physical threads

+
+
+
+ +
+ +
+
+ + + + diff --git a/_extra_zh/zh/_static/tirx-layout-demo/layout-demo.js b/_extra_zh/zh/_static/tirx-layout-demo/layout-demo.js new file mode 100644 index 00000000..ec117acc --- /dev/null +++ b/_extra_zh/zh/_static/tirx-layout-demo/layout-demo.js @@ -0,0 +1,805 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * TIRx tensor-layout visualizer. + * + * The page shell, CSS vocabulary, and the draw()/hover/arrow interaction + * pattern are adapted from the team's own course material + * (mlsyscourse/slides-modern-gpu-programming, data-layout/site/demo/ + * tile_distributed.html). The TIRx S/R/O parser and the logical->physical + * mapper below are original and mirror tvm/python/tvm/tirx/layout.py + * (_flatten_coord / _split_coord and the TileLayout forward mapping). + * This file is NOT derived from any third-party layout demo. + */ + +'use strict'; + +// ── TIRx named axes (from tvm/python/tvm/tirx/layout.py _AXIS_NAMES, plus the +// device axis `pid` used by distributed layouts) ────────────────────────────── +const AXIS_ORDER = [ + 'pid', 'bx', 'by', 'bz', 'cbx', 'cby', 'cbz', 'tx', + 'warpid', 'laneid', 'wgid', 'tid_in_wg', 'wid_in_wg', 'tid', + 'm', 'P', 'F', 'Bank', 'TCol', 'TLane', +]; +// "Owner" axes name a physical unit that owns data (threads, devices); the rest +// (m, P, F, Bank, TCol, TLane) are storage/memory coordinates within an owner. +const OWNER_AXES = new Set([ + 'pid', 'bx', 'by', 'bz', 'cbx', 'cby', 'cbz', 'tx', + 'warpid', 'laneid', 'wgid', 'tid_in_wg', 'wid_in_wg', 'tid', +]); +const KNOWN_AXES = new Set(AXIS_ORDER); +const MAX_ELEMENTS = 1024; // render cap + +function isOwnerAxis(a) { return OWNER_AXES.has(a); } +function product(arr) { return arr.reduce((a, b) => a * b, 1); } + +// ── Parser ─────────────────────────────────────────────────────────────────── +// Grammar mirrors layout.py: S[shape:stride] + R[shape:stride] + offset, +// stride/offset terms are "n@axis" (a bare int defaults to axis "m"). + +function splitTopLevel(s, sep) { + const out = []; + let depth = 0, cur = ''; + for (const ch of s) { + if (ch === '(' || ch === '[') depth++; + else if (ch === ')' || ch === ']') { + depth--; + if (depth < 0) throw new Error('unmatched closing bracket or parenthesis'); + } + if (ch === sep && depth === 0) { out.push(cur); cur = ''; } + else cur += ch; + } + if (depth !== 0) throw new Error('unmatched opening bracket or parenthesis'); + out.push(cur); + return out; +} + +function stripParens(s) { + s = s.trim(); + if (s.startsWith('(') && s.endsWith(')')) return s.slice(1, -1); + return s; +} + +function parseIntStrict(s) { + const t = s.trim(); + if (!/^-?\d+$/.test(t)) throw new Error(`expected integer, got "${t}"`); + return parseInt(t, 10); +} + +function parseTerm(tok) { + const t = tok.trim(); + if (t.includes('@')) { + const parts = t.split('@'); + if (parts.length !== 2) throw new Error(`bad term "${t}"`); + const numPart = parts[0].trim(); + const axis = parts[1].trim(); + if (!KNOWN_AXES.has(axis)) throw new Error(`unknown axis "${axis}"`); + const stride = numPart === '' ? 1 : parseIntStrict(numPart); + return { stride, axis }; + } + if (/^-?\d+$/.test(t)) return { stride: parseInt(t, 10), axis: 'm' }; + if (KNOWN_AXES.has(t)) return { stride: 1, axis: t }; + throw new Error(`bad term "${t}"`); +} + +function defaultStrides(extents) { + const n = extents.length; + const strides = new Array(n).fill(1); + for (let i = n - 2; i >= 0; i--) strides[i] = strides[i + 1] * extents[i + 1]; + return strides; +} + +function parseExtents(s) { + // Parse a comma-separated extent list and reject non-positive / oversized + // extents, so bad input (e.g. R[1000000:...] or a 0/negative extent) can't + // NaN-propagate or freeze the tab with huge loops in physOwners(). + const extents = splitTopLevel(stripParens(s), ',').map(parseIntStrict); + for (const e of extents) { + if (e <= 0) throw new Error(`extent must be positive, got ${e}`); + if (e > MAX_ELEMENTS) throw new Error(`extent ${e} exceeds maximum ${MAX_ELEMENTS}`); + } + return extents; +} + +function parseBracket(inner) { + const parts = splitTopLevel(inner, ':'); + if (parts.length === 1) { + const extents = parseExtents(parts[0]); + const strides = defaultStrides(extents); + return extents.map((e, i) => ({ extent: e, stride: strides[i], axis: 'm' })); + } + if (parts.length !== 2) throw new Error('layout bracket must be "shape : stride"'); + const extents = parseExtents(parts[0]); + const terms = splitTopLevel(stripParens(parts[1]), ',').map(parseTerm); + if (extents.length !== terms.length) { + throw new Error(`shape has ${extents.length} dims but stride has ${terms.length}`); + } + return extents.map((e, i) => ({ extent: e, stride: terms[i].stride, axis: terms[i].axis })); +} + +function bracketBody(piece, prefix) { + const open = piece.indexOf('['); + const close = piece.lastIndexOf(']'); + if (open < 0 || close < 0 || close < open) throw new Error(`malformed ${prefix}[...]`); + return piece.slice(open + 1, close); +} + +function parseSwizzlePrefix(src) { + // Optional "Swizzle(per_element, swizzle_len, atom_len[, inner]) [∘|o|*] ". + const m = src.match(/^Swizzle\s*\(([^)]*)\)\s*(?:∘|o|\.|\*)?\s*([\s\S]*)$/i); + if (!m) return { swizzle: null, rest: src }; + const a = m[1].split(',').map((s) => s.trim()).filter((s) => s.length); + if (a.length < 3) throw new Error('Swizzle needs (per_element, swizzle_len, atom_len)'); + const per_element = parseIntStrict(a[0]); + const swizzle_len = parseIntStrict(a[1]); + const atom_len = parseIntStrict(a[2]); + if (per_element < 0 || swizzle_len < 0 || atom_len < swizzle_len + || per_element >= 31 || atom_len >= 31) { + // atom_len/per_element feed 32-bit bitwise shifts in swizzleAddr; cap < 31. + throw new Error('swizzle requires 0≤per_element<31, swizzle_len≥0, swizzle_len≤atom_len<31'); + } + const inner = a[3] === undefined ? true : (a[3] === 'true' || a[3] === '1'); + return { swizzle: { per_element, swizzle_len, atom_len, inner }, rest: m[2].trim() }; +} + +function parseLayout(srcRaw) { + const { swizzle, rest } = parseSwizzlePrefix(srcRaw.trim()); + const src = rest; + const layout = { shard: [], replica: [], offset: {}, swizzle }; + let sawShard = false; + for (let piece of splitTopLevel(src, '+')) { + piece = piece.trim(); + if (piece === '') continue; + if (piece.startsWith('S[')) { + layout.shard = parseBracket(bracketBody(piece, 'S')); + sawShard = true; + } else if (piece.startsWith('R[')) { + layout.replica = parseBracket(bracketBody(piece, 'R')); + } else { + const t = parseTerm(piece); + layout.offset[t.axis] = (layout.offset[t.axis] || 0) + t.stride; + } + } + if (!sawShard) throw new Error('layout needs a shard term, e.g. S[...]'); + return layout; +} + +// ── Mapper (mirrors layout.py _flatten_coord / _split_coord + forward map) ───── + +function flattenCoord(coord, shape) { + let flat = 0; + for (let i = 0; i < shape.length; i++) flat = flat * shape[i] + coord[i]; + return flat; +} + +function splitCoord(flat, extents) { + const n = extents.length; + const res = new Array(n); + let remaining = flat; + for (let i = n - 1; i >= 0; i--) { + if (i === 0) res[0] = remaining; + else { res[i] = remaining % extents[i]; remaining = Math.floor(remaining / extents[i]); } + } + return res; +} + +function coordFromFlat(flat, shape) { return splitCoord(flat, shape); } + +function forwardBase(coord, shape, layout) { + const flat = flattenCoord(coord, shape); + const comps = splitCoord(flat, layout.shard.map((it) => it.extent)); + const phys = {}; + for (let k = 0; k < layout.shard.length; k++) { + const it = layout.shard[k]; + phys[it.axis] = (phys[it.axis] || 0) + comps[k] * it.stride; + } + for (const axis of Object.keys(layout.offset)) { + phys[axis] = (phys[axis] || 0) + layout.offset[axis]; + } + return phys; +} + +// Replica broadcasts the same logical element onto multiple physical owners: +// L(x) = { D(x) + r + O | r in R }. +function physOwners(coord, shape, layout) { + let owners = [forwardBase(coord, shape, layout)]; + for (const rep of layout.replica) { + const next = []; + for (const o of owners) { + for (let k = 0; k < rep.extent; k++) { + const o2 = Object.assign({}, o); + o2[rep.axis] = (o2[rep.axis] || 0) + k * rep.stride; + next.push(o2); + } + } + owners = next; + } + return owners; +} + +function axesUsed(layout) { + const s = new Set(); + for (const it of layout.shard) s.add(it.axis); + for (const it of layout.replica) s.add(it.axis); + for (const a of Object.keys(layout.offset)) s.add(a); + return AXIS_ORDER.filter((a) => s.has(a)); +} + +function coordStr(phys, axes) { + return axes.map((a) => `${a}=${phys[a] || 0}`).join(' '); +} + +// Swizzle a linear memory address (mirrors src/tirx/ir/layout/swizzle_layout.cc +// SwizzleLayoutNode::Apply): low `per_element` bits are kept; above them, the +// swizzle bits are XOR'd to scatter bank conflicts. +function swizzleAddr(m, sw) { + const base = 1 << sw.per_element; + const innerMask = (1 << sw.swizzle_len) - 1; + const outerMask = innerMask << sw.atom_len; + const x = Math.floor(m / base); + const fx = sw.inner ? (x ^ ((x & outerMask) >> sw.atom_len)) + : (x ^ ((x & innerMask) << sw.atom_len)); + return fx * base + (m % base); +} + +// Resolve the swizzle from the dtype + mode dropdowns (mirrors +// tma_utils.mma_atom_layout): per_element = bit_length(128//bits) - 1, +// swizzle_len = mode, atom_len = 3. Falls back to a typed Swizzle(...) prefix. +const SWIZZLE_LEN = { none: 0, '32': 1, '64': 2, '128': 3 }; +function computeSwizzle() { + const mode = swmodeSel ? swmodeSel.value : 'off'; + if (mode === 'off') return (ST.layout && ST.layout.swizzle) ? ST.layout.swizzle : null; + const bits = +(dtypeSel ? dtypeSel.value : 16) || 16; + const per_element = Math.floor(128 / bits).toString(2).length - 1; + return { + per_element, swizzle_len: SWIZZLE_LEN[mode] || 0, atom_len: 3, inner: true, + bits, mode, + }; +} + +// ── State + recompute ──────────────────────────────────────────────────────-- +function getComputedStyleVar(name) { + return getComputedStyle(document.documentElement).getPropertyValue(name).trim(); +} +const PALETTE = Array.from({ length: 8 }, (_, i) => + getComputedStyleVar(`--color-group-${i}`) || '#5b9bd5'); +function paletteColor(v) { const n = PALETTE.length; return PALETTE[((v % n) + n) % n]; } + +const ST = { + shape: [4, 8], + layout: null, + error: null, + tooBig: false, + banks: 32, + swizzle: null, + gridAxes: [], yAxis: null, xAxis: null, cellAxes: [], + yVals: [], xVals: [], + byFlat: [], // flat -> { owners:[phys], keys:[gridKey], color } + byCell: new Map(), // "y#x" -> [{flat, slot}] +}; +let hovFlat = null; +let drawing = false; + +function mk(t, c) { const d = document.createElement(t); if (c) d.className = c; return d; } +function gridKey(phys, axes) { return axes.map((a) => phys[a] || 0).join(','); } + +function recompute() { + ST.error = null; ST.tooBig = false; + try { + ST.shape = splitTopLevel(stripParens(shapeInput.value), ',').map(parseIntStrict); + if (ST.shape.length === 0 || ST.shape.some((x) => x <= 0)) throw new Error('shape must be positive ints'); + ST.layout = parseLayout(exprInput.value); + } catch (e) { ST.error = e.message; return; } + + const total = product(ST.shape); + if (total > MAX_ELEMENTS) { ST.tooBig = true; return; } + + ST.shapeTotal = total; + ST.shardTotal = product(ST.layout.shard.map((it) => it.extent)); + ST.mismatch = ST.shardTotal !== total; + ST.swizzle = computeSwizzle(); + // elements that share one 4-byte bank word (e.g. 2 fp16, 4 fp8, 1 fp32) + ST.elemsPerBank = ST.swizzle ? Math.max(1, Math.round(4 / ((ST.swizzle.bits || 32) / 8))) : 1; + + // 1) owners per element; in swizzle mode, also map the memory address through + // the swizzle and derive synthetic line/bank coordinates. + ST.byFlat = new Array(total); + for (let flat = 0; flat < total; flat++) { + const owners = physOwners(coordFromFlat(flat, ST.shape), ST.shape, ST.layout); + if (ST.swizzle) { + // A shared-memory bank is 4 bytes; an element occupies dtype_bytes, so the + // bank word index is floor(element_addr * dtype_bytes / 4). 32 banks per line. + const bytes = (ST.swizzle.bits || 32) / 8; + for (const o of owners) { + const sm = swizzleAddr(o.m || 0, ST.swizzle); + const word = Math.floor((sm * bytes) / 4); + o.__sm = sm; o.__word = word; o.bank = word % ST.banks; o.line = Math.floor(word / ST.banks); + } + } + ST.byFlat[flat] = { owners }; + } + + // 2) choose grid + color axes + if (ST.swizzle) { + ST.gridAxes = ['line', 'bank']; ST.yAxis = 'line'; ST.xAxis = 'bank'; + ST.cellAxes = []; ST.colorAxis = 'bank'; + } else { + const used = axesUsed(ST.layout); + const owners = used.filter(isOwnerAxis); + ST.gridAxes = owners.length ? owners : used.filter((a) => !isOwnerAxis(a)); + ST.yAxis = ST.gridAxes[0] || null; + ST.xAxis = ST.gridAxes[1] || null; + ST.cellAxes = used.filter((a) => !ST.gridAxes.includes(a)); + // color axis = first grid axis from shard/offset, so a replica-only row axis + // doesn't collapse every element to one color. + const shardOffsetAxes = new Set(ST.layout.shard.map((it) => it.axis)); + for (const a of Object.keys(ST.layout.offset)) shardOffsetAxes.add(a); + ST.colorAxis = ST.gridAxes.find((a) => shardOffsetAxes.has(a)) || ST.gridAxes[0] || null; + } + + // 3) build cells, hover keys, colors + ST.byCell = new Map(); + const yset = new Set(), xset = new Set(), cset = new Set(); + for (let flat = 0; flat < total; flat++) { + const rec = ST.byFlat[flat]; + const keys = []; + for (const o of rec.owners) { + const y = ST.yAxis ? (o[ST.yAxis] || 0) : 0; + const x = ST.xAxis ? (o[ST.xAxis] || 0) : 0; + yset.add(y); xset.add(x); + keys.push(gridKey(o, ST.gridAxes)); + const ck = y + '#' + x; + if (!ST.byCell.has(ck)) ST.byCell.set(ck, []); + ST.byCell.get(ck).push({ flat, slot: ST.swizzle ? ('addr ' + o.__sm) : coordStr(o, ST.cellAxes) }); + } + rec.keys = keys; + const cv = ST.colorAxis ? (rec.owners[0][ST.colorAxis] || 0) : 0; + rec.color = paletteColor(cv); + cset.add(cv); + } + ST.yVals = [...yset].sort((a, b) => a - b); + ST.xVals = [...xset].sort((a, b) => a - b); + ST.colorVals = [...cset].sort((a, b) => a - b); +} + +// ── Display geometry for the logical grid ──────────────────────────────────-- +function logicalGridDims() { + if (ST.shape.length === 2) return { rows: ST.shape[0], cols: ST.shape[1] }; + if (ST.shape.length === 1) return { rows: 1, cols: ST.shape[0] }; + return { rows: 1, cols: product(ST.shape) }; // N-D -> flat strip +} + +// ── Draw ──────────────────────────────────────────────────────────────────── +function resetFit() { + const p = document.getElementById('panels'); + if (p) { p.style.transform = 'none'; p.style.marginBottom = ''; } +} +function fitEmbed() { + if (!document.body.classList.contains('lock')) return; + const p = document.getElementById('panels'); + if (!p) return; + const natural = p.offsetWidth; + const pad = 2 * parseFloat(getComputedStyle(document.body).paddingLeft || '0'); + const avail = document.documentElement.clientWidth - pad; + if (avail > 0 && natural > avail) { + const sc = avail / natural; + p.style.transformOrigin = 'top left'; + p.style.transform = 'scale(' + sc + ')'; + p.style.marginBottom = (-(p.offsetHeight * (1 - sc))) + 'px'; + } +} +function postHeight() { + if (window.parent === window) return; + const h = Math.ceil(document.body.scrollHeight); + window.parent.postMessage({ tirxLayoutDemoHeight: h + 4 }, '*'); +} +function draw() { + drawing = true; + resetFit(); + const status = document.getElementById('status'); + const g0 = document.getElementById('g0'); + const phys = document.getElementById('phys'); + const fb = document.getElementById('fb'); + const lg = document.getElementById('lg'); + + if (ST.error) { + status.innerHTML = `parse error: ${escapeHtml(ST.error)}`; + g0.innerHTML = ''; phys.innerHTML = ''; lg.innerHTML = ''; + fb.innerHTML = '
Fix the layout expression to continue.
'; + setTimeout(() => { drawing = false; }, 0); return; + } + if (ST.tooBig) { + status.innerHTML = `${product(ST.shape)} elements exceeds the ${MAX_ELEMENTS} render cap — use a smaller shape.`; + g0.innerHTML = ''; phys.innerHTML = ''; lg.innerHTML = ''; + fb.innerHTML = '
Shape too large to visualize.
'; + setTimeout(() => { drawing = false; }, 0); return; + } + + status.innerHTML = `ok   ` + + `${product(ST.shape)} logical elements  |  ` + + `${ST.yVals.length * (ST.xVals.length || 1)} physical cells`; + if (ST.mismatch) { + status.innerHTML += `  ` + + `⚠ shard total ${ST.shardTotal} ≠ shape total ${ST.shapeTotal} — mapping may be ill-formed`; + } + if (ST.swizzle) { + const s = ST.swizzle; + const label = s.mode ? (s.mode === 'none' ? 'no swizzle' : s.mode + 'B swizzle') : 'swizzle'; + status.innerHTML += `  ` + + `${label}${s.bits ? ', ' + s.bits + '-bit' : ''} → Swizzle(${s.per_element},${s.swizzle_len},${s.atom_len})`; + } + document.getElementById('n0').textContent = `logical shape (${ST.shape.join(', ')})`; + document.getElementById('nphys').textContent = + (ST.xAxis ? `rows = ${ST.yAxis}, cols = ${ST.xAxis}` : `tiles = ${ST.yAxis || '(none)'}`) + + (ST.cellAxes.length ? ' · in-cell: ' + ST.cellAxes.join(', ') : ''); + + const hovKeys = hovFlat !== null ? new Set(ST.byFlat[hovFlat].keys) : null; + drawLogical(hovKeys); + drawPhysical(hovKeys); + drawFormula(); + drawArrow(); + drawLegend(); + fitEmbed(); + postHeight(); + setTimeout(() => { drawing = false; }, 0); +} + +function sharesOwner(flat, hovKeys) { + if (!hovKeys) return false; + return ST.byFlat[flat].keys.some((k) => hovKeys.has(k)); +} + +function drawLogical(hovKeys) { + const g = document.getElementById('g0'); + g.innerHTML = ''; + const { rows, cols } = logicalGridDims(); + g.style.gridTemplateColumns = '30px repeat(' + cols + ', 46px)'; + g.appendChild(mk('div', 'hdr')); + for (let c = 0; c < cols; c++) { const h = mk('div', 'hdr'); h.textContent = 'c' + c; g.appendChild(h); } + for (let r = 0; r < rows; r++) { + const rl = mk('div', 'rl'); rl.textContent = (rows > 1) ? ('r' + r) : ''; g.appendChild(rl); + for (let c = 0; c < cols; c++) { + const flat = r * cols + c; + const d = mk('div', 'cell'); + d.dataset.flat = flat; + d.textContent = flat; + d.style.background = ST.byFlat[flat].color; + d.style.color = '#fff'; + if (hovFlat !== null) { + if (flat === hovFlat) d.classList.add('hov'); + else if (!sharesOwner(flat, hovKeys)) d.classList.add('dm'); + } + g.appendChild(d); + } + } +} + +function cellEntries(y, x) { return ST.byCell.get(y + '#' + x) || []; } + +function makeSlot(entry) { + const s = mk('div', 'gcell'); + s.style.background = ST.byFlat[entry.flat].color; + s.style.color = '#fff'; + s.dataset.flat = entry.flat; + s.textContent = entry.flat; + s.title = entry.slot ? `element ${entry.flat} @ ${entry.slot}` : `element ${entry.flat}`; + if (hovFlat !== null) { + if (entry.flat === hovFlat) s.classList.add('hov'); + else s.classList.add('dm'); + } + return s; +} + +function drawPhysical(hovKeys) { + const wrap = document.getElementById('phys'); + wrap.innerHTML = ''; + if (!ST.yAxis) { wrap.textContent = '(no physical axes)'; return; } + + if (ST.xAxis) { + // 2D table: rows = yAxis values, cols = xAxis values. + const table = mk('div', 'phys-table' + (ST.swizzle ? ' bank-mode' : '')); + // In bank mode a cell is one 4-byte bank word holding elemsPerBank elements + // laid out horizontally; otherwise a fixed 54px cell. + const colW = ST.swizzle ? (ST.elemsPerBank * 48 + 8) : 54; + table.style.gridTemplateColumns = '44px repeat(' + ST.xVals.length + ', ' + colW + 'px)'; + table.appendChild(corner()); + for (const x of ST.xVals) table.appendChild(axHdr(String(x), false, `${ST.xAxis}=${x}`)); + for (const y of ST.yVals) { + table.appendChild(axHdr(String(y), true, `${ST.yAxis}=${y}`)); + for (const x of ST.xVals) { + const cell = mk('div', 'pcell'); + const entries = cellEntries(y, x); + if (hovFlat !== null && entries.some((e) => e.flat === hovFlat)) cell.classList.add('hov-cell'); + else if (hovFlat !== null && entries.length) cell.classList.add('dm-cell'); + for (const e of entries) cell.appendChild(makeSlot(e)); + table.appendChild(cell); + } + } + wrap.appendChild(table); + } else { + // 1D: wrapped list of owner tiles, one per yAxis value. + const list = mk('div', 'phys-1d'); + for (const y of ST.yVals) { + const tile = mk('div', 'thread-tile'); + const lbl = mk('div', 'thread-lbl'); lbl.textContent = `${ST.yAxis}=${y}`; tile.appendChild(lbl); + const slots = mk('div', 'thread-slots'); + const entries = cellEntries(y, 0).slice().sort((a, b) => a.flat - b.flat); + if (hovFlat !== null && entries.some((e) => e.flat === hovFlat)) tile.classList.add('hov-tile'); + else if (hovFlat !== null && entries.length) tile.classList.add('dm-tile'); + for (const e of entries) slots.appendChild(makeSlot(e)); + tile.appendChild(slots); + list.appendChild(tile); + } + wrap.appendChild(list); + } +} + +function corner() { + const d = mk('div', 'ax-hdr corner'); + d.textContent = '↘'; + if (ST.yAxis) d.title = `rows = ${ST.yAxis}` + (ST.xAxis ? `, cols = ${ST.xAxis}` : ''); + return d; +} +function axHdr(text, isRow, title) { + const d = mk('div', 'ax-hdr' + (isRow ? ' row' : '')); + d.textContent = text; + if (title) d.title = title; + return d; +} + +function drawFormula() { + const fb = document.getElementById('fb'); + if (hovFlat === null) { fb.innerHTML = '
Click a logical element to see its mapping.
'; return; } + const flat = hovFlat; + const coord = coordFromFlat(flat, ST.shape); + const comps = splitCoord(flat, ST.layout.shard.map((it) => it.extent)); + const perAxis = {}; + const termStrings = []; + for (let k = 0; k < ST.layout.shard.length; k++) { + const it = ST.layout.shard[k]; + perAxis[it.axis] = (perAxis[it.axis] || 0) + comps[k] * it.stride; + termStrings.push(`${comps[k]}·${it.stride}@${it.axis}`); + } + const offStrings = []; + for (const axis of Object.keys(ST.layout.offset)) { + perAxis[axis] = (perAxis[axis] || 0) + ST.layout.offset[axis]; + offStrings.push(`${ST.layout.offset[axis]}@${axis}`); + } + const owners = ST.byFlat[flat].owners; + let html = `
element ${flat} at logical (${coord.join(', ')})` + + `  →  shard split = (${comps.join(', ')})
`; + html += '
'; + html += `terms: ${termStrings.join(' + ')}` + + (offStrings.length ? ` + offset[${offStrings.join(', ')}]` : '') + '
'; + const baseParts = AXIS_ORDER.filter((a) => perAxis[a] !== undefined).map((a) => `${perAxis[a]}@${a}`); + html += `base location: ${baseParts.join(' , ')}`; + if (ST.swizzle) { + const o0 = owners[0]; + const sw = ST.swizzle; + const bytes = (sw.bits || 32) / 8; + html += `
swizzle(${sw.per_element},${sw.swizzle_len},${sw.atom_len}): ` + + `m=${o0.m || 0} → elem ${o0.__sm} → byte ${o0.__sm * bytes} → ` + + `bank ${o0.bank}, line ${o0.line} (${bytes}-byte dtype, 4-byte banks ×32)`; + } + if (owners.length > 1) { + html += `
owners (×${owners.length} via replica): ` + + owners.map((o) => '{ ' + coordStr(o, ST.gridAxes) + ' }').join(' , '); + } + html += '
'; + fb.innerHTML = html; +} + +// ── Arrow overlay (adapted from the course draw pattern) ───────────────────-- +const panels = document.getElementById('panels'); +const arrowSvg = document.getElementById('arrow'); + +function drawArrow() { + arrowSvg.innerHTML = ''; + if (hovFlat === null) return; + const leftCell = document.querySelector('#g0 .cell.hov'); + const rightCells = document.querySelectorAll('#phys .gcell.hov'); + if (!leftCell || rightCells.length === 0) return; + const pr = panels.getBoundingClientRect(); + const ns = 'http://www.w3.org/2000/svg'; + const defs = document.createElementNS(ns, 'defs'); + const marker = document.createElementNS(ns, 'marker'); + marker.setAttribute('id', 'ah'); marker.setAttribute('markerWidth', '8'); marker.setAttribute('markerHeight', '6'); + marker.setAttribute('refX', '7'); marker.setAttribute('refY', '3'); marker.setAttribute('orient', 'auto'); + const poly = document.createElementNS(ns, 'polygon'); + poly.setAttribute('points', '0 0, 8 3, 0 6'); poly.setAttribute('fill', '#222'); + marker.appendChild(poly); defs.appendChild(marker); arrowSvg.appendChild(defs); + const a = leftCell.getBoundingClientRect(); + const x1 = a.left + a.width / 2 - pr.left, y1 = a.top + a.height / 2 - pr.top; + rightCells.forEach((rc) => { + const b = rc.getBoundingClientRect(); + const x2 = b.left + b.width / 2 - pr.left, y2 = b.top + b.height / 2 - pr.top; + const mx = (x1 + x2) / 2, my = Math.min(y1, y2) - 24; + const path = document.createElementNS(ns, 'path'); + path.setAttribute('d', `M${x1},${y1} Q${mx},${my} ${x2},${y2}`); + path.setAttribute('marker-end', 'url(#ah)'); + arrowSvg.appendChild(path); + }); +} + +function clearHov() { + document.querySelectorAll('.cell.hov, .gcell.hov').forEach((d) => d.classList.remove('hov')); + arrowSvg.innerHTML = ''; + hovFlat = null; +} + +panels.addEventListener('click', (e) => { + const cell = e.target.closest('.cell') || e.target.closest('.gcell'); + if (!cell || cell.dataset.flat === undefined) return; + const flat = +cell.dataset.flat; + if (hovFlat === flat) { clearHov(); draw(); return; } + clearHov(); + hovFlat = flat; + draw(); +}); + +// ── Legend ────────────────────────────────────────────────────────────────── +function swatchEl(color) { const w = mk('div', 'swtch'); w.style.background = color; return w; } + +function drawLegend() { + const lg = document.getElementById('lg'); + lg.innerHTML = ''; + + // Color key: an actual swatch-per-value table using the same palette as the grid. + const r0 = mk('div', 'leg-row'); + const lead = mk('div', 'li'); + lead.innerHTML = `color = ${ST.colorAxis || 'physical'} value:`; + r0.appendChild(lead); + const vals = ST.colorVals || []; + if (vals.length <= 8) { + for (const v of vals) { + const li = mk('div', 'li'); + li.appendChild(swatchEl(paletteColor(v))); + li.appendChild(document.createTextNode(String(v))); + r0.appendChild(li); + } + } else { + for (let k = 0; k < 8; k++) { + const li = mk('div', 'li'); + li.appendChild(swatchEl(PALETTE[k])); + li.appendChild(document.createTextNode('≡' + k)); + r0.appendChild(li); + } + const note = mk('div', 'li'); + note.textContent = `(${ST.colorAxis} mod 8; ${vals.length} values)`; + r0.appendChild(note); + } + lg.appendChild(r0); + + const r1 = mk('div', 'leg-row'); + const c2 = mk('div', 'li'); c2.textContent = 'number = logical element index'; r1.appendChild(c2); + if (ST.layout.replica.length) { + r1.appendChild(mk('div', 'leg-sep')); + const c3 = mk('div', 'li'); + c3.textContent = 'replica → identical copies: the same color appears in multiple physical cells'; + r1.appendChild(c3); + } + lg.appendChild(r1); + + const r2 = mk('div', 'leg-row'); + const m = mk('div', 'li'); + const owners = axesUsed(ST.layout).filter(isOwnerAxis); + const mem = axesUsed(ST.layout).filter((a) => !isOwnerAxis(a)); + m.textContent = 'owner axes: ' + (owners.join(', ') || '(none — pure memory layout)') + + ' · memory axes: ' + (mem.join(', ') || '(none)'); + r2.appendChild(m); + lg.appendChild(r2); +} + +function escapeHtml(s) { + return s.replace(/[&<>"]/g, (c) => ({ '&': '&', '<': '<', '>': '>', '"': '"' }[c])); +} + +// ── Presets + controls ─────────────────────────────────────────────────────── +// Case-study presets use scaled-down shapes so every element renders; the +// mapping semantics match the full-size examples in the docs. +const PRESETS = [ + { label: 'Shard → lanes (intro)', shape: '4, 8', expr: 'S[(4,8):(8@laneid,1@laneid)]' }, + { label: 'Shard + registers', shape: '4, 8', expr: 'S[(4,2,4):(8@laneid,1@m,1@laneid)]' }, + { label: 'Shard + replica', shape: '4, 8', expr: 'S[(4,8):(8@laneid,1@laneid)] + R[2:1@warpid]' }, + { + label: 'Tensor-core tile (doc example)', shape: '8, 16', + expr: 'S[(8,2,4,2):(4@laneid,1@warpid,1@laneid,1)] + R[2:4@warpid] + 5@warpid', + }, + { label: 'Distributed 2×2 GPU mesh (pid)', shape: '4, 4', expr: 'S[(2,2,2,2):(1@pid,2@m,2@pid,1@m)]' }, + { label: 'Mesh + replica (pid)', shape: '4, 4', expr: 'S[(2,2,4):(1@pid,2@m,1@m)] + R[2:2@pid]' }, + { label: 'Accelerator scratchpad (P/F)', shape: '4, 8', expr: 'S[(2,4,4):(4@F,1@P,1@F)]' }, + { label: 'Blackwell tensor memory (TLane/TCol)', shape: '4, 8', expr: 'S[(2,4,4):(4@TCol,1@TLane,1@TCol)]' }, + { label: 'SMEM, no swizzle (bank conflicts)', shape: '8, 64', expr: 'S[(8,64):(64@m,1@m)]', dtype: 16, mode: 'none' }, + { label: 'SMEM swizzle 128B (fp16)', shape: '8, 64', expr: 'S[(8,64):(64@m,1@m)]', dtype: 16, mode: '128' }, + { label: '1-D shard', shape: '8', expr: 'S[8:4@laneid]' }, + { label: 'Extents only (default strides)', shape: '8, 4', expr: 'S[(8,4)]' }, +]; + +const DTYPES = [ + { label: 'float16 (16-bit)', bits: 16 }, + { label: 'bfloat16 (16-bit)', bits: 16 }, + { label: 'float8 (8-bit)', bits: 8 }, + { label: 'float32 (32-bit)', bits: 32 }, + { label: 'tfloat32 (32-bit)', bits: 32 }, + { label: 'float64 (64-bit)', bits: 64 }, +]; +const SWMODES = [ + { label: 'off (general layout)', value: 'off' }, + { label: 'none — raw banks', value: 'none' }, + { label: '32B swizzle', value: '32' }, + { label: '64B swizzle', value: '64' }, + { label: '128B swizzle', value: '128' }, +]; + +const shapeInput = document.getElementById('shape'); +const exprInput = document.getElementById('expr'); +const presetSel = document.getElementById('preset'); +const dtypeSel = document.getElementById('dtype'); +const swmodeSel = document.getElementById('swmode'); + +function applyPreset(i) { + const p = PRESETS[i]; + shapeInput.value = p.shape; + exprInput.value = p.expr; + if (dtypeSel) dtypeSel.value = String(p.dtype || 16); + if (swmodeSel) swmodeSel.value = p.mode || 'off'; + refresh(); +} +function refresh() { clearHov(); recompute(); draw(); } + +function init() { + PRESETS.forEach((p, i) => { + const o = document.createElement('option'); + o.value = i; o.textContent = p.label; presetSel.appendChild(o); + }); + DTYPES.forEach((d) => { + const o = document.createElement('option'); + o.value = d.bits; o.textContent = d.label; dtypeSel.appendChild(o); + }); + SWMODES.forEach((s) => { + const o = document.createElement('option'); + o.value = s.value; o.textContent = s.label; swmodeSel.appendChild(o); + }); + presetSel.addEventListener('change', () => applyPreset(+presetSel.value)); + shapeInput.addEventListener('input', refresh); + exprInput.addEventListener('input', refresh); + dtypeSel.addEventListener('change', refresh); + swmodeSel.addEventListener('change', refresh); + window.addEventListener('resize', () => { resetFit(); fitEmbed(); postHeight(); }); + window.addEventListener('load', () => { resetFit(); fitEmbed(); postHeight(); }); + // Deep-linking for embeds: ?preset=¬itle + const params = new URLSearchParams(location.search); + let presetIdx = 0; + const want = params.get('preset'); + if (want !== null) { + const slug = (x) => x.toLowerCase().replace(/[^a-z0-9]+/g, '-').replace(/^-|-$/g, ''); + if (/^\d+$/.test(want)) { + presetIdx = Math.min(PRESETS.length - 1, Math.max(0, parseInt(want, 10))); + } else { + const w = slug(want); + const found = PRESETS.findIndex((p) => slug(p.label).includes(w)); + if (found >= 0) presetIdx = found; + } + } + presetSel.value = presetIdx; + if (params.has('notitle')) document.body.classList.add('notitle'); + if (params.has('lock')) document.body.classList.add('lock'); + applyPreset(presetIdx); +} + +init(); diff --git a/_extra_zh/zh/_static/tirx-layout-demo/viz-base.css b/_extra_zh/zh/_static/tirx-layout-demo/viz-base.css new file mode 100644 index 00000000..624194ce --- /dev/null +++ b/_extra_zh/zh/_static/tirx-layout-demo/viz-base.css @@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * Derived from mlsyscourse/slides-modern-gpu-programming + * (data-layout/site/viz-base.css). Shared base styles for the + * TIRx layout visualization; not derived from any third-party demo. + */ + +/* Shared base styles for all viz HTMLs + * Colors: --bg gray, --accent blue, --surface white + * Fonts: Inter (body), SF Mono (code/notation) + */ + +:root { + --bg:#fff; --surface:#fff; --border:#dfe1e6; --text:#222; --dim:#888; --accent:#3b82f6; + + /* ── Group palette (tiles, lanes, banks, sectors) ─────── */ + --color-group-0: #5b9bd5; + --color-group-1: #ed9a3c; + --color-group-2: #d95555; + --color-group-3: #45b5a5; + --color-group-4: #5fb85f; + --color-group-5: #e0b828; + --color-group-6: #9d6eb5; + --color-group-7: #e87888; + + /* ── Interaction ──────────────────────────────────────── */ + --color-hover-bg: #dbeafe; + --color-hover-text: #1e40af; + + /* ── Status ───────────────────────────────────────────── */ + --color-good: #2e7d32; + --color-bad: #c62828; + + /* ── Boundaries & neutrals ────────────────────────────── */ + --color-boundary: #8C1515; + --color-cell-bg: #f0f1f3; + --color-cell-bg-alt: #e8eaed; +} +* { box-sizing:border-box; margin:0; padding:0; } +body { background:var(--bg); color:var(--text); font-family:'Inter','SF Pro','Segoe UI',system-ui,sans-serif; padding:24px; } +body.figure h1, body.figure .sub { display:none; } +body.notitle h1, body.notitle .sub { display:none; } +h1 { text-align:center; font-size:15px; font-weight:700; margin-bottom:2px; } +.sub { text-align:center; color:var(--dim); font-size:13px; margin-bottom:20px; font-family:'SF Mono','Fira Code',monospace; } + +/* Controls */ +.controls { display:flex; gap:18px; justify-content:center; align-items:center; margin-bottom:10px; flex-wrap:wrap; } +.lbl { font-size:13px; color:var(--dim); margin-right:4px; font-weight:600; } +.bg { display:inline-flex; gap:2px; } +.btn { + padding:5px 11px; border:1px solid var(--border); background:var(--surface); color:var(--text); + cursor:pointer; border-radius:5px; font-size:10px; font-family:inherit; font-weight:500; transition:all .12s; +} +.btn:hover { background:#eef0f4; } +.btn.on { background:var(--accent); border-color:var(--accent); color:#fff; } +.btn.on:hover { background:#2563eb; } + +/* Side-by-side panels */ +.panels { display:grid; grid-template-columns:1fr 1fr; gap:20px; max-width:1100px; margin:0 auto; position:relative; } +.panel { background:var(--surface); border-radius:10px; padding:16px 14px; border:1px solid var(--border); + box-shadow:0 1px 3px rgba(0,0,0,.06); } +.panel h2 { text-align:center; font-size:10px; font-weight:700; margin-bottom:2px; } +.panel .nota { text-align:center; font-size:8px; color:var(--accent); font-family:'SF Mono','Fira Code',monospace; + margin-bottom:8px; font-weight:600; } + +/* Grid cells */ +.grid { display:grid; gap:3px; } +.hdr { font-size:7px; color:var(--dim); text-align:center; padding:2px 0; font-weight:600; } +.rl { font-size:7px; color:var(--dim); display:flex; align-items:center; justify-content:flex-end; padding-right:3px; font-weight:600; } +.cell { + aspect-ratio:1; border-radius:5px; display:flex; flex-direction:column; align-items:center; justify-content:center; + font-size:10px; font-weight:700; border:2.5px solid transparent; transition:all .18s; + min-width:0; cursor:pointer; line-height:1.15; +} +.cell.hov { border-color:#222; border-width:3px; z-index:2; box-shadow:0 0 8px rgba(59,130,246,.4); } +.cell.dm { opacity:.25; } + +/* Arrow SVG overlay */ +.arrow-svg { position:absolute; top:0; left:0; width:100%; height:100%; pointer-events:none; z-index:10; overflow:visible; } +.arrow-svg path { fill:none; stroke:#222; stroke-width:1.5; } +.arrow-svg text { font-size:8px; font-weight:600; fill:#222; font-family:inherit; } + +/* Formula bar */ +.formula-bar { max-width:1100px; margin:18px auto 0; background:var(--surface); border-radius:10px; + padding:12px 18px; border:1px solid var(--border); box-shadow:0 1px 3px rgba(0,0,0,.06); } +.formula-bar .ftitle { font-size:15px; font-weight:700; margin-bottom:6px; } +.formula-bar .fcontent { font-size:15px; font-family:'SF Mono','Fira Code',monospace; color:var(--accent); line-height:1.6; } + +/* Legend */ +.leg { display:flex; flex-direction:column; align-items:center; gap:4px; margin-top:14px; } +.leg-row { display:flex; gap:12px; justify-content:center; flex-wrap:wrap; } +.leg-group { display:flex; gap:10px; align-items:center; } +.leg-sep { width:1px; height:16px; background:var(--border); } +.li { display:flex; align-items:center; gap:4px; font-size:9px; color:var(--dim); } +.swtch { width:14px; height:14px; border-radius:3px; } + +@media(max-width:700px) { .panels { grid-template-columns:1fr; } } diff --git a/_extra_zh/zh/_static/tirx-layout-demo/viz-base.js b/_extra_zh/zh/_static/tirx-layout-demo/viz-base.js new file mode 100644 index 00000000..c5ef10f1 --- /dev/null +++ b/_extra_zh/zh/_static/tirx-layout-demo/viz-base.js @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + * + * Derived from mlsyscourse/slides-modern-gpu-programming + * (data-layout/site/viz-base.js). Shared behavior for the TIRx + * layout visualization; not derived from any third-party demo. + */ + +// Shared behavior for all viz HTMLs +document.addEventListener('DOMContentLoaded', function() { + var p = new URLSearchParams(location.search); + if (p.has('notitle')) document.body.classList.add('notitle'); + + // Forward arrow keys to parent (reveal.js) when embedded + if (window.parent !== window) { + document.addEventListener('keydown', function(e) { + if ([37, 38, 39, 40, 27, 32].indexOf(e.keyCode) !== -1) { + // Left, Up, Right, Down, Escape, Space + window.parent.postMessage({ type: 'revealKey', keyCode: e.keyCode }, '*'); + } + }); + } +}); + +// Auto-height: when embedded in the book, measure our own content height and post +// it to the parent so it can size the iframe to fit (no inner scrollbar). This +// demo is responsive (it fills the iframe width), so only the HEIGHT needs to +// follow content. Push-based, so it catches our own DOM changes (editing the +// layout, clicking a cell, switching presets) that an outside observer can miss. +(function () { + if (window.parent === window) return; + var lastH = 0; + function report() { + var b = document.body, de = document.documentElement; + var h = (b ? b.scrollHeight : 0) || (de ? de.scrollHeight : 0) || 0; + if (h && Math.abs(h - lastH) > 1) { + lastH = h; + window.parent.postMessage({ type: 'demoHeight', height: h }, '*'); + } + } + var scheduled = false; + function schedule() { + if (scheduled) return; + scheduled = true; + requestAnimationFrame(function () { scheduled = false; report(); }); + } + try { new ResizeObserver(schedule).observe(document.documentElement); } catch (e) {} + try { + new MutationObserver(schedule).observe(document.documentElement, { + subtree: true, childList: true, attributes: true, characterData: true + }); + } catch (e) {} + document.addEventListener('DOMContentLoaded', schedule); + window.addEventListener('load', schedule); + window.addEventListener('click', function () { setTimeout(schedule, 0); }, true); + [100, 300, 600, 1200].forEach(function (t) { setTimeout(schedule, t); }); +})(); diff --git a/_extra_zh/zh/code-highlight.css b/_extra_zh/zh/code-highlight.css new file mode 100644 index 00000000..bc953fa1 --- /dev/null +++ b/_extra_zh/zh/code-highlight.css @@ -0,0 +1,139 @@ +/* ── Syntax highlight — shared across all demo slides ── + * + * Usage: + * + * 
…
or 
…
+ * + * Token classes: + * .kw — keyword (if, def, for, while, with, return, True, False) + * .fn — callable (function/method calls — the prominent highlight) + * .str — string literal + * .num — number literal + * .cmt — comment + * .op — operator (+, -, ==, =, :) + * .typ — type name + * .dec — decorator (@) + * + * Convention: only wrap the callable name in .fn, NOT the namespace prefix. + * ✓ Tx.copy_async(...) + * ✗ Tx.copy_async(...) + */ + +/* ── Dark theme (dark bg, bright tokens) ────────────── */ +.code-dark { + background: #0f172a; color: #e2e8f0; + font-family: 'SF Mono','Fira Code',monospace; + font-size: 12px; line-height: 1.6; + border-radius: 8px; padding: 12px 16px; + overflow-x: auto; +} +.code-dark .kw { color: #c084fc; } /* purple — keywords */ +.code-dark .fn { color: #60a5fa; } /* blue — callables (prominent) */ +.code-dark .str { color: #34d399; } /* green — strings */ +.code-dark .num { color: #fbbf24; } /* yellow — numbers */ +.code-dark .cmt { color: #64748b; font-style: italic; } /* gray — comments */ +.code-dark .op { color: #94a3b8; } /* slate — operators */ +.code-dark .typ { color: #7dd3fc; } /* cyan — types */ +.code-dark .dec { color: #c084fc; } /* purple — decorators */ + +/* ── Light theme (white bg, muted tokens) ───────────── */ +.code-light { + background: #f8fafc; color: #1e293b; + font-family: 'SF Mono','Fira Code',monospace; + font-size: 12px; line-height: 1.6; + border-radius: 8px; padding: 12px 16px; + border: 1px solid #e2e8f0; + overflow-x: auto; +} +.code-light .kw { color: #7c3aed; } /* purple — keywords */ +.code-light .fn { color: #2563eb; } /* blue — callables (prominent) */ +.code-light .str { color: #059669; } /* green — strings */ +.code-light .num { color: #d97706; } /* amber — numbers */ +.code-light .cmt { color: #6b7280; font-style: italic; } /* gray — comments */ +.code-light .op { color: #64748b; } /* slate — operators */ +.code-light .typ { color: #0369a1; } /* cyan — types */ +.code-light .dec { color: #7c3aed; } /* purple — decorators */ + +/* ── Code block panel (used by createCodeBlock) ────── + * + * Structure: + *
— wrapper with rounded corners + *
Title
— optional title bar + * 
         — code container (scrollable)
+ *       (regions and lines rendered by createCodeBlock)
+ *
+ *
+ * + * Or without panel wrapper, just 
 works too.
+ */
+
+.cb-panel {
+  border-radius: 8px;
+  overflow: hidden;
+}
+.cb-panel > .code-dark,
+.cb-panel > .code-light {
+  border-radius: 0;
+}
+
+/* Header bar */
+.cb-panel .cb-header {
+  padding: 6px 14px;
+  font-size: 11px;
+  font-weight: 600;
+  font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
+}
+.cb-panel.dark .cb-header {
+  background: #1e293b;
+  color: #94a3b8;
+  border-bottom: 1px solid rgba(255,255,255,0.06);
+}
+.cb-panel.light .cb-header {
+  background: #f1f5f9;
+  color: #64748b;
+  border-bottom: 1px solid #e2e8f0;
+}
+
+/* Dark scrollbar */
+.code-dark {
+  scrollbar-width: thin;
+  scrollbar-color: rgba(255,255,255,0.15) transparent;
+}
+.code-dark::-webkit-scrollbar { width: 6px; height: 6px; }
+.code-dark::-webkit-scrollbar-track { background: transparent; }
+.code-dark::-webkit-scrollbar-thumb { background: rgba(255,255,255,0.15); border-radius: 3px; }
+.code-dark::-webkit-scrollbar-thumb:hover { background: rgba(255,255,255,0.25); }
+
+/* Light scrollbar */
+.code-light {
+  scrollbar-width: thin;
+  scrollbar-color: rgba(0,0,0,0.15) transparent;
+}
+.code-light::-webkit-scrollbar { width: 6px; height: 6px; }
+.code-light::-webkit-scrollbar-track { background: transparent; }
+.code-light::-webkit-scrollbar-thumb { background: rgba(0,0,0,0.12); border-radius: 3px; }
+.code-light::-webkit-scrollbar-thumb:hover { background: rgba(0,0,0,0.2); }
+
+/* Regions */
+.cb-region {
+  border-left: 3px solid transparent;
+  padding: 0.35rem 0 0.35rem 0.85rem;
+  margin: 0.15rem 0;
+  border-radius: 2px;
+  cursor: pointer;
+  transition: background 0.2s;
+}
+.code-dark .cb-region:hover  { background: rgba(255,255,255,0.05); }
+.code-dark .cb-region.active { background: rgba(255,255,255,0.07); }
+.code-light .cb-region:hover  { background: rgba(0,0,0,0.04); }
+.code-light .cb-region.active { background: rgba(0,0,0,0.06); }
+
+/* Lines */
+.cb-line {
+  display: block;
+  white-space: pre;
+}
+.cb-line.cb-focus {
+  margin-left: -0.85rem;
+  padding-left: 0.85rem;
+}
diff --git a/_extra_zh/zh/code-highlight.js b/_extra_zh/zh/code-highlight.js
new file mode 100644
index 00000000..1e36669c
--- /dev/null
+++ b/_extra_zh/zh/code-highlight.js
@@ -0,0 +1,224 @@
+/* ── Code highlight — shared TIRx/Python tokenizer + code block builder ──
+ *
+ * Usage:
+ *   
+ *   
+ *
+ * API:
+ *   highlightCode(text)  — tokenize text, return HTML with  tokens
+ *   escapeHtml(text)     — escape < > & for safe HTML insertion
+ *
+ *   createCodeBlock(container, code, options)
+ *     Renders highlighted code into container with optional regions + focus lines.
+ *     options:
+ *       blockDefs:    [{ key, start, end, color }]              — clickable regions
+ *       focusLines:   { key: [lineNos] } or { key: { lines, color } } — per-region line highlights
+ *       onBlockClick: function(key)                             — callback on region click
+ *
+ * Auto-init:
+ *   Elements with class "auto-hl" get their textContent highlighted on load.
+ *   Combine with code-dark / code-light for theme:
+ *     
Tx.copy_async(Asmem, A[...])

+ *
+ * Panel structure (optional wrapper for header + rounded corners):
+ *     

+ *       
Title

+ *       

+ *     

+ *
+ * Token classes (same as code-highlight.css):
+ *   .kw  — keyword      .fn  — callable     .str — string
+ *   .num — number        .cmt — comment      .op  — operator
+ *   .typ — type          .dec — decorator
+ *
+ * Namespace convention: only the callable name gets .fn, NOT the prefix.
+ *   Tx.copy_async  →  Tx.copy_async
+ *   T.ptx.tcgen05.mma  →  T.ptx.tcgen05.mma
+ */
+
+(function (root) {
+  "use strict";
+
+  function escapeHtml(text) {
+    return text
+      .replace(/&/g, "&")
+      .replace(//g, ">");
+  }
+
+  var KEYWORDS = /^(def|with|for|in|if|else|elif|and|or|not|return|True|False|None|range|class|import|from|as|pass|break|continue|while|try|except|finally|raise|yield|lambda|assert|del|global|nonlocal)$/;
+
+  function classifyToken(tok) {
+    if (tok.startsWith("#")) return "cmt";
+    if (tok.startsWith("@")) return "dec";
+    if (tok.startsWith('"') || tok.startsWith("'")) return "str";
+    if (/^\d[\d.]*$/.test(tok)) return "num";
+    if (KEYWORDS.test(tok)) return "kw";
+    if (/^(?:Tx|T)(\.[A-Za-z_]\w*)+$/.test(tok)) return "fn";
+    if (/^[A-Za-z_]\w*$/.test(tok)) return "fn";
+    return "";
+  }
+
+  var TOKEN_RE = /("(?:[^"\\]|\\.)*"|'(?:[^'\\]|\\.)*'|#.*$|@[A-Za-z_][\w.]*|\b(?:def|with|for|in|if|else|elif|and|or|not|return|True|False|None|range|class|import|from|as|pass|break|continue|while|try|except|finally|raise|yield|lambda|assert|del|global|nonlocal)\b|\b\d[\d.]*\b|(?:Tx|T)(?:\.[A-Za-z_]\w*)+|\b[A-Za-z_]\w*(?=\())/gm;
+
+  function highlightCode(text) {
+    TOKEN_RE.lastIndex = 0;
+    var out = "";
+    var last = 0;
+    var m;
+    while ((m = TOKEN_RE.exec(text)) !== null) {
+      out += escapeHtml(text.slice(last, m.index));
+      var tok = m[0];
+      var cls = classifyToken(tok);
+      if (cls === "fn" && tok.indexOf(".") !== -1) {
+        // Namespace.callable — only wrap the last component
+        var lastDot = tok.lastIndexOf(".");
+        out += escapeHtml(tok.substring(0, lastDot + 1));
+        out += '' + escapeHtml(tok.substring(lastDot + 1)) + "";
+      } else if (cls) {
+        out += '' + escapeHtml(tok) + "";
+      } else {
+        out += escapeHtml(tok);
+      }
+      last = TOKEN_RE.lastIndex;
+    }
+    out += escapeHtml(text.slice(last));
+    return out;
+  }
+
+  // Auto-init: highlight elements with class "auto-hl"
+  if (typeof document !== "undefined") {
+    document.addEventListener("DOMContentLoaded", function () {
+      var els = document.querySelectorAll(".auto-hl");
+      for (var i = 0; i < els.length; i++) {
+        els[i].innerHTML = highlightCode(els[i].textContent);
+      }
+    });
+  }
+
+  // ── createCodeBlock: render code with optional regions + focus lines ──
+
+  function hexToRgb(hex) {
+    return [
+      parseInt(hex.slice(1, 3), 16),
+      parseInt(hex.slice(3, 5), 16),
+      parseInt(hex.slice(5, 7), 16)
+    ];
+  }
+
+  function createCodeBlock(container, code, options) {
+    options = options || {};
+    var blockDefs = options.blockDefs;
+    var focusLines = options.focusLines;
+    var onBlockClick = options.onBlockClick;
+    var lines = code.split("\n");
+
+    if (blockDefs) {
+      // Render code split into clickable regions
+      container.innerHTML = blockDefs.map(function (b) {
+        var lineHtml = lines.slice(b.start - 1, b.end).map(function (line, idx) {
+          var lineNo = b.start + idx;
+          var content = line.length ? highlightCode(line) : " ";
+          return '
' + content + "
";
+        }).join("");
+        var style = b.color ? ' style="border-left-color:' + b.color + '"' : "";
+        return '
" + lineHtml + "
";
+      }).join("");
+
+      // Color map for hover/active/focus
+      var colorMap = {};
+      blockDefs.forEach(function (b) { if (b.color) colorMap[b.key] = b.color; });
+
+      var regions = container.querySelectorAll(".cb-region");
+      var allLines = container.querySelectorAll(".cb-line");
+      var activeKey = null;
+
+      function clearFocus() {
+        for (var i = 0; i < allLines.length; i++) {
+          allLines[i].classList.remove("cb-focus");
+          allLines[i].style.background = "";
+        }
+      }
+
+      function applyFocus(key) {
+        if (!focusLines || !focusLines[key]) return;
+        var fl = focusLines[key];
+        var lns = Array.isArray(fl) ? fl : fl.lines;
+        var color = Array.isArray(fl) ? null : fl.color;
+        if (!color && colorMap[key]) {
+          var rgb = hexToRgb(colorMap[key]);
+          color = "rgba(" + rgb.join(",") + ",0.18)";
+        }
+        if (!color) color = "rgba(255,255,255,0.15)";
+        for (var i = 0; i < lns.length; i++) {
+          var el = container.querySelector('.cb-line[data-ln="' + lns[i] + '"]');
+          if (el) {
+            el.classList.add("cb-focus");
+            el.style.background = color;
+          }
+        }
+      }
+
+      for (var r = 0; r < regions.length; r++) {
+        (function (el) {
+          var regionKey = el.dataset.region;
+
+          // Hover with region color
+          if (colorMap[regionKey]) {
+            var rgb = hexToRgb(colorMap[regionKey]);
+            var hoverBg = "rgba(" + rgb.join(",") + ",0.15)";
+            el.addEventListener("mouseenter", function () {
+              if (!el.classList.contains("active")) el.style.background = hoverBg;
+            });
+            el.addEventListener("mouseleave", function () {
+              if (!el.classList.contains("active")) el.style.background = "";
+            });
+          }
+
+          // Click to activate
+          el.addEventListener("click", function () {
+            if (regionKey === activeKey) return;
+            activeKey = regionKey;
+            for (var j = 0; j < regions.length; j++) {
+              regions[j].classList.remove("active");
+              regions[j].style.background = "";
+            }
+            el.classList.add("active");
+            if (colorMap[regionKey]) {
+              var rgb2 = hexToRgb(colorMap[regionKey]);
+              el.style.background = "rgba(" + rgb2.join(",") + ",0.15)";
+            }
+            clearFocus();
+            applyFocus(regionKey);
+            if (onBlockClick) onBlockClick(regionKey);
+          });
+        })(regions[r]);
+      }
+    } else {
+      // Simple highlighted code, no regions
+      container.innerHTML = lines.map(function (line, idx) {
+        var content = line.length ? highlightCode(line) : " ";
+        return '
' + content + "
";
+      }).join("");
+    }
+  }
+
+  // ── renderAnnotatedCode: highlight with per-line CSS classes ──
+
+  function renderAnnotatedCode(container, code, lineClasses) {
+    lineClasses = lineClasses || {};
+    var lines = code.split("\n");
+    container.innerHTML = lines.map(function (line, idx) {
+      var content = line.length ? highlightCode(line) : " ";
+      var cls = lineClasses[idx + 1];
+      if (cls) return '
' + content + "
";
+      return "
" + content + "
";
+    }).join("");
+  }
+
+  // Export
+  root.highlightCode = highlightCode;
+  root.escapeHtml = escapeHtml;
+  root.createCodeBlock = createCodeBlock;
+  root.renderAnnotatedCode = renderAnnotatedCode;
+})(typeof window !== "undefined" ? window : this);
diff --git a/_extra_zh/zh/demo/barrier_intro.html b/_extra_zh/zh/demo/barrier_intro.html
new file mode 100644
index 00000000..9973f691
--- /dev/null
+++ b/_extra_zh/zh/demo/barrier_intro.html
@@ -0,0 +1,387 @@
+
+
+
+
+异步协作：Barrier
+
+
+
+
+
+
+
异步协作：Barrier

+
硬件单元完成后 signal mbarrier；warp group 在继续前等待

+
+

+

+  
+  
GMEM
A, B

+  

+    
+    
TMA 加载

+  

+  
SMEM

+  

+    
+    
tcgen05.mma

+  

+  
TMEM
累加器

+  

+    
+    
tcgen05.ld

+  

+  
寄存器
cast

+  

+    
+    
存储

+  

+  
SMEM

+  

+    
+    
TMA 存储

+  

+  
GMEM
D

+
+  
+  

+
+  
+  

+    
Warp 0 — 生产者

+    
TMA 派发

+  

+  

+    
Warp 1 — 消费者

+    
tcgen05 派发

+  

+  

+    
Warp 0 — 下一阶段

+    
TMA 派发

+  

+
+  
+  

+
+  
+  
TMA
引擎

+  

+    
+    
complete-tx

+  

+  
mbar
TMA 就绪

+  
Tensor
Core

+  
mbar
SMEM 空闲

+

+
+
+
+  
+    
+      
+    
+    
+      
+    
+    
+      
+    
+  
+  
+  
+  
+  
+  
+  
+  
+  
+  
+  
+  
+
+

+
+
+

+  协作：
+

+

+
+
    
+  
  • mbarrier (mbar) — 用于协调线程和硬件单元之间的异步操作
    • 
+  
  • 不同 warp 可并行运行在操作的不同阶段
    • 
+

+
+
+
+
diff --git a/_extra_zh/zh/demo/cta_cluster.html b/_extra_zh/zh/demo/cta_cluster.html
new file mode 100644
index 00000000..47e16dbb
--- /dev/null
+++ b/_extra_zh/zh/demo/cta_cluster.html
@@ -0,0 +1,113 @@
+
+
+
+
+
+2-CTA Cluster
+
+
+
+
+
+
2-CTA Cluster：通过跨 CTA SMEM 读取实现协作 MMA

+
点击任意部分查看说明；两个 CTA 在 cluster 内共享 stored-B 行切片（DSMEM）

+
+

+  CLUSTER — 跨 SM 的 2 个 CTA · DSMEM
+  

+    

+      
CTA 0 · SM-0

+      
Asmem（本 CTA）A 行 0–127

+      
BsmemB stored 行 0–127

+      
D[0:128, 0:256]

+    

+    

+      
跨 CTA
读取 ↔

+    

+    

+      
CTA 1 · SM-1

+      
Asmem（本 CTA）A 行 128–255

+      
BsmemB stored 行 128–255

+      
D[128:256, 0:256]

+    

+  

+  
Cluster 输出：256 × 256 — 是单个 CTA 独立计算 128 × 128 的两倍

+

+
+

+  

+  

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/mbarrier_arrive_timeline.html b/_extra_zh/zh/demo/mbarrier_arrive_timeline.html
new file mode 100644
index 00000000..e2856809
--- /dev/null
+++ b/_extra_zh/zh/demo/mbarrier_arrive_timeline.html
@@ -0,0 +1,470 @@
+
+
+
+
+
+MBarrier arrive 时间线
+
+
+
+
+
+  
用于跨线程同步的 MBarrier

+
+  

+    
+    

+      

+
+        
+        
T0（生产者）

+        

+
+        
+        
MBar

+        

+
+        
+        
T1（消费者）

+        

+
+        
+        

+          时间线 
+        

+        

+      

+
+      
+      
+    

+    

+  

+
+  
+  

+    

+      
 工作 / 活跃
+    

+    

+      
 Signal 操作
+    

+    

+      
 阻塞 / 等待
+    

+    

+      
 Barrier 状态
+    

+    

+      
 空闲
+    

+    

+       Signal 并改变 mbar 状态
+    

+    

+      
 init 后同步线程
+    

+  

+
+  

+    

+      
设置（T0）

+      
mbarrier.init(mbar, 1) — 对所有运行只设置一次，设置后显式同步。

+    

+    

+      
生产者（T0）

+      
mbarrier.arrive(mbar) 递减 pending_count；当它变为 0 时，phase 完成。

+    

+    

+      
消费者（T1）

+      
mbarrier.try_wait(mbar, phase) 挂起 T1，直到 phase 完成。

+    

+  

+
+
+
+
diff --git a/_extra_zh/zh/demo/mbarrier_mechanism.html b/_extra_zh/zh/demo/mbarrier_mechanism.html
new file mode 100644
index 00000000..16dbb513
--- /dev/null
+++ b/_extra_zh/zh/demo/mbarrier_mechanism.html
@@ -0,0 +1,220 @@
+
+
+
+
+
+MBarrier 机制
+
+
+
+
+
+  
MBarrier：数据结构与 API

+
+  

+    

+      
MBarrier 对象（64-bit，位于 shared memory）

+      

+        
phase
0 或 1

+        
pending_count
剩余 arrival

+        
expected_count
预期总数

+        
tx-count
待完成字节

+      

+      

+        phase
+        当前 phase（0 或 1）。完成后自动翻转。
+      

+      

+        pending_count
+        phase 完成前仍需要的 arrival 数量。
+      

+      

+        expected_count
+        每个 phase 预期的 arrival 总数（由 init 设置）。
+      

+      

+        tx-count
+        尚未完成的异步字节数。由 arrive.expect_tx 增加，由硬件在传输完成时减少。
+      

+    

+
+    

+      
Phase 完成条件

+      

+        pending_count == 0

+        &&

+        tx-count ≤ 0
+      

+      

+        当 所有预期线程都已 arrive 且 所有预期异步字节都已传输完成 时，一个 phase 完成。
+        完成后，barrier 会自动重置：phase 翻转，pending_count 重置为 expected_count，tx-count 重置为 0。
+      

+    

+  

+
+  

+    
核心 API

+    
+      
+        
+      
+      
+        
+          
+          
+          
+        
+        
+          
+          
+          
+        
+        
+          
+          
+          
+        
+        
+          
+          
+          
+        
+        
+          
+          
+          
+        
+      
+    
指令角色对 MBarrier 的影响
mbarrier.init 设置用预期 arrival 数初始化 barrierphase=0, pending=count, tx=0
mbarrier.arrive 生产者线程向 barrier signal arrivalpending_count -= 1
mbarrier.arrive.expect_tx 生产者Arrive 并声明预期异步传输字节数pending -= 1, tx_count += txCount
tcgen05.commit 生产者Tensor Core signal MMA 完成在 mbarrier 上执行 arrive::one
mbarrier.try_wait 消费者等待（挂起线程）直到 phase 完成阻塞直到 phase 完成

+  

+
+  

+    三种 arrival 模式： (1) mbarrier.arrive — 线程直接递减 pending_count。
+    (2) arrive.expect_tx — 线程 arrive 并声明字节数；TMA 引擎在传输完成时自动递减 tx_count。
+    (3) tcgen05.commit — Tensor Core 在 MMA 完成时自动 arrive。
+  

+
+
+
diff --git a/_extra_zh/zh/demo/mbarrier_tcgen05_timeline.html b/_extra_zh/zh/demo/mbarrier_tcgen05_timeline.html
new file mode 100644
index 00000000..6889526f
--- /dev/null
+++ b/_extra_zh/zh/demo/mbarrier_tcgen05_timeline.html
@@ -0,0 +1,479 @@
+
+
+
+
+
+MBarrier tcgen05.commit 时间线
+
+
+
+
+
+  
用 MBarrier signal tcgen05 完成

+
+  

+    
+    

+      

+
+        
+        
T0（生产者）

+        

+
+        
+        
Tensor Core

+        

+
+        
+        
MBar

+        

+
+        
+        
T1（消费者）

+        

+
+        
+        

+          时间线 
+        

+        

+      

+
+      
+      
+    

+    

+  

+
+  
+  

+    

+      
 工作 / 活跃
+    

+    

+      
 Signal 操作
+    

+    

+      
 派发到 Tensor Core
+    

+    

+      
 阻塞 / 等待
+    

+    

+      
 Barrier 状态
+    

+    

+      
 空闲
+    

+    

+       Signal 并改变 mbar 状态
+    

+    

+      
 init 后同步线程
+    

+  

+
+  

+    

+      
生产者（T0）

+      
tcgen05.mma 发起异步矩阵乘。tcgen05.commit(mbar) 告诉 barrier 跟踪它；硬件完成后会执行 arrive::one。

+    

+    

+      
Tensor Core

+      
执行 T0 排队的 tcgen05.mma 和 tcgen05.commit。

+    

+    

+      
消费者（T1）

+      
mbarrier.try_wait(mbar, phase) 会挂起 T1，直到 phase 完成。结果随后可在 TMEM 中读取。

+    

+  

+
+
+
+
diff --git a/_extra_zh/zh/demo/mbarrier_tma_timeline.html b/_extra_zh/zh/demo/mbarrier_tma_timeline.html
new file mode 100644
index 00000000..6ae2808a
--- /dev/null
+++ b/_extra_zh/zh/demo/mbarrier_tma_timeline.html
@@ -0,0 +1,490 @@
+
+
+
+
+
+MBarrier TMA 时间线
+
+
+
+
+
+  
用 MBarrier signal TMA 完成

+
+  

+    
+    

+      

+
+        
+        
T0（生产者）

+        

+
+        
+        
TMA 引擎

+        

+
+        
+        
MBar

+        

+
+        
+        
T1（消费者）

+        

+
+        
+        

+          时间线 
+        

+        

+      

+
+      
+      
+    

+    

+  

+
+  
+  

+    

+      
 工作 / 活跃
+    

+    

+      
 Signal 操作
+    

+    

+      
 派发到 TMA
+    

+    

+      
 阻塞 / 等待
+    

+    

+      
 Barrier 状态
+    

+    

+      
 空闲
+    

+    

+       Signal 并改变 mbar 状态
+    

+    

+      
 init 后同步线程
+    

+  

+
+  

+    

+      
生产者（T0）

+      
arrive.expect_tx(4096) 将 pending 递减到 0，并设置 tx_count。cp.async.bulk 触发 TMA 传输。

+    

+    

+      
TMA 引擎

+      
执行从 T0 派发的 cp.async.bulk。完成时自动递减 tx_count。

+    

+    

+      
消费者（T1）

+      
try_wait 会挂起，直到 pending==0 且 tx==0。硬件在传输完成时自动递减 tx_count。

+    

+  

+
+
+
+
diff --git a/_extra_zh/zh/demo/phase_tracking.html b/_extra_zh/zh/demo/phase_tracking.html
new file mode 100644
index 00000000..94c82c34
--- /dev/null
+++ b/_extra_zh/zh/demo/phase_tracking.html
@@ -0,0 +1,201 @@
+
+
+
+
+
+Phase 跟踪
+
+
+
+
+
+  
Phase 跟踪

+  
Barrier 通过在 phase 0 ↔ 1 间翻转，在多次迭代之间复用

+
+  

+    
+    
+  

+
+  

+    
phase = 0
迭代 0

+    
phase = 1
迭代 1

+    
phase = 0
迭代 2

+    
phase = 1
迭代 3

+    
phase = 0
迭代 4

+    
phase = 1
迭代 5

+    
phase = 0
迭代 6

+    
phase = 1
迭代 7

+  

+
+  

+    

+      
工作方式

+      

+        每个 barrier 都有一个 phase bit（0 或 1）。所有预期 arrival 到达后，barrier 会翻转 phase 并重置，准备好服务下一次迭代；无需重新初始化。
+        

+          
1
 生产者 arrive 到 barrier（phase 0）

+          
2
 消费者调用 try_wait(phase=0) — 解除阻塞

+          
3
 Barrier 翻转到 phase 1，phase ^= 1

+        

+      

+    

+
+    

+      
为什么要跟踪 phase？

+      

+        在流水线 kernel 中，barrier 会在 K-loop 的每次迭代中复用。phase 跟踪解决的问题是：“这次 barrier 触发对应的是迭代 i，还是迭代 i-1？”
+        

+          phase_tma = 0

+          for k in range(K):

+            try_wait(tma_bar, phase_tma)

+            phase_tma ^= 1  # 翻转
+        

+      

+    

+  

+
+
+
+
+
diff --git a/_extra_zh/zh/demo/pipeline_arch.html b/_extra_zh/zh/demo/pipeline_arch.html
new file mode 100644
index 00000000..469773bc
--- /dev/null
+++ b/_extra_zh/zh/demo/pipeline_arch.html
@@ -0,0 +1,509 @@
+
+
+
+
+
+Blackwell 上的 GEMM 数据流水线
+
+
+
+
+
+  
Blackwell 上的 GEMM 数据流水线

+  
D = A × B — 点击一个操作，在架构图中查看对应的数据路径

+
+  
+  

+    

+      
GMEM
A, B

+    

+    

+      
+      
TMA 加载

+    

+    

+      
SMEM

+    

+    

+      
+      
tcgen05.mma

+    

+    

+      
TMEM
累加器

+    

+    

+      
+      
tcgen05.ld

+    

+    

+      
寄存器
cast

+    

+    

+      
+      
存储

+    

+    

+      
SMEM

+    

+    

+      
+      
TMA 存储

+    

+    

+      
GMEM
D

+    

+  

+
+  

+    

+    Blackwell 架构
+    

+  

+
+  
+  

+    

+    

+      
流式多处理器 (SM)

+
+      

+        

+          
Tensor Core (tcgen05)

+          
第 5 代 MMA

+        

+        

+        

+          
CUDA Core

+          
FP/INT 单元

+        

+      

+
+      

+        

+          
Shared Memory (SMEM)

+          
每个 SM 228 KB

+        

+        

+          
Tensor Memory (TMEM)

+          
128 条 lane

+        

+        

+          
寄存器文件

+        

+      

+
+      

+        

+          
TMA 引擎

+          
数据搬运器

+        

+        

+        

+      

+    

+    

+      
SM ...

+      × N
+    

+    

+
+    

+      
全局内存 (GMEM)

+    

+
+    
+    
+  

+
+  
+  

+    

+    
点击流水线中的一个操作以高亮对应的数据路径。

+  

+
+
+
+
diff --git a/_extra_zh/zh/demo/sf_tmem.html b/_extra_zh/zh/demo/sf_tmem.html
new file mode 100644
index 00000000..48093a83
--- /dev/null
+++ b/_extra_zh/zh/demo/sf_tmem.html
@@ -0,0 +1,148 @@
+
+
+
+
+TMEM 中的 Scale Factor
+
+
+
+
+
+
TMEM 中的 Scale Factor（nvfp4）

+
M = 128，一个 MMA-K block（SF_K = 4） ·  warpx4 broadcast R[4 : 32@TLane]

+
+

+  
+  

+    
逻辑 SFA  (M × SF_K)

+    
点击 cell = SFA[m, sfk]  ·  颜色 = m // 32 分组

+    

+  

+  

+    
TMEM  （128 条 lane × 16 字节）

+    
w0: m 0–31 · w1: m 32–63 · w2: m 64–95 · w3: m 96–127  ·  byte = sfk

+    

+  

+

+
+

+  
布局  SFA[m, sfk] → TMEM

+  

+    TLane = m mod 32  · 
+    word = m div 32  · 
+    byte = sfk  · 
+    TCol = word·4 + byte  · 
+    warpx4 → lanes { TLane, +32, +64, +96 }
+  

+  
点击左侧 cell…

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/sm_architecture.html b/_extra_zh/zh/demo/sm_architecture.html
new file mode 100644
index 00000000..77de85ca
--- /dev/null
+++ b/_extra_zh/zh/demo/sm_architecture.html
@@ -0,0 +1,397 @@
+
+
+
+
+
+Blackwell SM 架构
+
+
+
+
+
+  
Blackwell SM 架构

+
+  

+    

+    

+      
流式多处理器 (SM)

+
+      
+      

+        

+          
Tensor Core (tcgen05)

+          
第 5 代 MMA

+        

+        

+        

+          
CUDA Core

+          
FP/INT 单元

+        

+      

+
+      
+      

+        

+          
Shared Memory (SMEM)

+          
每个 SM 228 KB

+        

+        

+          
Tensor Memory (TMEM)

+          
TMEM — 128 条 lane

+        

+        

+          
寄存器文件

+          

+        

+      

+
+      
+      

+        

+          
TMA 引擎

+          
数据搬运器

+        

+        

+        

+      

+    

+    

+      
SM ...

+      × N
+    

+    

+
+    
+    

+      
全局内存 (GMEM)

+    

+
+    
+    
+  

+
+  
+  

+    

+    
点击硬件单元查看详情。实线箭头表示加载路径，虚线箭头表示存储路径。

+  

+
+
+
+
+
diff --git a/_extra_zh/zh/demo/swizzle_128B.html b/_extra_zh/zh/demo/swizzle_128B.html
new file mode 100644
index 00000000..e33b8886
--- /dev/null
+++ b/_extra_zh/zh/demo/swizzle_128B.html
@@ -0,0 +1,432 @@
+
+
+
+
+Swizzle 布局 — SWIZZLE_128B
+
+
+
+
+
+
+
Swizzle 布局 — SWIZZLE_128B

+
physical_sector = logical_sector XOR row  |  8 行 × 32 个 bank，每个 sector 4 个 bank

+
+

+  
读取
+    

+      
+      
+    

+  

+  
索引

+

+
+

+  
+  
不使用 Swizzle

+    

+    
32 个 Bank

+    

+    

+  

+  
SWIZZLE_128B

+    

+    
32 个 Bank

+    

+    

+  

+

+

+
+

+  
按 cycle 展示读取

+  

+  

+    
不使用 Swizzle

+    
SWIZZLE_128B

+  

+

+
+

+  1 个元素 = 1 个 bank（4 字节）。一个 32 线程 warp 读取 32 个元素；bank conflict 会被串行化成额外 cycle。

+  列 tile (8×4)：8 行 × 4 列（1 个 sector）= 32 次读取。
+  使用 swizzle：1 个 cycle；不使用：8 个 cycle（8 路冲突）。

+  行 tile (1×32)：1 行 × 32 列 = 32 次读取。
+  始终是 1 个 cycle（同一行内 32 个 bank 全部不同）。
+

+
+
+
+
diff --git a/_extra_zh/zh/demo/swizzle_8x8.html b/_extra_zh/zh/demo/swizzle_8x8.html
new file mode 100644
index 00000000..def1575e
--- /dev/null
+++ b/_extra_zh/zh/demo/swizzle_8x8.html
@@ -0,0 +1,384 @@
+
+
+
+
+Swizzle 布局
+
+
+
+
+
+
+
Swizzle 布局 — 8×8

+
mapped_col = logical_col XOR row  |  8×8 矩阵，1 个元素 = 1 个 bank

+
+

+  
读取
+    

+      
+      
+    

+  

+  
索引

+

+
+

+  
+  
不使用 Swizzle

+    
Bank 活跃情况

+    

+  

+  
使用 Swizzle（XOR）

+    
Bank 活跃情况

+    

+  

+

+

+
+

+  
按 cycle 展示读取

+  

+  

+    
不使用 Swizzle

+    
使用 Swizzle（XOR）

+  

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/swizzle_atom_general.html b/_extra_zh/zh/demo/swizzle_atom_general.html
new file mode 100644
index 00000000..fd2e7379
--- /dev/null
+++ b/_extra_zh/zh/demo/swizzle_atom_general.html
@@ -0,0 +1,472 @@
+
+
+
+
+Swizzle Atom — 全部格式
+
+
+
+
+
+
+
Swizzle Atom — 全部格式

+
swizzle atom 是内存中应用 swizzle 的基本连续区域

+
+

+  

+    格式
+    
+      
+      
+      
+      
+    
+  

+  
+
+

+
+

+  
读取
+    

+      
+      
+    

+  

+  
索引

+

+

+  
dtype
+    

+      
+      
+      
+    

+  

+  

+

+
+

+  
+  
不使用 Swizzle

+    
Bank sector 活跃情况（全部 8 个）

+    

+  

+  
使用 Swizzle（XOR）

+    
Bank sector 活跃情况（全部 8 个）

+    

+  

+

+

+

+
关键：使用 swizzle 后，读取任意 8×16B 列都没有冲突！

+
+

+  
按 cycle 展示读取

+  

+  

+    
不使用 Swizzle

+    
使用 Swizzle（XOR）

+  

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tcgen05_intro.html b/_extra_zh/zh/demo/tcgen05_intro.html
new file mode 100644
index 00000000..90692ea0
--- /dev/null
+++ b/_extra_zh/zh/demo/tcgen05_intro.html
@@ -0,0 +1,331 @@
+
+
+
+
+Tensor Core：tcgen05.mma
+
+
+
+
+
+
+
tcgen05.mma

+
+

+  
转置 A
+    

+      
+      
+    

+  

+  
转置 B
+    

+      
+      
+    

+  

+

+

+  
M（输出行）
+    

+      
+    

+  

+  
N（输出列）
+    

+      
+      
+      
+      
+    

+  

+  
K
+    

+      
+    

+  

+

+

+  
行分组

+  
列分组

+

+
+

+  

+    
A (SMEM)

+    
+    

+  

+  
×

+  

+    
B (SMEM)

+    

+    

+  

+  

+    

+      
Tensor Core

+      
tcgen05.mma

+    

+    
→

+  

+  

+    
C (TMEM)

+    

+    

+      
128 条 lane（M）

+      

+        

+        
列（N）— 由 tmem.alloc 分配

+      

+    

+  

+

+
+

+  K 迭代：
+  

+

+
+

+

+
+

+  

+    
SMEM 描述符（A, B）

+    
smem_desc = base_addr | leading_byte_offset | stride_byte_offset | start_addr_off

+  

+  

+    
tcgen05.mma PTX

+    
tcgen05.mma.cta_group::1.kind::f16 taddr, a_desc, b_desc, idesc, enable;

+  

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/thread_hierarchy.html b/_extra_zh/zh/demo/thread_hierarchy.html
new file mode 100644
index 00000000..8c49996f
--- /dev/null
+++ b/_extra_zh/zh/demo/thread_hierarchy.html
@@ -0,0 +1,112 @@
+
+
+
+
+
+Blackwell 线程层级
+
+
+
+
+
+
Blackwell 线程层级

+
点击一个层级：thread → warp → warpgroup → CTA → cluster → grid

+
+

+  

+    GRID — 一次 kernel launch
+    

+      CLUSTER — 跨 SM 的 CTA（DSMEM）
+      

+        

+          CTA — 线程块 · 一个 SM
+          

+            WARPGROUP — 4 个 warp · 128 个线程
+            

+              

+                WARP — 32 个线程（SIMT）
+                

+              

+              
warp 1

+              
warp 2

+              
warp 3

+            

+          

+        

+        
CTA 1
…

+      

+    

+  

+

+
+

+  

+  

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/thread_register.html b/_extra_zh/zh/demo/thread_register.html
new file mode 100644
index 00000000..99029747
--- /dev/null
+++ b/_extra_zh/zh/demo/thread_register.html
@@ -0,0 +1,252 @@
+
+
+
+
+线程 + 寄存器布局
+
+
+
+
+
+
+
Layout: S[(8, 4, 2) : (4@laneid, 1@laneid, 1@reg)]

+
点击 cell 查看分配关系

+
+

+  

+    
逻辑 8×8 矩阵

+    

+    

+  

+  

+    
线程分配

+    

+    

+    

+  

+  
+

+
+

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tile_distributed.html b/_extra_zh/zh/demo/tile_distributed.html
new file mode 100644
index 00000000..f0c45819
--- /dev/null
+++ b/_extra_zh/zh/demo/tile_distributed.html
@@ -0,0 +1,480 @@
+
+
+
+
+分布式布局
+
+
+
+
+
+
+
分布式布局 — 2×2 GPU Mesh 上的 8×8 矩阵

+
完全 shard 到 4 个 GPU  |  点击 cell 查看放置位置

+
+

+  
布局
+    

+      
+      
+      
+    

+  

+

+
+

+  

+    

+    

+    

+  

+  

+    
GPU Mesh（2×2）

+    

+    

+  

+  
+

+
+

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tiled_layout.html b/_extra_zh/zh/demo/tiled_layout.html
new file mode 100644
index 00000000..fbe7c961
--- /dev/null
+++ b/_extra_zh/zh/demo/tiled_layout.html
@@ -0,0 +1,278 @@
+
+
+
+
+Tiled 布局
+
+
+
+
+
+
+
S[(4, 2, 2, 4) : (16, 4, 8, 1)]

+
逻辑视图 vs 物理内存  |  点击 cell 查看映射

+
+

+  

+    
逻辑矩阵

+    

+    

+    

+    

+  

+  

+    
物理内存（2×4 tiled）

+    

+    

+    

+    

+  

+  
+

+
+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tiling_constraint.html b/_extra_zh/zh/demo/tiling_constraint.html
new file mode 100644
index 00000000..16480f68
--- /dev/null
+++ b/_extra_zh/zh/demo/tiling_constraint.html
@@ -0,0 +1,343 @@
+
+
+
+
+Swizzle 带来的 tiling 约束
+
+
+
+
+
+
+
Swizzle 带来的 tiling 约束

+
在 16×16 矩阵上使用 128B swizzle — 切成 16×8 分组即可消除 bank conflict

+
+

+  
是否 tile
+    

+      
+      
+    

+  

+  
行范围
+    

+      
+      
+    

+  

+

+

+  
列

+

+
+

+  

+    
原始 16×16 矩阵

+    

+      

+      
+    

+  

+  

+    
为 128B Swizzle 进行 tiling（两个 16×8 分组）

+    

+  

+

+

+  
Bank sector 活跃情况

+  

+  

+

+

+  
dtype
+    

+      
+      
+      
+    

+  

+  

+

+

+  
硬件会基于物理偏移计算 swizzle_col，等价于假设每行只有 8 个元素。若每行有 16 个元素，则每个原始行在 swizzle 计算中会被当成 2 行处理。

+  

+

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tirx_dispatch.html b/_extra_zh/zh/demo/tirx_dispatch.html
new file mode 100644
index 00000000..0470045b
--- /dev/null
+++ b/_extra_zh/zh/demo/tirx_dispatch.html
@@ -0,0 +1,159 @@
+
+
+
+
+
+TIRx：Scope、Layout、Dispatch
+
+
+
+
+
+
TIRx：Scope、Layout、Dispatch

+
点击一个设计元素，查看它控制 single-MMA GEMM 中的哪些代码行

+
+

+  
+  
+  
+

+
+

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tma_3d.html b/_extra_zh/zh/demo/tma_3d.html
new file mode 100644
index 00000000..e118e70b
--- /dev/null
+++ b/_extra_zh/zh/demo/tma_3d.html
@@ -0,0 +1,353 @@
+
+
+
+
+用 3D TMA 做 tiling 与 swizzle
+
+
+
+
+
+
+
用 3D TMA 做 tiling 与 swizzle

+
使用 3D TMA 将数据复制到 tiled shared memory

+
+

+  
Swizzle
+    

+      
+      
+    

+  

+  
列偏移

+

+
+

+  
+
+  

+    
Global Memory

+    
16×256 fp16（连续布局，row = 512B）

+    

+  

+
+  

+    
→

+    
3D TMA

+    

+  

+
+  

+    
Shared Memory（tiled & swizzle 后）

+    

+  

+

+
+

+
+

+  

+

+
+
+
+
diff --git a/_extra_zh/zh/demo/tma_intro.html b/_extra_zh/zh/demo/tma_intro.html
new file mode 100644
index 00000000..65a79ed4
--- /dev/null
+++ b/_extra_zh/zh/demo/tma_intro.html
@@ -0,0 +1,358 @@
+
+
+
+
+TMA：Tensor Memory Accelerator
+
+
+
+
+
+
+
TMA：从 Global Memory 到 Shared Memory

+
硬件 2D copy，可选 swizzle — 一条指令，无需线程搬运

+
+

+  
Swizzle
+    

+      
+      
+    

+  

+  
行偏移

+  
列偏移

+

+
+

+  
+
+  

+    
Global Memory

+    
16×128 fp16（每个 cell = 16B = 1×8 fp16）

+    

+  

+
+  

+    

+      
TMA 引擎

+      
数据搬运器

+    

+    
→

+    
cp.async.bulk.tensor.2d
+ SWIZZLE_128B

+  

+
+  

+    
Shared Memory

+    
8×8 个 sector（swizzle 后）

+    

+  

+

+
+

+
+

+  

+

+
+
    
+  
  • ▸ 单线程派发（非阻塞）
    • 
+  
  • ▸ 自动 swizzle — 数据到达时天然避免 bank conflict
    • 
+  
  • ▸ 双向：load（GMEM→SMEM）和 store（SMEM→GMEM）
    • 
+

+
+
+
+
diff --git a/_extra_zh/zh/viz-base.css b/_extra_zh/zh/viz-base.css
new file mode 100644
index 00000000..314d742d
--- /dev/null
+++ b/_extra_zh/zh/viz-base.css
@@ -0,0 +1,90 @@
+/* Shared base styles for all viz HTMLs
+ * Colors: --bg gray, --accent blue, --surface white
+ * Fonts: Inter (body), SF Mono (code/notation)
+ */
+
+:root {
+  --bg:#fff; --surface:#fff; --border:#dfe1e6; --text:#222; --dim:#888; --accent:#3b82f6;
+
+  /* ── Group palette (tiles, lanes, banks, sectors) ─────── */
+  --color-group-0: #5b9bd5;
+  --color-group-1: #ed9a3c;
+  --color-group-2: #d95555;
+  --color-group-3: #45b5a5;
+  --color-group-4: #5fb85f;
+  --color-group-5: #e0b828;
+  --color-group-6: #9d6eb5;
+  --color-group-7: #e87888;
+
+  /* ── Interaction ──────────────────────────────────────── */
+  --color-hover-bg: #dbeafe;
+  --color-hover-text: #1e40af;
+
+  /* ── Status ───────────────────────────────────────────── */
+  --color-good: #2e7d32;
+  --color-bad: #c62828;
+
+  /* ── Boundaries & neutrals ────────────────────────────── */
+  --color-boundary: #8C1515;
+  --color-cell-bg: #f0f1f3;
+  --color-cell-bg-alt: #e8eaed;
+}
+* { box-sizing:border-box; margin:0; padding:0; }
+body { background:var(--bg); color:var(--text); font-family:'Inter','SF Pro','Segoe UI',system-ui,sans-serif; padding:24px; }
+body.figure h1, body.figure .sub { display:none; }
+body.notitle h1, body.notitle .sub { display:none; }
+h1 { text-align:center; font-size:21px; font-weight:700; margin-bottom:2px; }
+.sub { text-align:center; color:var(--dim); font-size:14px; margin-bottom:20px; font-family:'SF Mono','Fira Code',monospace; }
+
+/* Controls */
+.controls { display:flex; gap:18px; justify-content:center; align-items:center; margin-bottom:10px; flex-wrap:wrap; }
+.lbl { font-size:13px; color:var(--dim); margin-right:4px; font-weight:600; }
+.bg { display:inline-flex; gap:2px; }
+.btn {
+  padding:5px 11px; border:1px solid var(--border); background:var(--surface); color:var(--text);
+  cursor:pointer; border-radius:5px; font-size:13px; font-family:inherit; font-weight:500; transition:all .12s;
+}
+.btn:hover { background:#eef0f4; }
+.btn.on { background:var(--accent); border-color:var(--accent); color:#fff; }
+.btn.on:hover { background:#2563eb; }
+
+/* Side-by-side panels */
+.panels { display:grid; grid-template-columns:1fr 1fr; gap:20px; max-width:1100px; margin:0 auto; position:relative; }
+.panel { background:var(--surface); border-radius:10px; padding:16px 14px; border:1px solid var(--border);
+  box-shadow:0 1px 3px rgba(0,0,0,.06); }
+.panel h2 { text-align:center; font-size:14px; font-weight:700; margin-bottom:2px; }
+.panel .nota { text-align:center; font-size:11px; color:var(--accent); font-family:'SF Mono','Fira Code',monospace;
+  margin-bottom:8px; font-weight:600; }
+
+/* Grid cells */
+.grid { display:grid; gap:3px; }
+.hdr { font-size:10px; color:var(--dim); text-align:center; padding:2px 0; font-weight:600; }
+.rl { font-size:10px; color:var(--dim); display:flex; align-items:center; justify-content:flex-end; padding-right:3px; font-weight:600; }
+.cell {
+  aspect-ratio:1; border-radius:5px; display:flex; flex-direction:column; align-items:center; justify-content:center;
+  font-size:13px; font-weight:700; border:2.5px solid transparent; transition:all .18s;
+  min-width:0; cursor:pointer; line-height:1.15;
+}
+.cell.hov { border-color:#222; border-width:3px; z-index:2; box-shadow:0 0 8px rgba(59,130,246,.4); }
+.cell.dm { opacity:.25; }
+
+/* Arrow SVG overlay */
+.arrow-svg { position:absolute; top:0; left:0; width:100%; height:100%; pointer-events:none; z-index:10; overflow:visible; }
+.arrow-svg path { fill:none; stroke:#222; stroke-width:1.5; }
+.arrow-svg text { font-size:11px; font-weight:600; fill:#222; font-family:inherit; }
+
+/* Formula bar */
+.formula-bar { max-width:1100px; margin:18px auto 0; background:var(--surface); border-radius:10px;
+  padding:12px 18px; border:1px solid var(--border); box-shadow:0 1px 3px rgba(0,0,0,.06); }
+.formula-bar .ftitle { font-size:13px; font-weight:700; margin-bottom:6px; }
+.formula-bar .fcontent { font-size:13px; font-family:'SF Mono','Fira Code',monospace; color:var(--accent); line-height:1.6; }
+
+/* Legend */
+.leg { display:flex; flex-direction:column; align-items:center; gap:4px; margin-top:14px; }
+.leg-row { display:flex; gap:12px; justify-content:center; flex-wrap:wrap; }
+.leg-group { display:flex; gap:10px; align-items:center; }
+.leg-sep { width:1px; height:16px; background:var(--border); }
+.li { display:flex; align-items:center; gap:4px; font-size:12px; color:var(--dim); }
+.swtch { width:14px; height:14px; border-radius:3px; }
+
+@media(max-width:700px) { .panels { grid-template-columns:1fr; } }
diff --git a/_extra_zh/zh/viz-base.js b/_extra_zh/zh/viz-base.js
new file mode 100644
index 00000000..ce3e1a3f
--- /dev/null
+++ b/_extra_zh/zh/viz-base.js
@@ -0,0 +1,66 @@
+// Shared behavior for all viz HTMLs
+document.addEventListener('DOMContentLoaded', function() {
+  var p = new URLSearchParams(location.search);
+  if (p.has('notitle')) document.body.classList.add('notitle');
+
+  // Directly opened demo pages are copied outside Sphinx templates, so they do
+  // not receive html_js_files. Load the same language switch used by book pages.
+  if (window.parent === window && !p.has('notitle')) {
+    var langSwitchPath = window.location.pathname.indexOf('/zh/demo/') >= 0
+      ? '../../_static/lang-switch.js'
+      : '../_static/lang-switch.js';
+    var s = document.createElement('script');
+    s.src = new URL(langSwitchPath, window.location.href).href;
+    document.head.appendChild(s);
+  }
+
+  // Forward arrow keys to parent (reveal.js) when embedded
+  if (window.parent !== window) {
+    document.addEventListener('keydown', function(e) {
+      if ([37, 38, 39, 40, 27, 32].indexOf(e.keyCode) !== -1) {
+        // Left, Up, Right, Down, Escape, Space
+        window.parent.postMessage({ type: 'revealKey', keyCode: e.keyCode }, '*');
+      }
+    });
+  }
+});
+
+// Auto-height: when embedded in the book (demo-embed.js), the demo measures its
+// OWN content height and posts it to the parent, which sizes the iframe to fit so
+// there is never an inner scrollbar. This is push-based on purpose — the demo
+// catches its own DOM changes (a click that appends rows, expands a panel, …),
+// which a parent watching the iframe's  from outside can miss. Measuring
+// body.scrollHeight (not documentElement, which is floored to the viewport) lets
+// the reported height grow AND shrink with the content.
+(function () {
+  if (window.parent === window) return;   // only when embedded
+  var lastH = 0;
+  function report() {
+    var b = document.body, de = document.documentElement;
+    var h = (b ? b.scrollHeight : 0) || (de ? de.scrollHeight : 0) || 0;
+    if (h && Math.abs(h - lastH) > 1) {
+      lastH = h;
+      window.parent.postMessage({ type: 'demoHeight', height: h }, '*');
+    }
+  }
+  var scheduled = false;
+  function schedule() {
+    if (scheduled) return;
+    scheduled = true;
+    requestAnimationFrame(function () { scheduled = false; report(); });
+  }
+  // documentElement exists even while we are still in , so observers can be
+  // attached immediately; the first read happens in the rAF after layout.
+  try { new ResizeObserver(schedule).observe(document.documentElement); } catch (e) {}
+  try {
+    new MutationObserver(schedule).observe(document.documentElement, {
+      subtree: true, childList: true, attributes: true, characterData: true
+    });
+  } catch (e) {}
+  document.addEventListener('DOMContentLoaded', schedule);
+  window.addEventListener('load', schedule);
+  // Clicks often trigger async content changes; re-measure right after.
+  window.addEventListener('click', function () { setTimeout(schedule, 0); }, true);
+  // Catch late settling (fonts, deferred render).
+  [100, 300, 600, 1200].forEach(function (t) { setTimeout(schedule, t); });
+})();
diff --git a/_templates/sbt-sidebar-nav.html b/_templates/sbt-sidebar-nav.html
new file mode 100644
index 00000000..1b39afbc
--- /dev/null
+++ b/_templates/sbt-sidebar-nav.html
@@ -0,0 +1,32 @@
+
diff --git a/conf.py b/conf.py
index 0a03f5a5..1315caf4 100644
--- a/conf.py
+++ b/conf.py
@@ -44,12 +44,13 @@
 html_title = project
 html_logo = "static/mlc-logo-with-text-landscape.svg"
 html_favicon = "static/mlc-favicon.ico"
+templates_path = ["_templates"]
 html_static_path = ["static"]
 # Interactive slide demos (self-contained HTML+CSS+JS) copied verbatim into the
 # site root, then embedded via 
+
+```
+*交互：一个 `mbarrier` 状态视图，展示 arrival counter、phase bit，以及 `init`、`arrive` 和 `wait` 操作；点击字段可以聚焦。*
+
+barrier 从初始化开始。在 `init` 期间，kernel 设置这个 barrier 应该期望多少个 arrival。
+barrier 从 phase 0 开始，counter 被加载为预期 arrival count。从那时起，barrier 就在等待所有必需的 producer
+或某个资源的使用者报告自己已经完成。
+
+arrival 会减少 barrier 仍在等待的工作量。kernel 的不同部分可以用不同方式 arrive 到 barrier，而这种区别很重要。
+
+对于 TMA load，常见 arrival path 是 tx-count arrival。像 `mbarrier.arrive.expect_tx(bytes)` 这样的操作会做两件事：
+第一，它算作 issuing thread 在 barrier 上的 arrival。第二，它记录 TMA engine 预计传输的字节数。
+barrier 不会仅仅因为 issuing thread 已经 arrive 就完成。它还会等待 TMA engine 随着传输结束把 byte count drain 掉。
+只有两个条件都满足时，phase 才会翻转：普通 arrival count 到达零，pending tx byte count 也到达零。
+
+这就是为什么不应该把 `expect_tx` 理解为“又一个普通 arrival”。它为异步 copy 设置 byte budget。
+硬件稍后通过 complete-tx update 记账实际 copy completion。只有 arrival 和 byte transfer 都完成时，barrier 才完成。
+
+对于 Tensor Core 工作，arrival path 不同。`tcgen05` MMA 不会仅仅因为 MMA 已经发射就自动推进 barrier。
+kernel 必须显式地把 barrier arrival 附着到 commit path 上，例如使用 `tcgen05.commit.mbarrier::arrive` 操作。
+当这个 committed group 完成时，Tensor Core 侧会执行 barrier arrival。如果 kernel 忘了这个 commit arrival，
+等待 barrier 的 consumer 会永远等下去。
+
+普通 thread 也可以直接 arrive 到 barrier。当普通 thread code 是 producer，或一组 thread 在宣布自己已经用完某个资源时，会使用这种方式。
+例如，consumer 读完 shared-memory buffer 后，可以 arrive 到一个 barrier，告诉 producer 这个 buffer 可以复用了。
+
+waiting 是同一协议的 consumer 侧。consumer 会等待，直到 barrier 完成当前 iteration 所期望的 phase。
+只有这时，读取数据或复用该 barrier 保护的资源才是安全的。
+
+重要的一点是，异步硬件不仅会跑在程序前面；它也会通过 barrier 把 completion 报告回来。
+TMA 可以 signal 一个 shared-memory tile 已经 ready。Tensor Core 工作可以 signal TMEM 结果已经 ready。
+普通 thread 可以 signal 某个 buffer 不再被使用。barrier 给这些情况统一了 producer-consumer 形状：
+producer arrive，consumer wait。
+
+## Phase Tracking
+
+barrier 通常不会只为一次使用而分配。pipelined K-loop 可能执行同一个 handoff 数百次，
+如果每次 iteration 都分配新的 shared-memory barrier，就不现实。因此，kernel 会保留一小组固定 barrier，
+并随着 loop 推进反复使用它们。
+
+phase bit 让这种复用变得安全。
+
+```{raw} html
+

+
+

+```
+*交互：一个在多个 pipeline iteration 中复用的 barrier，展示 phase bit 如何在每个完成轮次后翻转。*
+
+每当 barrier 完成当前轮的所有 arrival，它都会翻转 phase：phase 0 变成 phase 1，phase 1 变成 phase 0，如此往复。
+wait 操作会检查 consumer 期望的 phase。这个 expected phase 由 kernel 保存在寄存器中。
+当某个 stage 成功等待一轮后，kernel 会在下一轮使用该 barrier 前切换自己的本地 phase value。
+
+这防止 kernel 把旧 completion 误认为新 completion。假设某个 barrier 用于一次 TMA load 并已经完成。
+如果下一个 loop iteration 复用同一个 barrier 却不追踪 phase，consumer 可能观察到上一次 completion，
+并错误地认为新的 load 已经 ready。phase bit 把这两轮分开：iteration 0 等待一个 phase，
+iteration 1 等待相反 phase，iteration 2 再次等待第一个 phase，模式持续下去。
+
+在真实 pipeline 中，bookkeeping 通常按 stage 进行。kernel 有固定数量的 shared-memory stage，
+匹配固定数量的 barrier，以及寄存器中一小组 phase value。随着 loop 前进，每个逻辑 iteration 映射到一个物理 stage，
+phase value 告诉 wait 操作它正在等待这个物理 barrier 的哪一轮。
+
+这就是为什么后面的 GEMM 代码不需要每个 K tile 一个 barrier（{ref}`zh_chap_gemm_async`）。
+它需要每个 reusable stage 一个 barrier，再加上 phase tracking。stage index 选择 shared-memory buffer 和 barrier。
+phase value 区分这个 stage 当前使用和上一次使用。
+
+**可以让你的 agent 试试**：给它一个 two-stage pipeline，让它追踪四次 iteration。
+对每次 iteration，列出 stage index、本地 phase value、barrier 何时翻转，以及如果 stage 复用前没有切换 phase 会出什么问题。
+
+## 同步规则
+
+一旦 barrier 和 phase 机制清楚了，tensor-core kernel 中的同步 pattern 就相当机械。
+每当一条路径产生数据，或释放另一条路径将要消费的资源时，handoff 都必须显式完成。
+
+常见有三种情况。
+
+第一种情况是 thread code 为异步引擎产生数据。如果 thread 写 shared memory，后续 TMA store 或 MMA 指令会读取这块 shared memory，
+kernel 就必须在引擎读取前让 thread 写入可见。这需要合适的 thread-level synchronization 或 fence。
+精确指令取决于 handoff 的 scope，但原因始终相同：引擎不能在 producer thread 完成写入前观察 shared-memory buffer。
+
+第二种情况是 TMA 为 MMA 产生数据。TMA load 会异步填充 shared-memory tile。
+MMA 路径不能只因为 TMA 指令已经发射，就推断 tile 已经 ready。
+TMA 操作必须关联一个 `mbarrier`，而 MMA 路径必须在读取 tile 前 wait 这个 barrier。
+
+第三种情况是 MMA 为 epilogue 产生数据。`tcgen05` MMA 会异步把结果写入 TMEM。
+在 Tensor Core 完成相关工作之前，epilogue 不能安全读取 accumulator。
+因此 MMA commit path 会 arrive 到一个 completion barrier，epilogue 在读取 TMEM 前 wait 这个 barrier。
+
+```{raw} html
+

+
+

+```
+*交互：TMA load 通过 `mbarrier` signal completion。MMA 路径在读取 shared-memory tile 前等待 barrier。
+Tensor Core 到 epilogue 的 handoff 形状相同，只是执行 arrival 的不是 TMA，而是 Tensor Core commit path。*
+
+同一个思想也适用于 resource reuse。barrier 不只是 data-ready signal，也可以是“resource is free” signal。
+在旧 tile 的所有 consumer 都用完它之前，shared-memory stage 不能被覆写。
+在前一个使用者完成读写之前，TMEM region 不能被复用。在这些情况下，arrival 表示“我用完这个资源了”，
+wait 表示“现在可以安全地为下一个 stage 复用这个资源了”。
+
+这正是阅读 pipelined GEMM kernel 中同步逻辑的正确方式。wait 和 arrive 并不是作为 defensive programming 四处散落。
+每一个都标记一次具体 ownership transfer：tile 变得 ready、accumulator 变得可读，或 buffer 变得可复用。
+一旦识别出这些 handoff，control flow 就会容易跟随得多。
diff --git a/zh/chapter_background/index.md b/zh/chapter_background/index.md
new file mode 100644
index 00000000..c69a3541
--- /dev/null
+++ b/zh/chapter_background/index.md
@@ -0,0 +1,168 @@
+(zh_chap_background)=
+# GPU 执行模型
+
+:::{admonition} Overview
+:class: overview
+
+- 一个 kernel 会在一套线程层级（thread → warp → warpgroup → CTA → cluster → grid）上运行，并跨越不同的内存空间（寄存器、SMEM、GMEM、TMEM）。
+- 计算被划分到 CUDA core 和 Tensor Core；像 TMA 这样的专用引擎负责搬运供它们消费的数据。
+- 一个 kernel 本质上是一条 pipeline：它把数据暂存到这些内存空间中，并在彼此独立的计算引擎和数据移动引擎之间交接工作；反复出现的目标，是让这些引擎同时保持忙碌。
+:::
+
+要写出高速 GPU 程序，理解硬件本身以及代码如何在硬件上运行非常重要。本章概览 GPU 的执行模型：
+执行工作的线程层级，存放和移动数据的内存空间，以及承担重活的计算引擎和数据移动引擎。
+我们会先逐一介绍这些部件，然后把它们组合进一条 GEMM pipeline 中，从而看清数据和执行如何流经硬件。
+本书后续几乎每一种优化，本质上都是以某种方式在这些相同部件之间安排工作。
+
+现代 GPU 还包含许多专门化的硬件单元。为了先建立一个直观印象，在深入每个部件之前，下面的交互式演示展示了
+Blackwell streaming multiprocessor 内部的主要元素。你可以点击各个部分查看细节。
+
+```{raw} html
+

+
+

+```
+*交互：Blackwell SM，展示其中的 warp/warpgroup、shared memory、Tensor Memory，以及
+Tensor Core 和 TMA 引擎。*
+
+## 执行层级
+
+我们从真正执行工作的线程开始。GPU 并不会把成千上万个线程呈现为一个扁平的池子，而是把它们组织成嵌套层级。
+这样做的原因是，协作会同时发生在几个不同尺度上。每一层的存在，都是为了让某个尺度上的协作更廉价。
+下图展示了 Blackwell 上的线程层级；你可以点击每一层来高亮它。
+
+```{raw} html
+
+```
+*交互：点击某一层：thread → warp → warpgroup → CTA → cluster → grid。*
+
+- **Thread**：标量执行单元。每个 thread 都有自己的程序计数器和寄存器，并通过它在所属 warp 内的 lane ID 来标识。
+- **Warp**：以 SIMT（*single instruction, multiple threads*）方式执行的 32 个 thread。一个 warp 的各个 lane 会一起发射同一条指令，
+  但每个 lane 保留自己的寄存器，也可以被单独 mask 掉；这正是单个 warp 中不同 lane 能够走不同分支的原因。
+- **Warpgroup**：四个连续的 warp，也就是 128 个 thread。Hopper 引入 warpgroup，作为发射 warpgroup 级 MMA（`wgmma`）的单位；
+  在 Blackwell 上，它又承担了第二个角色：Tensor Memory 访问的协作单位。128 个 thread 会一起把一个 TMEM tile 移入或移出寄存器。
+- **CTA**（*Cooperative Thread Array*，也就是 CUDA 所说的 thread block）：硬件调度的基本单位。
+  一个 CTA 运行在单个 SM 上，并拥有该 SM 内一块私有的 shared-memory 分配。同一个 SM 上可以同时驻留多个 CTA；
+  这种情况下，它们会瓜分该 SM 的 shared-memory 容量。
+- **Cluster**：一组相互协作的 CTA，它们可能位于不同的 SM 上。cluster 内的 CTA 可以彼此同步，也可以读写彼此的 shared memory；
+  这种能力称为 distributed shared memory。
+
+这些层级值得多停留一下，因为不同于更早的架构，Blackwell 的关键操作**并不全都由同一组线程发射**。
+TMA copy 由单个 thread 发起，随后由硬件执行。TMEM 到寄存器的 load 是 warpgroup-distributed 的：
+四个 warp 共同协作，每个 warp 移动 TMEM tile 中属于自己的切片。`tcgen05` MMA 由一个被选出的 thread 提交，
+而 clustered MMA 会一次跨越两个 CTA。因此，每种操作都有自己的自然粒度；运行该操作的线程集合，
+就是我们所说的该操作的 **scope**。scope 是本书反复回到的三个设计元素（scope、layout 和 dispatch）中的第一个。
+
+## 内存空间
+
+这一层级中的线程能跑多快，取决于数据能多快到达它们手中。因此接下来我们看数据住在哪里。
+不存在一种既大又快的单一内存；物理规律迫使容量和速度之间做取舍。所以 GPU 提供的不是一种内存，而是多种内存，
+每一种都在不同的位置取得这种折中；kernel 的工作，就是让数据流经这些内存空间。每个空间都有自己的容量、
+延迟，以及关于谁能访问它的规则。
+
+| 内存 | 所属范围 | 作用 | 说明 |
+|--------|-----------|------|-------|
+| **Global (GMEM)** | 整个 device | 持久化 tensor 存储 | 大容量 HBM，由所有 SM 共享 |
+| **Shared (SMEM)** | 每个 CTA（一个 SM） | tile 暂存 | 低延迟 scratchpad；B200 上最高 228 KB/SM |
+| **Tensor Memory (TMEM)** | 每个 CTA | MMA accumulator 存储 | Blackwell 新增；供 `tcgen05` 使用 |
+| **Register File (RF)** | 每个 thread | 标量和每线程 tile fragment | 很快；保存 epilogue/临时值 |
+
+按顺序读，这些空间描述了一条路径。本书中几乎每个 kernel 的数据路径都是
+**GMEM → SMEM →（compute）→ registers → SMEM → GMEM**；对于 tensor-core kernel，TMEM 位于这条路径中间，
+在数学计算运行时保存 accumulator。
+
+在这四者中，**Tensor Memory (TMEM)** 是唯一一个在 Blackwell 之前没有对应物的空间；完整细节会留到
+{ref}`zh_chap_tensor_cores`。不过，现在先理解它的动机很有价值。早期 GPU 把大型 MMA accumulator 保存在寄存器中，
+而寄存器是稀缺资源，accumulator 会与其他值竞争。Blackwell 则把 `tcgen05` 的 accumulator 输出写入 TMEM：
+这是一个 CTA 作用域的二维 scratchpad，每个 CTA 有 128 个 lane，最多 512 个 32-bit column
+（这个数组物理上位于 SM 上）。随后 kernel 必须在 epilogue 之前显式地把 TMEM 读回寄存器。
+这个额外步骤并不是免费的，它带来的两个后果会贯穿全书。第一，TMEM read 是**显式且 warpgroup-distributed** 的，
+由一个 warpgroup 的四个 warp 协作完成。第二，TMEM 不同于寄存器，必须被**显式分配和释放**。
+
+### 跨 cluster 的 Distributed Shared Memory
+
+cluster 是这个层级中唯一一个成员可以跨越多个 SM 的层级；这种可达范围带来了一种其他层级没有的内存能力。
+一个 CTA 运行在一个 SM 上，并使用该 SM 的 shared memory，但单个 CTA 的 SMEM 预算有限，而大 tile 通常需要比一个 block
+单独能提供的更多 operand 存储，或者更多复用。Hopper 给出的答案是 **thread block cluster**：
+一组比独立 block 更紧密协作的 CTA；它们可以一起同步，也可以读写彼此的 shared memory，这种能力称为
+**distributed shared memory (DSMEM)**。Blackwell 保留了 cluster，并在此基础上加入动态调度
+（{ref}`zh_chap_clc`）和 2-CTA cooperative MMA。
+
+DSMEM 允许一个 CTA 直接寻址并访问另一个 peer CTA 的 shared memory。一个 thread 可以命名 peer 的 SMEM 中的某个位置，
+并把一个 tile 从自己的 SMEM 直接 bulk-copy 到对方 SMEM 中；当字节落地后，会触发 completion barrier
+（{ref}`zh_chap_async_barriers`）。第三部分中的 2-CTA cluster GEMM 正是建立在这个机制之上：它利用 DSMEM 在一对 CTA
+之间共享 operand tile，而不需要把数据绕回 global memory。
+
+下图展示了 CTA cluster 让额外的 DSMEM hop 成为可能；点击某个部分，可以看到每个 CTA 拥有什么，以及 cross-CTA read
+发生在哪里。
+
+```{raw} html
+

+
+

+```
+*交互：一个 2-CTA cluster，其中每个 CTA 拥有 A 的一半和 B 的一半，通过 cluster（DSMEM）读取对方的 B，
+两者共同产生一个 256×256 的输出 tile。*
+
+## 计算：CUDA Core 与 Tensor Core
+
+线程以及它们搬运的数据，最终必须在算术单元处相遇。一个 SM 提供的不是一种数学引擎，而是两种不同的数学引擎。
+二者之间的分工塑造了几乎每个 kernel 的写法，并且它们扮演互补角色。
+
+- **CUDA core** 是通用 SIMT ALU。它们运行标量和向量指令，用于处理索引算术、elementwise 计算、
+  reduction 和控制流，也就是围绕重型矩阵工作的“胶水逻辑”。
+- **Tensor Core** 是固定功能单元，在 *tile* 粒度执行 dense matrix multiply-accumulate，
+  用一条指令计算 $D = AB + C$。
+
+这种划分之所以重要，是因为 Tensor Core 提供的算术吞吐远高于 CUDA core，FLOP/s 通常高出一个数量级甚至更多。
+因此，dense linear algebra（GEMM、convolution 和 attention）只有运行在 Tensor Core 上，才能接近峰值性能。
+所以，获得性能在很大程度上就是让这些 Tensor Core 持续有数据可算。不同 GPU 世代之间变化的是 Tensor Core
+**如何**被编程，以及它们的结果**落在哪里**。Hopper 引入了异步 warpgroup MMA（`wgmma.mma_async`）；
+Blackwell 的第五代 Tensor Core，即 `tcgen05`，把 accumulator 放在 Tensor Memory 中，而不是寄存器中；
+我们会用 {ref}`zh_chap_tensor_cores` 专门介绍它。
+
+cluster 以两种方式扩展这些引擎，而这两种方式会在 GEMM 章节中反复出现。**2-CTA cooperative MMA**
+让两个 CTA 各自贡献自己的 SMEM operand，共同形成一个更大的 Tensor Core MMA tile。**TMA multicast**
+让数据移动引擎的一次 load 同时把同一个 GMEM tile 送到多个 CTA，消除本来由多次独立 load 造成的冗余 global traffic。
+二者都建立在前面介绍的 distributed shared memory 之上。
+
+## GEMM 数据 Pipeline
+
+到目前为止，我们已经分别介绍了各个硬件单元。为了看清它们如何协同工作，可以用一条典型的通用矩阵乘法（GEMM）
+pipeline 作为例子。下面的交互式演示展示了三阶段 GEMM tile pipeline 中涉及的单元；点击诸如 `tma load`
+这样的动作，可以高亮它穿过各硬件单元时所走的数据路径。
+
+```{raw} html
+

+
+

+```
+*交互：Blackwell 上的 load → MMA → epilogue pipeline；点击一个动作，追踪它跨硬件单元的数据路径。*
+
+一个 GEMM tile 会流经三个阶段。
+
+1. **Load。** 一个 TMA copy（{ref}`zh_chap_tma`）把 A 或 B operand tile 从 GMEM 流式搬入 SMEM。
+   一个 thread 发射这次 copy，并预先记录预计会到达多少字节。当字节落地时，TMA 引擎报告进度；
+   只有当所有预期字节都已送达后，completion barrier 才会翻转。
+2. **Compute。** 一个 `tcgen05` MMA（{ref}`zh_chap_tensor_cores`）从 SMEM 中读取 operand tile，
+   并把乘积累加进一个 TMEM tile。一个被选出的 thread 发射它；数学计算完成后，它会 signal 一个 barrier。
+3. **Epilogue。** warpgroup 把 TMEM accumulator 读回寄存器，把结果 cast 成输出 dtype，然后存到 GMEM；
+   这通常会先暂存到 SMEM，再发射一次 TMA store。
+
+这样写出来，三个阶段看上去是严格串行的；但慢 kernel 和快 kernel 的全部差异，就在于 **overlap**。
+朴素 kernel 确实会按顺序执行这些步骤（load、wait、compute、wait、store），于是每个引擎在等待前一个引擎时都会闲置。
+快速 kernel 则把它们 pipeline 起来：Tensor Core 正在计算 tile `k` 时，TMA 引擎已经在获取 tile `k+1`，
+epilogue 也正在忙着排空 tile `k-1`，因此三个引擎可以同时保持占用。让三个异步引擎安全地相互交接工作，
+正是 barrier 和 phase 模型（{ref}`zh_chap_async_barriers`）的职责；第三部分的 GEMM 阶梯就是建立在这个模型之上。
+
+## 接下来读什么
+
+现在我们已经看过高层图景，可以继续阅读深入解释主要机制的章节：
+
+- {ref}`zh_chap_tensor_cores` 详细解释 `tcgen05` 计算和 Tensor Memory。
+- {ref}`zh_chap_tma` 介绍基于 TMA 的异步数据移动。
+- {ref}`zh_chap_async_barriers` 介绍用于协调这些引擎的 mbarrier 和 phase 模型。
diff --git a/zh/chapter_clc/index.md b/zh/chapter_clc/index.md
new file mode 100644
index 00000000..671d7623
--- /dev/null
+++ b/zh/chapter_clc/index.md
@@ -0,0 +1,111 @@
+(zh_chap_clc)=
+# 进阶：Cluster Launch Control
+
+:::{admonition} 概览
+:class: overview
+
+- persistent kernel 会保持一组固定 CTA 或 CTA cluster 驻留（通常让规模大致达到每个 SM 一个活跃 work owner，但不依赖保证的 1:1 映射），并让它们循环处理许多 output tile，而不是每个 tile 启动一个 CTA。
+- Cluster Launch Control 是 Blackwell 的硬件机制，允许驻留 cluster 在运行时请求另一个 tile。它是一条围绕两条 PTX 指令构建的硬件 work-stealing 路径：一条指令请求工作，另一条读回请求是否成功。
+- 主要收益是更好的 tail behavior。当 tile 成本不均，或者 tile 数量不能均匀分配到可用 SM 时，提前完成的 CTA 可以拉取更多工作，而不是闲置。
+:::
+
+persistent GEMM 不会把 CUDA grid 当成固定的“每个 output tile 一个 CTA”的 launch。
+相反，它启动一组更小的、长生命周期的 CTA 或 CTA cluster。每一个计算一个 tile，前进到另一个 tile，再次计算，
+并持续执行，直到输出空间完成。这正是 {ref}`zh_chap_gemm_advanced` 中逐步构建的执行模式。
+
+一旦 kernel 是 persistent 的，主要调度问题就变得很简单：当一个 CTA 或 cluster 完成当前 tile 后，下一个 tile 从哪里来？
+
+最简单的答案是静态公式。例如，kernel 可以从 CTA id 计算 tile coordinate，然后按 grid stride 前进。
+这很容易实现，并且当所有 tile 成本大致相同、tile 数量能均匀分布到 GPU 上时效果很好。
+但 schedule 是在实际工作运行前决定的。如果少数 tile 花费更久，或者最后几个 tile 分配不均，
+有些 SM 会提前完成自己的份额，而另一些仍在处理 tail。
+
+Cluster Launch Control，即 CLC，会改变这个调度模型。persistent cluster 不再预先决定整个 assignment，
+而是可以向硬件 grid scheduler 请求另一个尚未 launch 的 cluster 的工作。如果请求成功，当前 cluster 接管那个 cluster coordinate，
+并计算对应 tile。如果请求失败，就没有更多工作可偷，loop 退出。
+
+这与 thread block cluster 本身不是一回事。thread block cluster（一起 launch 的 CTA，具有 cluster-level synchronization，
+并能访问 distributed shared memory）是在 Hopper 中引入的（{ref}`zh_chap_background`）。
+CLC 是 Blackwell 增加的机制，让这些 cluster coordinate 上的调度变为动态。
+cluster 已经是 launch 单位；CLC 让已经运行的 cluster 可以取消一个 pending launch，并继承它的坐标。
+
+## 两条指令
+
+Cluster Launch Control 通过两条 PTX 指令暴露。第一条指令向 grid scheduler 发送异步请求，第二条指令读取响应。
+
+请求指令是 `clusterlaunchcontrol.try_cancel.async`。
+
+`try_cancel` 会要求 scheduler 取消一个 pending cluster 的 launch，并把该 cluster 的坐标返回给调用者。
+响应会作为 16-byte record 写入 shared memory。由于请求是异步的，指令不会等待响应到达。
+completion 会通过 `mbarrier` 报告，使用与 TMA 相同的 barrier-and-phase 模型。
+
+这是一个重要细节，因为它意味着 CLC 没有引入新的等待模型。kernel 发射请求，把它关联到一个 barrier，
+随后在读取响应前等待 barrier。响应到达通过带 byte-count completion 的 barrier signal，
+整体风格与其他异步硬件操作相同（见 {ref}`zh_chap_async_barriers`）。
+
+一旦 barrier 触发，kernel 使用 query 指令。
+
+第一个 query 是 `clusterlaunchcontrol.query_cancel.is_canceled`。它返回一个 predicate，告诉 kernel cancel 是否成功。
+predicate 为 true 表示 scheduler 找到了一个 pending cluster launch、取消了它，并返回了其坐标。
+predicate 为 false 表示没有剩余 pending work 可取。
+
+只有当 `is_canceled` 为 true 时，kernel 才应该读取 coordinate。它通过
+`clusterlaunchcontrol.query_cancel.get_first_ctaid` 完成这件事，该指令提取被取消 cluster 的第一个 CTA id。
+这个 CTA id 是 coordinate vector，通常读作 `(x, y, z)`，kernel 会把它 decode 成接下来应该计算的 output tile。
+
+这个协议里没有数值形式的 sentinel tile id。kernel 根据 predicate 分支。如果 predicate 为 true，coordinate 有效。
+如果 predicate 为 false，work-stealing loop 结束。
+
+在底层，这个形状直接来自 CLC 正在做的事。硬件不是从软件队列中分配一个抽象 task；
+它是在取消一个尚未发生的 cluster launch。因此，成功响应包含一个真实 cluster coordinate；
+失败响应只是表示 launch queue 已经耗尽。
+
+## Work-Stealing Loop
+
+有了这两条指令，persistent scheduler 就变成一个短 loop。
+
+在 loop 的任意时刻，cluster 都有一个自己负责计算的 tile。在开始这个 tile 之前，它会为下一个 tile 发送 `try_cancel` 请求。
+请求异步运行。当 scheduler 处理这个请求时，cluster 计算自己的当前 tile。
+
+当前 tile 完成后，cluster 会等待与 `try_cancel` 响应关联的 `mbarrier`。
+响应 ready 后，它调用 `query_cancel.is_canceled`。如果 predicate 为 true，它调用 `query_cancel.get_first_ctaid`，
+decode 返回的 coordinate，并把它作为下一个 tile。如果 predicate 为 false，就没有剩余工作，cluster 退出。
+
+代码形态上，这个 loop 是：
+
+1. 为可能的下一个 tile 发射 `try_cancel`；
+2. 在请求 in flight 时计算当前 tile；
+3. 等待 response barrier；
+4. 查询 cancellation 是否成功；
+5. 要么用返回的 coordinate 继续，要么退出。
+
+请求的位置正是这个 loop 有用的原因。cluster 不会等当前 tile 完成后才请求更多工作。
+它先请求，再计算。这样就把 scheduler request 与有用工作 overlap 起来。
+当当前 tile 完成时，下一个 tile 的答案往往已经可用。
+
+这与 persistent kernel 在其他地方使用异步 copy 和 tensor-core barrier 的基本原因相同：
+kernel 避免把长延迟操作直接放到 critical path 上。CLC 把同样想法应用到 tile scheduling：
+提前请求下一份工作，计算当前工作，然后在需要时消费调度结果。
+
+## 与 Persistent GEMM 的关系
+
+{ref}`zh_chap_gemm_advanced` 中的 persistent GEMM 在主线讲解中使用 static scheduler。
+static scheduler 更容易解释，因为下一个 tile 可以直接从 loop state 计算出来。
+例如，像 `ClusterPersistentScheduler2D` 这样的 scheduler 可以在 output tile space 上用 grid-stride pattern 分配 tile。
+
+CLC 是这种 static assignment 的动态替代。outer loop 保持不变：每个 resident cluster 反复计算一个 output tile，
+然后前进到另一个 tile。变化的是下一个 tile 从哪里来。使用 static scheduler 时，下一个 tile 由公式计算。
+使用 CLC 时，下一个 tile 由硬件 work stealing 返回。
+
+这种差异在 launch tail 附近最重要。在 static schedule 中，剩余工作可能分布不均。
+有些 SM 可能已经耗尽 assigned tile，而其他 SM 仍有几个 tile。使用 CLC 时，提前完成的 cluster 会请求另一个 pending cluster coordinate。
+只要 launch queue 中还有工作，提前完成者就会继续拉取更多 tile。
+
+当 tile cost 不均匀时，这也很重要。一些 GEMM tile 可能因为边界、masking、sparsity、grouped scheduling，
+或主矩阵乘法周围的 fused work 而走不同路径。static schedule 在观察到这些成本之前，就假设 tile assignment 足够好。
+CLC 不需要这个假设。它只在某个 cluster 变得可用之后，才分配更多工作。
+
+因此，在 TIRx 中，CLC 可以暴露为 dynamic tile scheduler。编程模型不需要改变 tile 的计算。
+tile body 仍然是 static scheduler 使用的同一个 persistent GEMM body。scheduler 从“用公式计算我的下一个 tile coordinate”
+变成“向硬件请求下一个可用 cluster coordinate”。结果仍然是同一个 persistent loop，
+但工作分布由硬件驱动，而不是由固定 launch-time schedule 决定。
diff --git a/zh/chapter_data_layout/index.md b/zh/chapter_data_layout/index.md
new file mode 100644
index 00000000..5b299eba
--- /dev/null
+++ b/zh/chapter_data_layout/index.md
@@ -0,0 +1,256 @@
+(zh_chap_data_layout)=
+# 数据布局及其记法
+
+:::{admonition} 概览
+:class: overview
+
+- *数据布局*把 tensor 的逻辑索引映射到物理位置，并决定 coalescing、bank conflict，以及某个引擎能否读取一个 tile。
+- 本书用一种记法书写布局：`S[(shape) : (strides)]`，并配合 named axes（`@laneid`、`@TLane` 等）以及用于 broadcast 或复制数据的 replication 项 `R[...]`。
+- Swizzle 是一种对地址做 XOR 重映射的机制，用来消除 shared-memory bank conflict。
+:::
+
+同一组数字，如果以不同的物理排列写入内存，在同一块 GPU 上的运行速度可能相差一个数量级。
+
+原因在于，tensor 的逻辑索引并不会说明它的字节实际位于哪里。硬件对这个位置极其敏感：
+它决定了 32 个 lane 的 load 是 coalesce 成一次 transaction，还是散成 32 次；
+决定了它们的地址是落在不同 memory bank 中，还是碰撞并串行化；
+甚至决定了一个 tile 的字节排列是否能被 Tensor Core 读取。
+
+机器学习程序通常用逻辑 shape 来描述 tensor。**数据布局**补上缺失的物理部分：
+它说明具有逻辑索引 `(i, j, …)` 的元素住在哪里，是在内存、寄存器，还是某种其他硬件存储中。
+
+本章介绍现代 GPU 编程中出现的主要布局。为了让讨论可控，我们发展出一种紧凑的**记法**，
+用它描述机器学习系统会遇到的多种场景。最后我们会讨论 **swizzling**：
+它是一种让同一个 tile 的按行访问和按列访问都能同时高效的机制。
+
+## Shape–Stride 模型
+
+在进入 GPU 特有布局之前，值得先从最简单的布局开始，因为本章后面的所有内容都建立在它之上。
+从核心上说，一个 layout 只有两部分：一个 **shape**，以及一组匹配的 **strides**。
+我们把这对信息写成 `S[(shape) : (strides)]`；要找到某个逻辑索引的位置，只需把该索引与 strides 做点积。
+例如，一个 row-major 的 4×4 矩阵可以写成：
+
+```text
+S[(4, 4) : (4, 1)]        addr(i, j) = i·4 + j·1
+```
+
+这不过是经典 shape/stride 模型的一种紧凑写法（也是 CuTe 记法的 row-major 简化版），后续一切都从它构建出来。
+
+事实上，你几乎肯定已经用过这个模型。任何写过 PyTorch 或 NumPy 的人都用过，因为这些库里的 tensor
+本质上就是一个 shape，加上一组作用于扁平 storage buffer 的 stride：
+
+```python
+import torch
+t = torch.arange(12).reshape(3, 4)
+t.shape        # torch.Size([3, 4])
+t.stride()     # (4, 1)        ← exactly S[(3, 4) : (4, 1)]
+```
+
+一旦你这样看待 tensor，就会明白为什么许多“reshape”操作根本不触碰数据。
+它们只是重写 strides，并返回同一份 storage 上的一个 **view**。最清楚的例子是 transpose，也就是 permute：
+
+```python
+tt = t.permute(1, 0)               # or t.T
+tt.shape                           # torch.Size([4, 3])
+tt.stride()                        # (1, 4)        ← strides swapped, no data moved
+tt.data_ptr() == t.data_ptr()      # True, same bytes
+```
+
+这里，`t.permute(1, 0)` 是同一块内存上的 `S[(4, 3) : (1, 4)]`：
+transpose 纯粹是 stride 的变化，没有移动任何一个字节。对 contiguous tensor 做 `reshape` 或 `view` 也是同样故事：
+在旧 storage 上给出新的 shape 和新的 strides。（NumPy 的行为完全相同；唯一差别是它的 `.strides`
+以字节为单位，而不是以元素为单位。）
+
+GPU 上的 layout 正是这样工作的。本章剩余内容其实都是同一个思想的各种变体：
+一个 tile 的映射（无论映射到内存，还是通过稍后介绍的 named axes 映射到 lane 和寄存器）都是固定 buffer 上的一条 stride 规则，
+所以重新排列 tile 通常是改变 *layout*，而不是 copy。不过，我们也要小心这种推理的边界。
+zero-copy 的故事对于单一线性地址空间上的逻辑 view 非常清晰；但在 GPU 上，只有当新的 view 与既有字节排列和 ownership
+安排兼容时才成立。一旦你改变某个元素由哪个 thread 或 register 拥有，或者改变 SMEM swizzle，
+通常就需要真实的数据移动：load、store、shuffle、`ldmatrix`、transpose。
+
+## Tile Layout
+
+到目前为止，我们描述的是整个 tensor 的 layout。不过，GPU kernel 很少一次操作整张矩阵；
+它们处理更小的 tile，而这些 tile 会由硬件的不同部分 load、transform 和 compute。
+好消息是，tiling 并不需要新概念。它仍然只是一个 layout，只是现在多写了几个维度。
+把一个 8×8 矩阵切成 2×4 的 tile，就得到一个 4-D layout，其坐标为
+`(tile_row, row_in_tile, tile_col, col_in_tile)`，并选择 strides 让每个 tile 保持 contiguous：
+
+```text
+S[(4, 2, 2, 4) : (16, 4, 8, 1)]
+```
+
+一个逻辑 `(i, j)` 会先变成 `(i//2, i%2, j//4, j%4)`，然后通过 strides 计算地址。
+值得注意的是，这个记法完全不需要特殊的“tile”概念就能表达 tiling：
+它仍然是前面的 shape–stride 模型，只是把索引拆成了外层和内层坐标。
+
+下面的交互式可视化展示了逻辑矩阵索引如何被分解成 tile 坐标，然后映射到物理地址。
+
+```{raw} html
+
+```
+*交互：点击一个单元格，查看它的 tiled index 和 address。*
+
+## Named Axes
+
+到目前为止，`S[...]` 中的每个 stride 都表示线性内存中的 offset，而我们也把 address 当成内存中的位置。
+但在 GPU 上，数据可以住在不止一个地方：除了内存，一个 tile 也可能分散在 warp lane、thread register，
+或者 TMEM lane 和 column 之中。为了统一描述这些情况，我们用 **named axes** 扩展记法。
+思路是让每个 stride 系数携带一个轴标签，说明它沿着哪个空间移动：
+`@m` 表示普通内存，`@laneid` 表示 warp lane，`@reg` 表示寄存器，`@warpid` 表示 warp，
+`@TLane` / `@TCol` 表示 TMEM 坐标。有了这些标签，单个 layout 不仅能描述数据位于内存何处，
+还能描述它如何分布在负责操作它的硬件资源上。
+
+一旦显式标出 memory tag，内存中一个 row-major 的 8×16 tile 就只是：
+
+```text
+S[(8, 16) : (16@m, 1@m)]
+```
+
+当 layout 描述的不是内存中的排列，而是*跨 thread 分布*的数据时，这些 tag 就开始发挥价值。
+以 `S[(8, 4, 2) : (4@laneid, 1@laneid, 1@reg)]` 为例：它不是指向线性内存，
+而是把行和列映射到 lane ID 以及每个 lane 的一个寄存器。这里的 `laneid` 表示 warp 内的 lane index，
+即 `thread_index % warp_size`。这正是你会在 {ref}`zh_chap_layout_generations` 中遇到的 tensor-core register fragment。
+
+下面的交互式可视化展示了 layout 如何把 tensor 元素分布到 warp lane 和 per-lane register 上，
+而不是把它们放在线性内存中。
+
+```{raw} html
+
+```
+*交互：一个位于 `@laneid` 和 `@reg` 上的 layout；点击一个单元格，查看哪个 lane/register 持有它。*
+
+## Distributed Layout
+
+named axes 之所以有用，是因为它们让我们能在系统的许多层级上统一描述 placement，
+甚至包括*跨整个 device* 的 placement。我们刚刚把它们用于单个 GPU 内部的 lane 和 register，
+但同一个思想也可以向外延伸：像 `@gpuid_x` 和 `@gpuid_y` 这样的轴可以说明数据位于 GPU mesh 的哪里，
+于是这个记法也能捕捉分布式训练和推理中出现的 sharding pattern。
+这些轴尚未捕捉到的一件事是 *replication*，也就是数据被复制到不止一个位置。
+因此我们加入记法 `R[n : stride]`，其中 `R` 标记 replicated dimension。
+例如，`R[2 : 1@gpuid_x]` 描述沿 `@gpuid_x` 轴的 replication。把二者合在一起，
+一个表达式就能同时把 tensor shard 到 2×2 GPU mesh 上，并沿一个轴复制它：
+
+```text
+S[(2, 4, 8) : (1@gpuid_y, 8@m, 1@m)] + R[2 : 1@gpuid_x]
+```
+
+下面的演示在一个小型 GPU mesh 上展示这种 partition-and-replication 组合模式。
+点击任意单元格，可以看到哪个 device 持有它；也可以观察 `@gpuid_x` replication 如何把相同副本放到配对 device 上。
+按钮可以在 fully-sharded、shard + replica 和 shard + offset layout 之间切换。
+
+```{raw} html
+
+```
+*交互：一个分布在 2×2 GPU mesh 上的 layout；点击一个单元格，查看哪些 device 持有它。*
+
+### Kernel 内部的 Replication Pattern：TMEM 中的 Scale Factor
+
+我们刚刚为 GPU mesh 引入的 replication dimension `R[...]`，并不只与多个 device 有关。
+同一个结构也能描述完全发生在单个 kernel 内部的事情：硬件把数据*跨 lane broadcast*。
+Blackwell 的 block-scaled MMA（{ref}`zh_chap_layout_generations`）就是一个很好的例子。
+它的 scale factor 位于 TMEM 中，其中一个 128-row scale vector 只存储在 **32 个 TMEM lane** 中：
+逻辑行 `r` 会去到 TMEM lane `r % 32`，而 `r // 32` 沿 column 方向展开。
+这 32 个已存储的 TMEM lane 随后会**沿 TMEM `TLane` 轴复制**，从 32 个扩展到 128 个 TMEM lane，
+让读取 warpgroup 中四个 warp 的每一个，都能在自己的 32-lane TMEM window 中找到一份副本。
+这是一种 `warpx4` broadcast，我们用 replication dimension 来书写它。读取本身由这些 warp 的 thread 执行：
+
+```text
+S[(32, …) : (1@TLane, …)] + R[4 : 32@TLane]
+```
+
+这会给出四个副本，副本之间相隔 32 个 TMEM lane：TMEM lane `l`、`l+32`、`l+64` 和 `l+96`
+都持有同一个 scale。和之前一样，replication dimension 不携带新数据；它只是说“同一个值，位于四个 TMEM-lane 位置上”，
+就像刚才 `@gpuid_x` 把一行 broadcast 到 GPU mesh 上一样。
+
+下面的交互式演示把两个步骤放在一起展示：先紧凑 pack 到 32 个 TMEM lane 中，然后通过 `warpx4`
+broadcast 到 128 个读取 lane。
+
+```{raw} html
+
+```
+*交互：点击一个 scale factor `SFA[m, sf]`；它会 pack 到 TMEM 的 lane `m mod 32`、column `(m // 32)·4 + sf`，
+然后沿 `TLane` 轴通过 `warpx4` broadcast 到四个 lane 副本（`l`、`l+32`、`l+64`、`l+96`），每个 warp 的 32-lane window 各一份。*
+
+每个 column 内部的 byte packing（`scale_vec` 的 1X/2X/4X 模式）以及 `cta_group::2` split
+会在 {ref}`zh_chap_layout_generations` 中介绍。
+
+已经熟悉 CuTe 的读者，可以把本章记法理解为它的一个 row-major 变体：
+我们在其上扩展了显式的 hardware-named axes，以及专门的 replication 结构。
+
+## Swizzle Layout
+
+本章最后一种 layout 是为了解决一个具体的硬件问题。GPU 上的 shared memory 被组织成多个 memory bank；
+当不同 lane 落到不同 bank 上时，访问最快。相反，如果多个 lane 访问的是*同一个* bank 内的不同地址，
+硬件别无选择，只能把它们串行化，于是我们就要付出 **bank conflict** 的代价。
+
+在 tensor 程序中，这很难避免，因为内存访问并不是纯线性顺序。处理矩阵时，我们经常需要读取同一个 tile 的行切片和列切片，
+这就产生了真实的张力：对按行访问高效的 layout，往往会让按列访问产生 bank conflict；偏向列的 layout 又会伤害行访问。
+**Swizzling** 正是为打破这种张力而设计的技术。
+
+swizzle 背后的想法是置换地址映射，通常做法是把 column index 与 row 做 XOR，
+让按行和按列访问最终都分散到多个 bank 上。它提供的 conflict-free 保证是有条件的：
+只对匹配的元素宽度、swizzle mode 和访问模式（也就是某个引擎的 descriptor 所期望的模式）成立，
+并不适用于任意元素宽度或对齐方式。
+
+下面第一个交互式演示把这一点具体化。点击一个 column index，观察每个元素落到哪个 bank 中：
+左侧朴素 row-major tile 中，一列会把全部八个元素汇入同一个 bank，因此 read 会串行化为八个 cycle；
+右侧的 XOR-swizzled layout 中，同一列会分散到八个不同 bank，只需一个 cycle 即可读取。
+
+```{raw} html
+
+```
+*交互：一个 8×8 tile；朴素 row-major 中按列访问会产生 bank conflict，XOR swizzle 后则 conflict-free。*
+
+这个小小的 8×8 例子抓住了核心思想，但真实 GPU 内存的 bank 数量远多于这个玩具图所暗示的数量。
+为了让 swizzling 在完整尺度上工作，我们不会把整个 tile 当成一个单体对象。相反，我们把内存切成小 segment，
+并在每个 segment 内应用 swizzle pattern。实践中最常见的情况是 `SWIZZLE_128B`，
+它围绕 128-byte segment 组织，使同样的 row/column-remapping 技巧能够自然适配 32-bank memory system。
+
+下面的交互式演示展示一个具体的硬件 swizzle：`SWIZZLE_128B`。这样在推广到多种格式之前，
+你可以先看到逐 segment 重复的 pattern。
+
+```{raw} html
+
+```
+*交互：128-byte segment 内部的 `SWIZZLE_128B` pattern；逐步查看 read cycle，观察 `physical_sector = logical_sector XOR row` 如何把每一列分散到不同 bank。*
+
+同一个想法可以扩展到这个 128-byte 情况之外。为了简化可视化，接下来我们会用一个单色块表示一个 segment，
+而不是画出每个 bank。一般来说，硬件会定义一个小的重复 **atom**，permutation 会应用在这个 atom 上；
+不同 swizzle mode 会选择不同 atom 大小。`SWIZZLE_128B` 使用 8 × 128 B atom，
+`SWIZZLE_64B` 使用 8 × 64 B atom，`SWIZZLE_32B` 使用 8 × 32 B atom；
+随后整个 tile 会由当前使用的 atom 平铺而成。
+
+最后一个交互式演示允许你在这些格式之间切换（包括 16 B interleaved mode）、选择数据类型，
+并悬停任意单元格，直接检查一个 atom 内部的元素排列。对于推理某条 load/store 指令期望哪种 swizzle，
+这正是合适的细节层级。
+
+```{raw} html
+
+```
+*交互：选择一种 swizzle format（以及数据类型），查看它的 atom shape（8 × N B）；悬停某个单元格，查看其中元素如何被置换。*
+
+应该选择哪种模式？经验法则是优先选择 tile 能填满的*最大* atom。一个 N-byte atom 要求 tile 的 contiguous dimension
+至少为 N 字节，并且是 N 的倍数。因此，`SWIZZLE_128B` 只在一行跨度至少为 128 字节，
+也就是 64 个 `float16` 元素时适用。如果能适配，它就是默认选择，因为它的 8 × 128 B atom 覆盖完整的 128-byte bank line，
+从而一次把一列分散到全部 32 个 bank 上，在 fp16 中可以同时对 8 行和 8 列提供 conflict-free 访问。
+不过，当问题 shape 迫使 contiguous dimension 较小时，tile 就无法再填满 128 B atom；
+此时你会降到 `SWIZZLE_64B` 或 `SWIZZLE_32B`，也就是该行仍然能覆盖的最大 atom。
+
+你永远不需要手工算出这些置换后的地址；但有必要精确说明 swizzle 与 `S[...]` 记法之间的关系：
+它*不是*那个 affine map 的一部分，而是叠加在其上的一个独立、非仿射层。`S[...]` layout 把元素放到线性内存
+（`@m`）地址上，随后 swizzle 置换该地址。在 TIRx layout API 中，这写作
+`ComposeLayout(swizzle, tile)`（{ref}`zh_chap_tirx_layout_api`）。你的任务只是为每个会接触这个 tile 的 op
+选择一种一致的模式，然后让 composed layout 完成其余工作。
+
+硬件填充的也是这个 composed layout，这正是 swizzling 和 tiling 汇合的地方。TMA descriptor 是多维的，
+所以单个三维 box 可以同时描述 tile 的 atom tiling 以及每个 atom 内的 swizzle；
+一次 TMA load 随后会按 atom 布置 tile，并在写入 shared memory 时完成 swizzle（{ref}`zh_chap_tma`），
+不需要单独的 swizzling pass。每个引擎要求*哪一种* swizzle 是 generation-specific 的，这正是下一章的主题。
diff --git a/zh/chapter_flash_attention/index.md b/zh/chapter_flash_attention/index.md
new file mode 100644
index 00000000..dc5773b5
--- /dev/null
+++ b/zh/chapter_flash_attention/index.md
@@ -0,0 +1,607 @@
+(zh_chap_flash_attention)=
+# Flash Attention 4
+
+:::{admonition} 概览
+:class: overview
+
+- Attention 会运行两个 MMA，并在它们之间插入 softmax，因此它不能像 GEMM 那样简单地重复同一个 MMA。
+- 这个 kernel 把 Part I 中的硬件原语（TMA、`tcgen05`、TMEM、barrier）和 Part III 中的 GEMM 技术组合起来：warp 角色分工、online-softmax rescaling、causal masking 和 GQA。
+:::
+
+Attention 是决定 transformer 能不能高效运行的核心 kernel，也是前面构建的所有机制终于汇合的地方。我们为 GEMM 组装过的每个组件都会延续到这里：TMA tile 搬运、`tcgen05` MMA、TMEM、warpgroup register tile，以及显式 barrier。
+
+难点在于，attention 不是简单重复一个 MMA。它是两个 MMA，中间夹着真正的计算：online softmax、causal masking，以及把早期和后续 block 维持在同一尺度下的 rescaling。
+
+新的复杂度正藏在这个中间阶段。普通矩阵乘只需要往 accumulator 里加；attention 则必须在新的 key 和 value 持续流入时，回头重访并重新缩放已经算过的结果。softmax 本身也运行在两个 Tensor Core MMA 之间的 CUDA core 上，因此指数运算和逐行规约直接落在关键路径上。
+
+这就是为什么 attention 优化很大一部分其实是 softmax 优化：重写 `exp`，并把 softmax 与 MMA 重叠起来，而不是让 MMA 停下来等待它。
+
+本章的目标不是从零重新推导 Flash Attention。我们会保留足够的算法视角，让 kernel 能读懂；然后把注意力放在真正新的部分：这个算法如何落到 TIRx 上。
+
+最清晰的入口，是跟随一个 tile 在 kernel 中的流动。`Q`、`K` 和 `V` 作为输入 tile 进入 kernel，从 GMEM 加载到 SMEM。score MMA 将 `Q` 和 `K` 相乘，得到 TMEM 中的 score tile `S`。softmax 把 `S` 变成 numerator tile `P`，value MMA 再组合 `P` 和 `V` 来更新输出 accumulator `O`。
+
+到这里为止，它看起来像是两个矩阵乘粘在一起。但它有一个 GEMM 不必处理的转折：每当运行中的 softmax 最大值发生变化，已经累计到 `O` 里的结果就突然处在了错误的尺度上。它必须先被重新缩放，下一次 value MMA 才能安全地加进去。下面几节会先追踪这条路径，然后再展示 TIRx 如何把每个阶段交给对应的 warpgroup，并把这些阶段串起来。
+
+## 算法形状
+
+在把 tile 放进内存之前，我们需要先看清这些 tile 服务的算法。对于一个 query block，Flash Attention 计算：
+
+$$O = \text{softmax}(QK^{\top} / \sqrt{d})V$$
+
+按字面理解，这个公式会先形成完整的 score 矩阵 `S = QKᵀ`，对它做 softmax，然后再乘以 `V`。这恰恰是我们不能采用的做法，因为完整的 `S` 非常巨大。seq=4096 时，每个 head 大约有 16M 个元素，fp32 下约 64 MB，远远超过 SMEM 或单个 128×512 TMEM 区域的容量。片上根本没有地方放它。Flash Attention 的答案是完全不物化 `S`。它改为按 block 流式读取 `K/V`，并维护三个逐行运行状态，用来概括目前为止看过的全部内容：
+
+- `row_max`：到目前为止看到的最大 score。
+- `row_sum`：softmax 分母的运行和。
+- `O`：运行中的输出 accumulator。
+
+流式更新负责让这些状态在新 block 到来时仍然正确。微妙之处在于，每处理一个 block，运行中的最大值都可能升高；一旦升高，所有在旧最大值尺度下计算出的内容都落在了错误尺度上。因此在加入新的贡献之前，我们先把旧状态拉回到新的尺度：
+
+```text
+S = Q_block @ K_block.T
+m_new = max(row_max, rowmax(S))
+scale = exp((row_max - m_new) / sqrt(d))
+P = exp((S - m_new) / sqrt(d))
+row_sum = row_sum * scale + rowsum(P)
+O = O * scale + P @ V_block
+row_max = m_new
+```
+
+单个 `scale` 因子在这里一鱼两吃：它同时重新缩放运行中的分母和运行中的输出，让早先 block 与后续 block 的贡献最终都落在同一个尺度上。
+
+上面的伪代码使用自然 `exp`，并显式写出 `/sqrt(d)`，因为这样最好读；但 kernel 采用了更便宜的路径。它把 `1/sqrt(d)` 和 `log2(e)` 合并成一个常量 `scale_log2 = log2(e)/sqrt(d)`，然后用硬件 `exp2` 在原始 score 上计算所有指数，利用恒等式 `exp(x/sqrt(d)) = exp2(x · scale_log2)`。动机很简单：在这类硬件上，`exp2` 比自然 `exp` 更快。
+
+继续往下之前，有一点值得钉牢：这里的 `P` 不是最终归一化后的 attention 矩阵。它只是当前 K/V block 的 softmax 分子。归一化被刻意推迟，只有最后一个 block 处理完之后，kernel 才写出 `O / row_sum`。
+
+对 TIRx 来说，知道算法算什么只是一半；另一半是 kernel 运行时每个 tile 住在哪里，因为这决定了 layout 和 barrier 代码。`S`、`P` 和 `O` 都是 tile 值，而且各自有自己的家：
+
+- `S` 是 score tile。score MMA 将它写入 TMEM。
+- `P` 是 softmax numerator tile。softmax 从 TMEM 把 `S` 读入寄存器，计算 `P = exp((S - m_new) / sqrt(d))`，再把 `P` 写回 TMEM。
+- `O` 是输出 accumulator tile。value MMA 从 TMEM 读取 `P`、从 SMEM 读取 `V`，然后累计到 TMEM 中的 `O`。
+
+前面提到的 rescale 也是一个 tile 操作，而不是一段标量记账：当 `row_max` 变化时，旧的 `O` 会从 TMEM 读出，在寄存器中相乘，再写回 TMEM，然后下一次 value MMA 才会继续累加。后续每一节都会沿着同样的结构展开：tile 的位置、硬件路径，以及证明下一个 consumer 可以运行的 barrier。
+
+## Tile-Primitive 图
+
+有了运行状态及其位置之后，我们可以把算法展开成一串具体的 tile 移动。对于一个 K/V block，kernel 从上到下走过这条 tile 路径：
+
+```text
+Q, K, V 位于 GMEM
+  -> Q, K, V 位于 SMEM        由 TMA 加载完成
+  -> S 位于 TMEM              由分数 MMA 完成：QK^T
+  -> P 位于 TMEM              由 softmax 分子计算完成：TMEM -> RF -> TMEM
+  -> O 位于 TMEM              由值 MMA 完成：P V
+  -> O 位于 GMEM              由归一化、SMEM 暂存和 TMA 存储完成
+```
+
+它和 GEMM 的区别归结为一行。GEMM 是重复一条 MMA 链；FA4 有两个 MMA 阶段，中间坐着 softmax。后面几乎所有复杂度，都是这个额外阶段带来的后果。
+
+如果把上面的短路径展开成显式的 producer-consumer 边，就得到完整图：
+
+| 阶段 | Tile 移动或计算 | TIRx primitive | 硬件路径 |
+|-------|------------------|----------------|----------|
+| 加载 Q/K/V | GMEM tile -> SMEM tile | `Tx.copy_async(..., dispatch="tma")` | TMA load |
+| Score MMA | SMEM 中的 Q 与 K -> TMEM 中的 score tile `S` | `Tx.warp.gemm_async(..., dispatch="tcgen05")` | `tcgen05.mma` |
+| Softmax 读取 | TMEM 中的 `S` -> warpgroup register tile | `Tx.wg.copy_async(reg, tmem)` | `tcgen05.ld` |
+| Softmax 写入 | 寄存器中的 numerator tile `P` -> fp16 TMEM view | `Tx.copy_async(tmem_as_f16, reg)` | TMEM store，然后 `tcgen05.wait.st()` |
+| Value MMA | TMEM 中的 `P` 与 SMEM 中的 V -> TMEM 中的输出 accumulator `O` | `Tx.warp.gemm_async(..., dispatch="tcgen05")` | 带 TMEM operand 的 `tcgen05.mma` |
+| Correction | TMEM 中的 `O` -> 寄存器 -> TMEM 中的 `O` | TMEM readback、寄存器乘法、TMEM store | `tcgen05.ld` / TMEM store |
+| Epilogue | TMEM 中的最终 `O` -> 寄存器 -> SMEM -> GMEM | TMEM readback、`Tx.copy`、TMA store | `tcgen05.ld` + TMA store |
+
+新增的行是 softmax 和 correction。二者都会增加 TMEM -> register -> TMEM 流量，也都会在 score MMA 与 value MMA 之间制造额外的交接。
+
+**试着让你的 agent 做一遍**：让它只追踪上面的短路径。对每条箭头，指出 producer 阶段、consumer 阶段、源 tile、目标 tile 和硬件路径。然后再问哪些箭头在 GEMM 章节里并不存在。
+
+## Warp 角色与 Scope
+
+数据路径理清之后，自然的下一个问题是：每个阶段到底由谁来跑。这里每个 CTA 有 4 个 warpgroup，总共 512 个线程；它们不是按接触的数据划分，而是按 warpgroup 执行的工作类型划分：
+
+- WG3 驱动硬件引擎：TMA load、MMA 和 TMA store。
+- WG0、WG1、WG2 执行这些引擎调用之间的寄存器重计算：softmax、correction 和 epilogue。
+
+精确的角色表如下：
+
+| Owner | 角色 | 做什么 |
+|-------|------|--------|
+| WG3, warp 1 | TMA load | 从 GMEM 把 Q、K、V tile 加载到 SMEM |
+| WG3, warp 0 | MMA | 发出 score MMA 和 value MMA |
+| WG3, warp 2 | TMA store | 把最终 O tile 从 SMEM 存回 GMEM |
+| WG0 | Q stage 0 的 softmax | 从 TMEM 读取 S，计算 P，把 P 写回 TMEM |
+| WG1 | Q stage 1 的 softmax | 对第二个 Q pipeline stage 执行同样工作 |
+| WG2 | Correction 与 epilogue | 重新缩放 TMEM 中的 O，做归一化，暂存输出 |
+
+很容易把“两个 Q stage”误读成两个 attention head，但它们不是。它们只是 Q pipeline 中的两个槽位：WG0 拥有一个，WG1 拥有另一个，因此两个 Q tile 可以同时在路上。这就是 softmax 工作出现两份的原因，一份在 WG0，一份在 WG1。
+
+代码用符号坐标选出这些角色：
+
+```python
+wg_id = T.warpgroup_id([4])
+warp_id = T.warp_id_in_wg([4])
+```
+
+读 kernel 时，先找角色分支。它会告诉你分支内部每个 tile primitive 归哪个团队所有。
+
+- WG3 warp 1 启动 TMA load 命令。一个被选出的 lane 发出 copy，TMA 引擎移动 tile。
+- WG3 warp 0 发出 `tcgen05.mma` 指令。
+- WG0 和 WG1 在完整 warpgroup scope 下运行 softmax。
+- WG2 在完整 warpgroup scope 下执行 correction 与 epilogue。
+
+一个不对称性最终塑造了整个 barrier 图：每个 MMA，无论 score 还是 value，都只由 WG3 warp 0 发出。WG0 和 WG1 从不发出 MMA。它们只消费 score tile、运行 softmax，并把 `P` 写回 TMEM。
+
+正是这种分离，让 softmax 周围必须有 barrier。`s_ready` 把 score tile 从 MMA warp 交给 softmax；`p_o_rescale` 则交付 `P`，以及一个对 value MMA 来说安全的 `O` 槽位：要么已经完成 rescale，要么因为不需要 rescale 而被释放。后面几节我们会反复回到这两个名字。
+
+## 阅读代码片段
+
+本章的代码片段摘自 [`flash_attention4.py`](https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/attention/flash_attention4.py)，所以它们不可避免地会引用一些我们没有完整复现的 kernel 内部名称。自解释的名称（`wg_id`、`warp_id`、`BLK_M`/`BLK_N`、`HEAD_DIM`、`kv_stage`、各个 `SMEM_PIPE_DEPTH_*` / `TMEM_PIPE_DEPTH` 深度、`should_accumulate`，以及这里为 1 的 `CTA_GROUP`）会在第一次相关时引入。其余名字先在下表给出一行释义，这样当代码片段突然把一个陌生名字放到你面前时，你有地方可查：
+
+| 名称 | 含义 |
+|------|------|
+| `q_stage`, `i_q` | Q pipeline stage，取 0 或 1，也就是哪个 Q tile 槽位（`SMEM_PIPE_DEPTH_Q = 2`）。在 WG0/WG1 softmax 内部，warpgroup 自己的 `wg_id`（0 或 1）就是同一个 stage 索引，因此 `S_region[q_stage]`、`P_region[wg_id]` 和 `O_region[i_q]` 都选择同一个 Q stage |
+| `MMA_N` | TMEM 列中的 score/output tile 宽度（128） |
+| `MMA_K` | `P`/`V` 列方向的 MMA inner-K 步长（16）；`K_SPLIT = 6 * MMA_K = 96` |
+| `K_SPLIT` | value-MMA 调度的切分点（见“两段 MMA 阶段”）；第一段 value MMA 覆盖列 `0:K_SPLIT`（`6 * MMA_K = 96`） |
+| `should_rescale` | WG2 的逐行标志：旧 `O` 是否需要在下一次 value MMA 前 rescale（通过 `any_sync` 在 warpgroup 内规约） |
+| `rescale_threshold` | 跳过小幅 row-max 变化的阈值；当前 kernel 使用 `8.0`，跳过 rescale 时会把 `acc_scale` 精确设为 `1.0` |
+| `scale_log2` | log2 单位下的 softmax scale，即 `log2(e)/√d`，因此 `P = exp2((S - m) · scale_log2)` |
+| `acc_scale` | softmax 通过 SMEM mailbox 传给 WG2 的逐行 rescale 因子 |
+| `chunk_start`/`chunk_end`, `p_start`/`p_end` | 正在读取/写入的 32 宽 softmax chunk 的列范围 |
+
+## 两段 MMA 阶段
+
+对每个流式 K/V tile，Flash Attention 都会运行两个 MMA 阶段，并由 softmax 把它们桥接起来：
+
+```text
+Q, K -> 分数 MMA -> S
+S    -> softmax   -> P
+P, V -> 值 MMA   -> O
+```
+
+可以把它看成三个 producer 串成的一条 pipeline。第一个 MMA 产生 attention score `S`，softmax 把 `S` 转成 numerator `P`，第二个 MMA 消费 `P` 来更新输出 accumulator `O`。按 `row_sum` 归一化会一直推迟到 epilogue，等每个 K/V tile 都贡献完之后再做。
+
+下面每个 tile op 都会使用 GEMM 步骤中同样的 **scope / layout / dispatch** 卡片，只额外加一行 **handoff**，用来指出把 tile 交给下一个角色的 barrier。
+
+计算代码从不直接使用裸 TMEM 列号。kernel 会把唯一的 TMEM 分配切成按 stage 索引的视图（`S_region`、`P_region`、`O_region`），然后用 pipeline stage 访问它们（`S_region[q_stage]`、`O_region[i_q]`、`P_region[i_q, 0:K_SPLIT]`）。这些视图由 [TMEM 布局与复用](#tmem-布局与复用) 一节中的 `T.TMEMStages` 定义；现在只要把每个 region 理解为同一块物理 TMEM 的一个具名切片就够了。
+
+### Score MMA
+
+两个阶段中的第一个是 score MMA，也就是打开每个 K/V iteration 的矩阵乘。它计算：
+
+$$S = Q_{\text{block}}K_{\text{block}}^{\top}$$
+
+并把 `128 x 128` score tile 写入 TMEM：
+
+```python
+Tx.warp.gemm_async(
+    S_region[q_stage],
+    Q_smem[q_stage, 0:BLK_M, 0:HEAD_DIM],
+    K_smem[kv_stage, 0:BLK_N, 0:HEAD_DIM],
+    dispatch="tcgen05",
+    cta_group=CTA_GROUP,
+)
+if T.ptx.elect_sync():
+    s_ready.arrive(q_stage)
+```
+
+我们可以问 GEMM 章节对每个 tile op 问过的同样四个问题：谁运行它、tile 住在哪里、如何 dispatch，以及如何交接：
+
+> **Tile-primitive 解读：Score MMA**
+> - Scope：WG3 warp 0 发出它；一个 elected lane 到达 `s_ready`。
+> - Layout：Q、K 在 SMEM → TMEM 中的 `S`（`S_region[q_stage]`）。
+> - Dispatch：`tcgen05`。
+> - Handoff：`s_ready`（→ softmax）。
+
+被选中的单个线程到达 `s_ready`，就是整个交接。它宣告这个 score tile 已经完成，softmax warpgroup 现在可以读取它了。
+
+### 两个 MMA 之间的 Softmax
+
+两个 MMA 之间坐着 softmax，它把 score tile `S` 转成 numerator tile `P`。它的解读卡如下：
+
+> **Tile-primitive 解读：Softmax**
+> - Scope：WG0（Q stage 0）/ WG1（Q stage 1），完整 warpgroup。
+> - Layout：TMEM 中的 `S` → 寄存器 → fp16 TMEM 中的 `P`（`P_region[wg_id]`）。
+> - Dispatch：用 `tcgen05.ld` 读取，用 TMEM store 写入；中间在寄存器中做逐行计算。
+> - Handoff：等待 `s_ready`；到达 `p_o_rescale`（前 96 列）和 `p_ready_2`（最后 32 列）。
+
+这个阶段完全没有 GEMM 对应物。WG0/WG1 等待 `s_ready` 上的 score tile 到达，然后每次按寄存器大小的 chunk 从 TMEM 读出：
+
+```python
+Tx.copy_async(
+    s_chunk[:, chunk_start : chunk_end],
+    S_region[wg_id, chunk_start : chunk_end],
+)
+```
+
+这是 warpgroup scope 下的一次 TMEM-to-register tile 读取。score 进入寄存器后，softmax warpgroup 按顺序做三件事：
+
+1. 计算 row max 和 row sum；
+2. 计算 softmax numerator tile `P`；
+3. 以 fp16 把 `P` 写回 TMEM。
+
+最后一步形如：
+
+```python
+Tx.copy_async(
+    P_region[wg_id, p_start : p_end],
+    p_chunk[:, p_start : p_end],
+)
+```
+
+为什么已经在寄存器里算完了，还要把 `P` 写回 TMEM？因为 value MMA 需要把 `P` 当作一个 *tile operand*，而 MMA 不能把分散在每个线程里的标量寄存器直接当成矩阵来读。在这个 kernel 中，MMA 可读的 `P` 形态就是 `P_region`，它是 fp16 TMEM alias `tmem_as_f16` 上的一个视图。所以这次写回不是多余搬运；它是在把 `P` 放进下一个 MMA 唯一能消费的形态。
+
+### Value MMA
+
+第二个阶段，也是每个 K/V iteration 的收尾阶段，是 value MMA。它计算：
+
+$$O = O + P_{\text{block}}V_{\text{block}}$$
+
+这个 MMA 运行时，`O` 已经被放进了当前 K/V block 需要的正确状态：第一个 block 上完成初始化，后续 block 上完成 rescale。因此 MMA 只需要累加。它和 GEMM 的区别在于 operand 的位置：A operand 是 TMEM 中的 `P`，B operand 是 SMEM 中的 `V`，accumulator `O` 也在 TMEM 中：
+
+```python
+# 第一段 sub-MMA：列 0:K_SPLIT（P 的前 96 列 / V 的对应行）。
+Tx.warp.gemm_async(
+    O_region[i_q],
+    P_region[i_q, 0:K_SPLIT],
+    V_smem[kv_stage, 0:K_SPLIT, 0:HEAD_DIM],
+    transB=True,
+    accum=should_accumulate,
+    dispatch="tcgen05",
+    cta_group=CTA_GROUP,
+)
+# 第二段 sub-MMA 形式相同，accum=True，由 p_ready_2 gate 控制，
+# 覆盖剩余列 K_SPLIT:BLK_N。
+```
+
+> **Tile-primitive 解读：Value MMA**
+> - Scope：WG3 warp 0。
+> - Layout：TMEM 中的 `P` + SMEM 中的 V → TMEM 中的 `O`（`O_region[i_q]`）。
+> - Dispatch：带 TMEM operand 的 `tcgen05`。
+> - Handoff：等待 `p_o_rescale`、`p_ready_2`、`kv_load.full`；到达 `o_ready`（→ epilogue）。
+
+这种 operand 放置是两个 MMA 在硬件上的差异：
+
+- Score MMA 从 SMEM 读取两个 operand：Q 和 K。
+- Value MMA 从 TMEM 读取一个 operand：`P`。
+- Value MMA 从 SMEM 读取另一个 operand：V。
+- 结果累计到 TMEM 中的 `O`。
+
+`accum=should_accumulate` 标志实现了算法中的“初始化还是相加”选择：query block 的第一个 K/V tile 上为 false，之后每个 tile 上为 true。
+
+你可能还会注意到，value MMA 不是一次性跑完，而是切成 `96 + 32` 的调度：
+
+1. Softmax 以四个 32 列 chunk 写入 `P`。
+2. 前三个 chunk 一就绪，value MMA 就开始处理 `P` 的前 96 列和 `V` 的匹配行。
+3. 最后 32 列等待 `p_ready_2`。
+4. 第二个 MMA 消费最后这个 chunk 并完成 tile。
+
+这样切分是为了让 Tensor Core 保持忙碌。如果 value MMA 作为单条指令运行，整个阶段都要等四个 32 列 `P` chunk 全部完成指数计算并写回后才能开始。先对前三个 chunk 发起 MMA，可以把最后一个 chunk 的 `exp` 和 TMEM 写入，与已经在飞行中的 96 宽 MMA 重叠起来，把本来会空转的时间变成有用工作。
+
+## TMEM 布局与复用
+
+`S`、`P` 和 `O` 都必须共享一个 `128 x 512` TMEM 分配；它们被打包进同一块空间的方式，正是这个 kernel 中 barrier 与 layout 不可分割的原因：
+
+下图直接展示了这种打包：score slot、numerator slot 和 output slot 全部共享同一块 TMEM 分配，因此 barrier 协议负责让这种复用合法。
+
+![TMEM 布局](../img/tmem_layout_v3.png)
+
+可以把图读成一组 tile 槽位：
+
+- Score slot 保存 `S = QK^T`。
+- Numerator slot 保存 softmax 指数化后的 `P` tile。
+- Output slot 保存 fp32 `O` accumulator。
+
+它们不是彼此独立的 buffer，而是同一块分配中的区域；这种共享不是风格选择，而是容量限制迫出来的。Q pipeline 深度为 2 时，两个 `S` slot（2 × MMA_N = 256 列）和两个 `O` slot（2 × MMA_N = 256 列）已经占满了全部 512 个 fp32 列。没有剩余空间给 `P`，所以 `P` 只能通过更窄的 fp16 view alias 到同一批字节上。安全性的唯一来源，是每个 region 都严格在前一个 consumer 完成之后才复用；这个时序正是 barrier 保证的。因此在 FA4 里，barrier 不只是调度机制；它们本身就是 layout 合法性的条件。
+
+aliasing 技巧通过 `T.TMEMPool` 搭起来。kernel 先拿一个 fp32 view（`tmem`）用于 score 和 output accumulator，然后把 pool base 倒回 0，再在同一批物理字节上拿第二个 fp16 view（`tmem_as_f16`）：
+
+```python
+tmem_pool = T.TMEMPool(pool, total_cols=N_COLS_TMEM, cta_group=CTA_GROUP, tmem_addr=tmem_addr)
+tmem = tmem_pool.alloc((128, N_COLS_TMEM), "float32")
+tmem_pool.move_base_to(0)
+tmem_as_f16 = tmem_pool.alloc((128, N_COLS_TMEM * 2), "float16")
+tmem_pool.commit()
+```
+
+由于 fp16 元素宽度只有 fp32 的一半，fp16 view 会在同一批字节上暴露两倍数量的可索引列；`P` 正是住在这块空间里，而 fp32 layout 没有余量容纳它。拿到两个 view 后，kernel 使用 `T.TMEMStages` 把 `S`、`P` 和 `O` 槽位切成 staged region，这样计算代码就可以按 pipeline stage 索引，而不必直接操作裸列号：
+
+```python
+S_region = T.TMEMStages(tmem,        col_start=0,                       width=MMA_N, stages=SMEM_PIPE_DEPTH_Q, stride=MMA_N)
+O_region = T.TMEMStages(tmem,        col_start=MMA_N * SMEM_PIPE_DEPTH_Q, width=MMA_N, stages=SMEM_PIPE_DEPTH_Q, stride=MMA_N)
+P_region = T.TMEMStages(tmem_as_f16, col_start=MMA_N,                   width=BLK_N, stages=SMEM_PIPE_DEPTH_Q, stride=MMA_N * 2)
+```
+
+`P_region` stride 里的 `* 2`，是 aliasing 在代码中显形的一个地方。`S_region` 和 `O_region` 用 fp32 `tmem` 列计数，而 `P_region` 用 fp16 `tmem_as_f16` 列计数；fp16 列只有一半宽，所以 stage 到 stage 的移动需要双倍 stride，才能落在相同的物理字节上。不过 region 一旦定义好，计算代码就保持干净：写 `S_region[q_stage]`，读 `S_region[wg_id, ...]`，写 `P_region[wg_id, ...]`，累计到 `O_region[i_q]`，完全不用碰裸列号。
+
+**试着让你的 agent 做一遍**：让它解释这个 FA4 kernel 里的 fp32（`tmem`）和 fp16（`tmem_as_f16`）两个 view。哪些物理 TMEM 区域保存 `S`、`P` 和 `O`？为什么 `P_region` 的 stride 使用 `MMA_N * 2`？复用问题先留到下一节：看完 barrier 表之后，再检查每个 region 复用前必须等哪些 consumer 完成。
+
+## Barrier 如何连接各个角色
+
+这是整个 kernel 最难的部分，所以值得循序渐进。先从沿主计算路径移动数据的少数 barrier 入手，把其他部分都当作稍后可查的 bookkeeping。数据就绪交接包括：
+
+| Handoff | 含义 |
+|---------|------|
+| TMA load -> score/value MMA | Q、K 或 V 已经到达 SMEM，可以供 MMA 使用 |
+| score MMA -> softmax | `S` 已经在 TMEM 中就绪 |
+| softmax/correction -> value MMA | `P` 已经在 TMEM 中就绪，并且 `O` 可以安全累计 |
+| value MMA -> epilogue | 最终 `O` 已经在 TMEM 中就绪 |
+| epilogue -> TMA store | `O_smem` 已经可以存回 |
+
+不在这张表里的东西都是 pipeline bookkeeping：释放某个 SMEM、TMEM 或 staging buffer，让另一个角色可以复用它。有用的是，每个 barrier 不管携带的是数据还是 bookkeeping，都能用同一种方式阅读：一次 tile handoff。你问谁生产了数据、谁消费它，以及双方完成后哪个 buffer 变得可复用。
+
+下一张图把这些交接压缩成两个 MMA 阶段的精确 readiness gate：score MMA 等什么，value MMA 累计前又必须等什么。
+
+![Flash Attention 4 MMA 输入门控](../img/flash_attention_main_handoff.png)
+
+请把这张图读成一组正确性 gate，而不是调度表。它回答“这个 MMA 发射前必须满足什么”，但不说明具体时序。score MMA 等待 SMEM 中的 Q 和 K，然后产生 `S`。value MMA 同时等待三件事：SMEM 中的 V、softmax 产生的 `P` tile，以及一个由 WG2 释放或 rescale 完成的 `O` 槽位。softmax 到 value 的 gate 会分裂成两段，原因我们已经见过：`P` 的前 96 列就绪后 value MMA 就可以开始，`p_ready_2` 再释放最后 32 列。
+
+有一个 handoff 不符合 tile-readiness 的模板：softmax 到 correction 的边。它不是传递 tile，而是通过一个单槽 SMEM mailbox，把一个标量（K/V loop 中的 `acc_scale`，或 epilogue 中的最终 `row_sum`）传给 WG2。由于这个槽位每次 iteration 都会复用，因此必须由一对 `full`/`empty` barrier 保护：
+
+下图放大了这个 mailbox handshake，因此这对 barrier 应该被读作一个标量 producer-consumer 通道，而不是 tile-ready gate。
+
+![Flash Attention 4 softmax 缩放槽握手](../img/flash_attention_softmax_correction.png)
+
+把 `softmax_corr.full` 和 `softmax_corr.empty` 读成一对 producer-consumer barrier：
+
+1. Softmax 在复用 scale/sum 槽位前等待 `softmax_corr.empty`。
+2. Softmax 把 `acc_scale` 或最终 `row_sum` 写入这个槽位。
+3. Softmax 到达 `softmax_corr.full`。
+4. WG2 等待 `softmax_corr.full`，然后读取这个槽位。
+5. WG2 到达 `softmax_corr.empty`。
+6. softmax warpgroup 可以在下一阶段复用这个槽位。
+
+要特别小心 `softmax_corr.empty` 表示什么、又不表示什么。它只表示 WG2 已经消费了 scale/sum 槽位。它不说明 `P` 是否就绪，更绝对不是允许 value MMA 开始的 gate。真正的 gate 是 `p_o_rescale`，它在 `P` 的前 96 列写好、且 `O` 槽位可以安全累计时触发。混淆这两者，是产生错误结果的经典来源。
+
+掌握主路径后，完整 barrier 列表就可以作为参考：
+
+| Barrier | Producer -> consumer | 什么变得安全 |
+|---------|----------------------|----------------|
+| `q_load.full` | TMA load -> score MMA | Q SMEM tile 可以供 MMA 使用 |
+| `q_load.empty` | 这个 Q stage 的所有 score MMA -> TMA load | Q SMEM stage 可以复用于下一个任务 |
+| `kv_load.full` | TMA load -> score/value MMA | K 或 V SMEM tile 可以供 MMA 使用 |
+| `kv_load.empty` | score/value MMA -> TMA load | K/V SMEM stage 可以复用 |
+| `s_ready` | score MMA -> softmax | S TMEM tile 可以读取 |
+| `p_o_rescale` | softmax + WG2 -> value MMA | P 的前 96 列已经在 TMEM 中，且 O 槽位可以供 value MMA 安全累计 |
+| `p_ready_2` | softmax -> value MMA | P 的最后四分之一已经在 TMEM 中 |
+| `o_ready` | value MMA -> epilogue | 最终 O accumulator 已经就绪 |
+| `softmax_corr.full` | softmax -> WG2 | `acc_scale` 或最终 `row_sum` 已在 SMEM mailbox 中就绪 |
+| `softmax_corr.empty` | WG2 -> softmax | WG2 读取后，同一个 SMEM mailbox 槽位可以复用 |
+| `corr_epi.full` | epilogue -> TMA store | O_smem 已经可以存储 |
+| `corr_epi.empty` | TMA store -> epilogue | O_smem stage 可以复用 |
+
+和 GEMM 一样，你可以从信号生产者预测 barrier 类型：
+
+- TMA load 使用 `TMABar`，因为 TMA 引擎会按字节数统计自己的完成。
+- MMA 完成使用 `TCGen05Bar`，因为 `tcgen05.commit` 会 signal completion group。
+- 纯线程到线程的 handoff 使用 `MBarrier`，参与线程显式 arrive。
+
+softmax 到 value 的分裂式 handoff 值得仔细看。它使用两个 gate：
+
+- `p_o_rescale` 在 `P` 的前 96 列写好、且 `O` tile 可以安全累计后，允许 value MMA 开始。
+- `p_ready_2` 释放 `P` 的最后 32 列，对应上一节中的 `96 + 32` value-MMA 调度。
+
+第一个 K/V block 是简单情况。WG2 会预先 arrive `p_o_rescale`，因为还没有旧的 `O` tile 需要 rescale。
+
+后续 block 必须更谨慎。WG2 只有在跳过一次不必要的 rescale，或完成旧 `O` 的 rescale 之后，才会到达 `p_o_rescale`。跳过测试刻意保守：softmax 计算 log2 缩放后的 delta `(m_old - m_new) * scale_log2`；如果这个值仍然高于 `-rescale_threshold`，说明新的 max 变化还没大到值得 rescale，kernel 就保持旧 max，并把 `acc_scale` 精确设为 1.0。只有更大的 max 跳变才会走 `exp2` 路径，并要求 WG2 rescale `O`。
+
+随后 WG2 用 `any_sync` 在 warpgroup 内规约 `should_rescale`。如果没有任何行需要更新，它就保持 `O` 不动。这个跳过很重要，因为 rescale `O` 是覆盖整个 accumulator 的一次 TMEM -> RF -> TMEM read-modify-write；当阈值逻辑已经把 `acc_scale` 维持为 1.0 时，做这件事就是纯浪费。
+
+注意所有新增 barrier 都聚在同一个地方。`s_ready`、`p_o_rescale`、`p_ready_2`，以及 softmax/correction 这对 barrier，全都围绕 softmax。它们存在只有一个原因：score MMA 和 value MMA 不再相邻。寄存器计算、TMEM 重写和输出 rescale 现在插在两者之间，每一步都需要自己的 handoff。
+
+**试着让你的 agent 做一遍**：让它追踪一个 K/V block 经过 `s_ready`、`p_o_rescale`、`p_ready_2` 和 `o_ready`。对每个 barrier，问谁等待、谁 arrive、哪个 tile 变得可读，以及之后哪个存储可以复用。
+
+## Pipeline 结构
+
+barrier 告诉我们某个角色消费 tile 之前什么必须 *就绪*。但它们没有告诉我们实际上哪些东西会 *并发* 运行；这正是现在要讨论的问题。两者确实不同：一个正确性 gate 可能在 producer 真正运行之前很久或之后很久才满足。
+
+这里不存在单一的 pipeline 深度，因为不同 tile stream 以不同速率移动。因此 kernel 为每条 stream 保留独立的 ring：
+
+- Q pipeline 深度 2：一个 CTA 同时处理两个 Q stage。WG0 处理一个 stage，WG1 处理另一个。
+- KV pipeline 深度 3：K 和 V block 在内层循环中流动，同时复用同一批 Q stage。
+- TMEM pipeline 深度 2：每个 Q stage 都有自己的 S/P/O TMEM 槽位，并在匹配的 barrier 触发后复用。
+
+下图从正确性 gate 切换到时间线视角，展示这些独立 ring 进入稳定状态后，哪些角色大致可以同时活跃。
+
+![Flash Attention 4 流水线结构](../img/flash_attention_pipeline_v2.png)
+
+请把它读成时间线，而不是 barrier 图。它展示的是同一时刻大致有哪些角色处于活跃状态；前面的 barrier-flow 图才是检查精确 producer-consumer wait 的地方。两张图合起来，回答了本节开头提出的两个不同问题。
+
+每一行对应代码中的一个角色分支：
+
+- WG3 warp 1 发出 TMA load。
+- WG3 warp 0 发出 score MMA 和 value MMA。
+- WG0 与 WG1 为两个 Q stage 运行 softmax。
+- WG2 释放或 rescale `O`，稍后再归一化最终输出。
+- WG3 warp 2 发出 TMA store。
+
+沿着图从左到右，可以追踪一个有代表性的 pipeline wave。load warp 先从 `Q0`、`K[n-1]`、`Q1`、`V[n-1]` 开始，然后持续流式读取更低索引的 K/V block。MMA warp 发出最早的 score MMA，产生 `S0` 和 `S1`，WG0/WG1 再把它们转成 `P0` 和 `P1`。
+
+重要的是，MMA warp 不会先跑完所有 score MMA，再跑所有 value MMA。两个 Q stage 预热好之后，它会交错两种 MMA：当前 `V` block 的一次 value MMA，下一次 `K` block 的一次 score MMA，如此继续：
+
+```text
+计算 Q0*K[n-1] 的分数
+计算 Q1*K[n-1] 的分数
+用 P0*V[n-1] 更新输出
+计算 Q0*K[n-2] 的分数
+用 P1*V[n-1] 更新输出
+计算 Q1*K[n-2] 的分数
+用 P0*V[n-2] 更新输出
+...
+```
+
+这种交错正是图中 score、softmax、correction 和 value 各行会重叠，而不是整齐依次执行的原因。
+
+WG2 行标为 `release / rescale`，两半对应我们已经见过的两种情况。第一个 K/V block 上还没有旧 `O`，所以 WG2 只参与允许 value MMA 继续的 handoff；后续 block 上，它可能在 value MMA 累计之前先 rescale 旧的 `O`。归一化和 TMA store 只会发生一次，在 attention task 的最后一个 K/V block 之后。
+
+没有一个 GEMM 风格的单 pipeline 可以描述 FA4，因为 Q、K/V 和 TMEM 槽位都在独立调度上前进。TIRx 把这些调度显式保留下来，用独立的 tile buffer、`PipelineState` cursor 和 barrier phase 表达，而不是把 kernel 藏进一个巨大的 monolithic primitive。代价是移动部件更多；收益是复杂度仍然可见、可检查。
+
+## Rescaling 与 Writeback
+
+rescale 是必须的，不是可以丢掉的优化。online softmax 可能随着每个新 score tile 抬高逐行最大值；一旦发生，早先 block 累计到 `O` 中的内容就是按 *旧* 最大值缩放的。这样早先每一项都会大出一个 `exp(m_new - m_old)` 因子。跳过 correction 会让这些 block 权重过大，最终输出就是错的。修正方式是一次 TMEM → registers → TMEM tile 操作：
+
+$$O_{\text{old}} \leftarrow O_{\text{old}} \cdot e^{(m_{\text{old}} - m_{\text{new}}) / \sqrt{d}}$$
+
+工作分给两个角色完成。softmax 计算逐行 scale，并把它投递到 SMEM mailbox；WG2 等待 `softmax_corr.full`，把当前 `O` 从 TMEM 读出，乘上该 scale，再把 `O` 写回：
+
+```python
+RESCALE_TILE = T.meta_var(16)
+o_row = T.wg_reg_tile(RESCALE_TILE)
+Tx.copy_async(o_row, O_region[i_q, d_start : d_start + RESCALE_TILE])
+Tx.mul(o_row, o_row, acc_scale)
+Tx.copy_async(O_region[i_q, d_start : d_start + RESCALE_TILE], o_row)
+T.ptx.tcgen05.wait.st()
+```
+
+值得强调的是，这是覆盖整个 `O` accumulator 的一次完整 TMEM → registers → TMEM tile 操作，不是一点标量记账；它和其他阶段一样，也有自己的解读卡：
+
+> **Tile-primitive 解读：Correction（rescale）**
+> - Scope：WG2，完整 warpgroup。
+> - Layout：TMEM 中的 `O` → 寄存器 → TMEM 中的 `O`（`O_region[i_q]`）。
+> - Dispatch：用 `tcgen05.ld` 读取，用 TMEM store 写入；中间做寄存器乘法。
+> - 交接：等待 `softmax_corr.full`；到达 `p_o_rescale`（→ value MMA）和 `softmax_corr.empty`（→ softmax）。
+
+端到端追踪同步过程：
+
+1. Softmax 把 scale 值写入 SMEM。
+2. WG2 等待 `softmax_corr.full`。
+3. WG2 在 TMEM 中 rescale `O`。
+4. WG2 到达 `p_o_rescale`。
+5. WG3 的 value MMA 现在可以消费 `P`，并累计到 rescale 后的 `O` tile。
+
+WG2 读取之后，`softmax_corr.empty` 会释放 SMEM 槽位，循环随之闭合，softmax 可以在下一次 iteration 复用 mailbox。
+
+K/V loop 结束后，WG2 从 correction 切换到 epilogue。它等待最终的 `row_sum` 和 `o_ready`，从 TMEM 读取最终 `O`，乘以 `1 / row_sum`（也就是一开始推迟的归一化），转换成 fp16，并写入 `O_smem`。然后 WG3 的 TMA store warp 把 `O_smem` 搬回 GMEM。
+
+如果你打算扩展这个 kernel，有一个限制值得标出。它只计算 forward output，而训练时的 forward pass 通常还要保存 backward pass 需要的 log-sum-exp（LSE）。加入 LSE 时有一个缩放细节要记住：这个 kernel 把 `row_max` 保留为 *未缩放* 的原始 `QK^T` score 最大值，而 `row_sum` 累计的是 `exp((S - row_max) / sqrt(d))`。因此形成自然对数 LSE 时，必须把 `1/\sqrt{d}` 因子重新应用到 `row_max` 上：
+
+$$\mathrm{LSE}_i = \log(\mathrm{row\_sum}_i) + \mathrm{row\_max}_i / \sqrt{d}$$
+
+这个实现只输出 forward 结果，不写 LSE。
+
+## Causal Masking
+
+causal attention 增加了一个约束：一个 query 只能 attend 到自身位置及之前的 key。kernel 用两种互补方式满足它，一种便宜，一种精确。
+
+便宜的方式是直接跳过整块工作。很多 K/V block 完全位于对角线上方，对给定 Q block 没有任何贡献，因此 `get_n_block_max(...)` 会计算该 block 最多可能需要的最后一个 block，循环就根本不加载、不计算剩余部分。
+
+精确的方式处理跨过对角线的 block，也就是一部分列有效、一部分无效的情况。这些 block 仍然运行 score MMA，但 softmax 会在指数化之前把无效列 mask 掉。对每一行，它根据该行 query 位置和 block offset 推导一个列上限，保留不超过该上限的列，并在寄存器中把之后每一列设为 `-inf`，让这些列既不贡献 row max，也不贡献 `exp2` numerator。
+
+实现并不是逐元素分支，而是用 `mask_r2p(...)` 应用这个上限：它把上限转成整个 32 宽 score chunk 上的 bit mask，并一次性 mask 整个 chunk。完全位于对角线下方的 block 保留所有列，不需要 mask。
+
+从 tile-primitive 视角看，causal mode 完全不改写数据路径。它只是缩短 K/V trip count，并把一个 masking 步骤插入寄存器驻留的 softmax 中，位于 score MMA 和 `P` writeback 之间。
+
+## GQA 支持
+
+Grouped Query Attention 允许多个 query head 共享一个 K/V head。这能节省内存带宽，但也提出一个打包问题：如何只保留一个 K/V tile，同时让多个 query head 都使用它？这个 kernel 的答案是：一次处理一整组 query head，让它们共同对应一个调度出来的 `kv_head_idx`：
+
+```python
+GQA_RATIO = num_qo_heads // num_kv_heads
+SEQ_Q_PER_TILE = BLK_M // GQA_RATIO
+```
+
+技巧在于重新解释 128 行 Q tile。对于 `GQA_RATIO=4`，它们不再表示 128 个序列位置，而是表示 32 个序列位置 × 4 个 query head；这些 query head 被打包在一起，共乘同一个 K/V tile。行解码如下：
+
+```text
+seq_pos = row // GQA_RATIO
+q_head  = row % GQA_RATIO
+```
+
+Q load 用一个 3D view 表达这种打包。源数据是自然的 `Q[batch, seq, qo_head, dim]` 布局，目标则是 score MMA 随后会当作扁平 `128 x HEAD_DIM` operand 读取的同一个 SMEM tile。view 负责调和这两种形态，而且不需要任何额外 copy：
+
+```python
+Q_smem_3d = Q_smem.view(SMEM_PIPE_DEPTH_Q, SEQ_Q_PER_TILE, GQA_RATIO, HEAD_DIM)
+Tx.copy_async(
+    Q_smem_3d[i_q, :, :, :],
+    Q[batch_idx,
+      m_start : m_start + SEQ_Q_PER_TILE,
+      kv_head_idx * GQA_RATIO : (kv_head_idx + 1) * GQA_RATIO,
+      :],
+    **tma_copy_q,
+)
+```
+
+K 和 V 从不在内存中展开，而这正是 GQA 的意义：`kv_head_idx` 对应的单个 K/V tile，会被打包进 Q 行里的全部 `GQA_RATIO` 个 query head 复用。输出侧与输入侧镜像对应，epilogue 之后用匹配的 3D view，把打包行存回 `O[batch, seq, qo_head, dim]`。
+
+结果是，GQA 完全生活在 Q-load 和 O-store 边界上。在内部计算路径中，score MMA 仍然看到一个普通的 `128 x HEAD_DIM` Q tile，其余 tile-primitive 图完全不变。
+
+## Tile 调度
+
+scheduler 的工作是把每个 CTA 映射到一个 `(batch, kv_head, m_block)` attention task；合适的策略取决于 masking 是否让这些 task 代价相同：
+
+- 非 causal mode 使用 `FlashAttentionLinearScheduler`。每个 task 的工作量相同，因此一个固定 CTA 池按 `num_ctas` 前进，就足以把任务均匀摊开。
+- Causal mode 使用 `FlashAttentionLPTScheduler`，因为 causal masking 会让工作量极不均匀：靠近开头的 Q block 大约只 attend 一个 K/V block，而靠近结尾的 Q block 会 attend 所有 block。朴素切分会让某些 CTA 远晚于其他 CTA 完成，所以 longest-processing-time scheduler 会优先安排重任务以拉平完成时间，同时仍尽量保持相邻 batch/head task 聚在一起，利于 L2 locality。
+
+尽管两种 scheduler 不同，它们暴露的循环接口完全相同：
+
+```python
+while scheduler.valid():
+    m_block_idx = scheduler.m_block_idx
+    batch_idx = scheduler.batch_idx
+    kv_head_idx = scheduler.head_idx
+    # 用对应 K/V block 范围处理一个 Q block
+    scheduler.next_tile()
+```
+
+唯一的行为差异在于 `next_tile()` 做什么：非 causal mode 下，它会让 CTA 前进到另一个 task；causal mode 下，它会在当前 task 后结束循环。无论哪种方式，这都只是调度决策：它选择 CTA 拥有 *哪个* attention tile，而不改变这个 tile 如何计算。循环内部仍然运行同样的本地 primitive：TMA load、score MMA、softmax、value MMA、correction、TMA store。
+
+## 编译与验证
+
+上面的内容都是摘录；要把所有东西合在一起并真正运行 kernel，我们会从 `tirx-kernels` 导入真实实现、编译它，并与 torch reference 对比。完整 kernel 位于 `tirx-kernels` 仓库中的 [`flash_attention4.py`](https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/attention/flash_attention4.py)，本章讲过的所有部件都组装在这个文件里。它和 GEMM 验证 cell 有两点不同：Flash Attention 的入口更丰富（`get_flash_attention4_kernel`），而且它为内建 profiler 多接收一个 `profiler_buf` 参数。整章只需要运行这一格：
+
+```python
+import torch
+import torch.nn.functional as F
+import tvm
+from tirx_kernels.attention.flash_attention4 import (
+    get_flash_attention4_kernel, PROFILER_BUFFER_SIZE)
+
+B, S, Hq, Hkv, D = 1, 1024, 32, 8, 128   # GQA：32 个 query head 共享 8 个 KV head
+Q = torch.randn(B, S, Hq, D, dtype=torch.float16, device="cuda")
+K = torch.randn(B, S, Hkv, D, dtype=torch.float16, device="cuda")
+V = torch.randn(B, S, Hkv, D, dtype=torch.float16, device="cuda")
+O = torch.empty(B, S, Hq, D, dtype=torch.float16, device="cuda")
+prof = torch.zeros(PROFILER_BUFFER_SIZE, dtype=torch.uint64, device="cuda")
+
+kernel = get_flash_attention4_kernel(B, S, S, Hq, Hkv, D, is_causal=False)
+target = tvm.target.Target("cuda")
+with target:
+    ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+ex.mod(Q, K, V, O, prof)   # 和其他章节一样，ex.mod 直接接收 torch tensor
+torch.cuda.synchronize()
+
+# torch reference；enable_gqa 允许 32 个 query head 共享 8 个 KV head
+qt, kt, vt = (x.transpose(1, 2).float() for x in (Q, K, V))
+ref = F.scaled_dot_product_attention(qt, kt, vt, enable_gqa=True).transpose(1, 2).half()
+torch.testing.assert_close(O, ref, rtol=1e-2, atol=1e-2)
+print(f"FA4: B={B} S={S} Hq={Hq} Hkv={Hkv} D={D}, non-causal -> PASS")
+```
+
+**预期输出**：`... -> PASS`。kernel 用 fp32 累计 online softmax，但它和高精度 reference 之间仍然存在几类近似：输入和 operand 的 fp16 存储与舍入；基于 `exp2` 的 softmax 重写（把每个指数改写成 `scale_log2 = log2(e)/√d` 形式）；online-softmax 的重排和逐行 rescaling，它按运行尺度分 block 求和，而不是一次性求和；最后还有 writeback 时对 `O` 的 fp16 cast。这里选择的 `rtol`/`atol` 和源 kernel 自己的测试一致，是为了同时覆盖这些因素相对 torch reference 的差异，而不只是单独覆盖 fp16 舍入。因此如果你看到真正失败，而不是边界附近的小偏差，就应该把它当作指向 softmax 路径的路标：漏掉了 `s_ready` / `p_o_rescale` / `p_ready_2` wait，或者 `row_max` / `row_sum` 更新没有被 rescale 步骤正确应用。这些正是本章用 barrier 反复处理的 handoff。
+
+## 与 GEMM 的差异
+
+下表沿发生变化的轴比较 FA4 与 GEMM：
+
+| 方面 | GEMM | Flash Attention 4 |
+|------|------|-------------------|
+| MMA 阶段 | 一个重复的 MMA | score MMA 和 value MMA |
+| MMA 之间的工作 | 除 pipeline handoff 外没有额外工作 | online softmax、masking 和 O rescaling |
+| 运行状态 | 只有 accumulator | row max、row sum、O accumulator |
+| 主要中间值 | accumulator TMEM tile | S、P 和 O TMEM tile region |
+| Warp 角色 | TMA producer、MMA consumer、writeback | TMA load、MMA、softmax、correction、TMA store |
+| Barrier | 主要是 load/compute/writeback handoff | 额外的 score/softmax/value/correction handoff |
+| 调度单元 | 输出矩阵 tile | attention task：`(batch, kv_head, m_block)` |
+
+这些差异全部可以追溯到本章开头那条结构性变化：第二个 MMA，以及夹在两个 MMA 之间的 softmax。另一方面，底层 TIRx contract 完全没有改变：
+
+- tile primitive 说明哪个 tile 在移动或计算；
+- 周围的 scope 说明哪些线程协作；
+- layout 说明 tile 住在哪里；
+- barrier 说明下一个角色什么时候可以消费它。
+
+所以 FA4 比 GEMM 更难，并不是因为它依赖不同硬件，而是因为 tile 值更多、它们之间的 handoff 也更多。
+
+## 练习
+
+1. 和 GEMM 相比，FA4 在两个 MMA 阶段之间新增了什么 tile handoff？请说出 producer、TMEM tile 和 consumer。
+2. 为什么 softmax 要把 numerator tile `P` 写回 TMEM，而不是只把它留在寄存器中供 value MMA 使用？
+3. 任选 `p_o_rescale` 或 `p_ready_2`。这个 barrier 精确证明了什么？如果 value MMA 跳过这次等待，可能出什么错？
+
+**试着让你的 agent 做一遍**：任选一个没有注释过的 tile primitive，例如 epilogue 里的 `Tx.copy_async`、fp32 -> fp16 的 `Tx.cast`，或第二段 `gemm_pv` sub-MMA。让它写出 scope / layout / dispatch / handoff 卡片，然后对照源码里的 guard、allocation 和 wait 检查答案。
diff --git a/zh/chapter_gemm_advanced/index.md b/zh/chapter_gemm_advanced/index.md
new file mode 100644
index 00000000..3237131c
--- /dev/null
+++ b/zh/chapter_gemm_advanced/index.md
@@ -0,0 +1,964 @@
+(zh_chap_gemm_advanced)=
+# 用 Warp Specialization 和 Cluster 扩展 GEMM
+
+:::{admonition} 概览
+:class: overview
+
+- pipelined GEMM 仍然让一个 warpgroup 按顺序做 load、MMA 和 writeback；本章会移除这个瓶颈。
+- Step 7 把 warp 专门化为不同角色，Step 8 加入 2-CTA cluster，Step 9 加入多个 consumer。
+- 每一步都会移除一个串行瓶颈，最终接近 state-of-the-art 吞吐。
+:::
+
+上一章的 pipelined GEMM（{ref}`zh_chap_gemm_async`）已经很快，但它仍然要求一个 warpgroup 做所有事情：
+发射 load、运行 MMA，然后写回结果。即使用了 software pipeline，这一组 thread 仍然成为三个引擎的汇合点。
+
+症状很容易看见。Tensor Core 运行时 TMA unit 安静下来；结果 drain 到内存时 Tensor Core 安静下来；
+每个引擎都通过同一组 thread 等待其他引擎。越过这个问题的方法，是停止让一组 thread 做所有事情。
+
+我们会通过三步逐渐扩大协作范围来追求这个想法。Step 7（{ref}`zh_chap_warp_specialization`）
+把 warp 专门化为 producer、consumer 和 writeback 角色。Step 8（{ref}`zh_chap_cta_cluster`）
+把两个 CTA 组成一个 cluster，并跨它们的 shared memory 共享 operand。Step 9（{ref}`zh_chap_multi_consumer`）
+加入第二个 MMA consumer，让一个 staged tile 喂给两倍数学工作。
+
+把这三步看作同一种 pattern 在不同尺度上的展开，会很有帮助。Step 7 把完整 pipeline 保持在一个 CTA 内部：
+TMA 和 MMA 共享一个 warpgroup，而 writeback 在另一个 warpgroup 中运行。
+Step 8 把协作扩大到 CTA 之间，产生一个跨越两个 CTA 的 256×256 tile。
+Step 9 进一步提高 compute density：cluster output 增长到 512×256，每个 staged B tile 被两个 consumer 复用，
+我们也到达教程中最密集的变体。
+
+贯穿这一切，有一件事保持不变。SMEM、TMEM 和 register layout 仍然遵守前两章建立的 contract；
+变化的是*谁协作*，而不是数据如何布局。Step 8 是协作 scope 第一次扩展到单个 CTA 之外，
+因此它的 operand tile 会切分到两个 CTA 的 shared memory 中，一个 layout 会沿 `cbx` cluster axis 跨越两个 CTA。
+
+
+(zh_chap_warp_specialization)=
+## Step 7：Warp Specialization + Pipeline
+
+single-warpgroup kernel 留下性能的原因很简单：每个 thread 走同一条路径，先 load，再 compute，再 write。
+因此它在 loading 时，Tensor Core 无事可做；它在 computing 时，TMA engine 无事可做。
+修复方式是 *warp specialization*。我们不再让一组 thread 轮流做每项工作，而是把每项工作交给专门的 warp，
+并让这些 warp 同时运行，再由 software pipeline 缝合起来。这是 GEMM 路径中最大的架构变化，
+本章剩余内容都建立在它之上。这里的 benchmark 使用 M=N=K=4096。
+
+> **这一步改变什么：Scope**
+> - Scope：一个 warpgroup 按顺序走 load → MMA → writeback，变成三个并发角色（TMA producer、MMA consumer、writeback），由 full/empty barrier 连接。
+> - Layout：不变，与 Step 6 相同的 SMEM stage 和 TMEM accumulator。
+> - Dispatch：不变，TMA load、`tcgen05` MMA。
+
+**主题。**
+
+- Warp specialization：把不同 warp/warpgroup 专门用于不同任务
+
+- 高层 barrier 抽象：`TMABar`、`TCGen05Bar`、`MBarrier`
+
+- `PipelineState` 用于自动 stage/phase 管理
+
+- `warpgroup_sync` barrier ID 用于按 warpgroup 同步
+
+（multi-stage SMEM pipeline 和 persistent `ClusterPersistentScheduler2D` 从 Step 5–6 原样复用；这里只新增 scope split。）
+
+### 从顺序到并发
+
+在介绍角色和 barrier 之前，先隔离 warp specialization 要移除的 scheduling bottleneck 会很有帮助。
+下图用 Step-4 风格的 sequential timeline 作为 Step 4-6 中 specialization 前 kernel 的紧凑参考，
+并把它放在 Step 7 warp-specialized schedule 上方，让 engine utilization 的差异一眼可见。
+
+![Warp 专门化时间线](../img/warp_specialization_timeline.png)
+
+上方是 specialization 前的 single-warpgroup pattern：同一个未专门化的 thread group 同时拥有 load path 和 MMA path，
+因此一个引擎活跃时，另一个引擎很容易闲置。Step 5 和 Step 6 用 double buffering 和 persistent scheduling 改进了这个 baseline，
+但它们还没有把 loading 和 compute 分成独立 producer/consumer 角色。
+下方的 specialization 打破了这种轮流执行。TMA producer 在 MMA consumer 忙于计算时 prefetch 下一个 tile，
+writeback 则自行推进。producer warp 3 在 consumer warp 0 仍在处理当前 MMA 时发射下一次 load，
+因此两个引擎都不必等待对方。load/MMA handoff 使用两个 barrier：
+
+- **`tma2mma`**（TMA → MMA）：signal 已载入的 SMEM 数据已经 ready，可供 MMA 消费。
+- **`mma2tma`**（MMA → TMA）：signal MMA 已经读完一个 buffer，因此 TMA 可以为下一次 load 复用它。
+
+图中有个细节第一眼可能像错误：`mma2tma` 箭头会跨过一个 stage。原因是 ring buffer。
+`PIPE_DEPTH=2` 时有两个 SMEM buffer，stage 0 和 stage 1；TMA Load k=0 填充 buffer 0，TMA Load k=1 填充 buffer 1。
+当 MMA Compute k=0 读完 buffer 0 时，它 signal `mma2tma` 表示 buffer 空闲；
+但真正想重新使用 buffer 0 的 load 是 TMA Load k=2，而不是 k=1（它使用 buffer 1）。
+这就是为什么 MMA Compute k=0 的 `mma2tma` 箭头一路指向 TMA Load k=2。
+release 跳过一个 stage，只是因为 ring 有两个 slot。
+
+### Warp Roles
+
+timeline 展示了我们*为什么*拆分工作；下一个问题是*谁*做每一部分。
+specialization 把三个工作（load、compute、writeback）分配给特定 warp，让它们能同时运行。
+当 `WG_NUMBER=2` 时，kernel 使用两个 warpgroup（角色表中缩写为 WG）：
+
+| Actor | Location | Job |
+|-------|----------|-----|
+| **TMA Producer** | Warpgroup 1, warp 3 | 通过 TMA 持续 load A 和 B tile |
+| **MMA Consumer** | Warpgroup 1, warp 0 | 数据 ready 后立即运行 MMA |
+| **Writeback** | Warpgroup 0（全部 warp） | 读取 TMEM 结果，写入 GMEM |
+
+### 4 个 Barrier
+
+三个并发 actor 需要四个 barrier，而这四个 barrier 正好分成两个相反方向。
+forward path（TMA → MMA → Writeback）signal 数据 *readiness*；它的信息是“你等的 tile 到了”。
+backward path（Writeback → MMA → TMA）signal buffer *release*：“你想要的 slot 又空了”。
+一旦知道命名约定，名字就能自己读懂：每个都是 `source2destination`，所以 `tma2mma`
+就是 TMA signal MMA 的 barrier。
+
+| Barrier | Type | Direction | Meaning |
+|---------|------|-----------|---------|
+| **tma2mma** | `TMABar` | TMA -> MMA | “SMEM data is ready” |
+| **mma2tma** | `TCGen05Bar` | MMA -> TMA | “SMEM buffer can be reused” |
+| **mma2ld** | `TCGen05Bar` | MMA -> Writeback | “TMEM results are ready” |
+| **ld2mma** | `MBarrier` | Writeback -> MMA | “TMEM is free for next tile” |
+
+为什么每个 barrier 会有它自己的 *type*？type 来自 producer 如何宣布自己完成。
+**TMA Load** 使用 `TMABar`，即带 byte counting 的 mbarrier：当 transfer 的字节落地后，
+TMA 硬件自己 arrive 到 barrier，因此 consumer 能知道数据 ready，而不需要任何 thread poll。
+**TMA Store** 不能使用这个机制（store 没有人需要通知），所以它们退回到
+`cp_async.bulk.commit_group()` + `wait_group(0)`，issuing thread 只是在等待自己的写入 drain。
+**MMA operation** 使用 `TCGen05Bar`，当 MMA 完成时，`tcgen05.commit()` 指令会 signal 这个 barrier。
+
+这里有一个小细节会在 Step 8 产生回报。`arrive` 调用传入 `cta_mask=0`，
+因为在 single-CTA kernel 中没有其他 CTA 需要 signal。当 Step 8 形成 cluster 时，
+这个参数会变成非零，并成为唤醒协作 CTA 的机制。
+
+### PipelineState
+
+四个 barrier 会告诉角色 buffer *何时* ready；但还需要有东西追踪 pipeline 循环时每个角色位于*哪个* buffer。
+这正是 `PipelineState` 管理的 bookkeeping。ring buffer 同时携带两份 bookkeeping：
+当前位于哪个 slot，以及正在等待这个 slot 的 barrier 的哪个 “phase”。
+在 pipelined loop 中手动追踪二者，正是容易滋生 off-by-one 错误的事情；
+这里的 off-by-one 会让整个 kernel deadlock。`PipelineState` 存在的目的就是把二者绑在一起，免得你手动管理：
+
+```python
+tma_ps = PipelineState(PIPE_DEPTH, phase=1)   # Producer starts ready (phase=1)
+# tma_ps.stage = current stage index
+# tma_ps.phase = current phase (0 or 1)
+tma_ps.advance()                          # Advance to next stage
+```
+
+initial `phase` 会决定某个角色的第一次 `wait` 是让它运行，还是让它阻塞。
+pipe 两端的正确答案正好相反，这就是容易绊倒人的地方：
+- `phase=1`（producer）-> 第一次 `wait(phase=1)` 看到 barrier 仍在 phase 0；由于 0 != 1，它会**立即通过**。
+  这正是我们想要的，因为 buffer 一开始是空的，producer 应该可以立刻开始填充。
+
+- `phase=0`（consumer）-> 第一次 `wait(phase=0)` 看到 barrier 位于 phase 0；由于 0 == 0，它会**阻塞**。
+  这同样是我们想要的，因为还没有数据，consumer 在 producer arrive 前没有东西可读。
+
+如果给两端相同的 starting phase，你会得到 deadlock，或者更糟，silent corruption。
+所以这个选择值得认真做对。
+
+### `warpgroup_sync` Barrier IDs
+
+specialization 引入了一个很容易踩到的同步危险。一旦每个 warpgroup 运行不同代码路径，
+熟悉的 `cta_sync()` 就会 deadlock：它使用硬件 barrier #0，并要求*每个* CTA thread arrive；
+但在 warpgroup branch 内，只有一部分 thread 存在。我们需要的是作用域为单个 warpgroup 的 barrier。
+GPU 给了我们 16 个 named barrier（ID 0–15），所以 kernel 会使用 `warpgroup_sync(10)`，
+它只同步一个 warpgroup 内的 thread。当多个 warpgroup 都需要各自同步时（multi-consumer Step 9 中就会这样），
+它们通过 `warpgroup_sync(wg_id + 10)` 使用不同 ID，避免在同一个硬件 barrier 上碰撞。
+
+**实现。**
+
+这里使用 `PIPE_DEPTH=2`，这是仍然能让 load 和 compute overlap 的最小深度。
+更深的 pipeline 可以隐藏更多内存延迟，直到 SMEM 预算限制为止；下面的 *When Step 7 misbehaves* 会详细讨论这个取舍。
+现在所有部件都已具备（角色、四个 barrier、`PipelineState` 和 warpgroup-scoped sync），我们可以组装完整 kernel：
+
+```python
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+from tvm.tirx.lang.pipeline import TMABar, TCGen05Bar, MBarrier, PipelineState
+from tvm.tirx.lang.tile_scheduler import ClusterPersistentScheduler2D
+
+SM_COUNT = 148  # Number of SMs on NVIDIA B200 GPU
+F16_SIZE = 2
+
+def hgemm_v7(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    K_TILES = K // BLK_K
+    PIPE_DEPTH = 2
+    WG_NUMBER = 2
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_N, BLK_K))
+    D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_N))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        bx = T.cta_id([SM_COUNT])
+        wg_id = T.warpgroup_id([WG_NUMBER])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        # --- Allocation ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        tma2mma = TMABar(pool, PIPE_DEPTH)
+        mma2tma = TCGen05Bar(pool, PIPE_DEPTH)
+        mma2ld  = TCGen05Bar(pool, 1)
+        ld2mma  = MBarrier(pool, 1)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout)
+        Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout)
+
+        # --- Barrier init ---
+        tma2mma.init(1)
+        mma2tma.init(1)
+        mma2ld.init(1)
+        ld2mma.init(128)   # all 128 Warpgroup 0 threads arrive
+        pool.commit()
+
+        # --- TMEM alloc + fence ---
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+            (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]))
+
+        # --- Tile scheduler ---
+        tile_scheduler = ClusterPersistentScheduler2D(
+            "ts", num_m_tiles=M // BLK_M, num_n_tiles=N // BLK_N,
+            l2_group_size=8, num_clusters=SM_COUNT)
+        tile_scheduler.init(bx)
+        m_st = T.meta_var(tile_scheduler.m_idx * BLK_M)
+        n_st = T.meta_var(tile_scheduler.n_idx * BLK_N)
+
+        # =============================================
+        # Warpgroup 1: TMA Producer (warp 3) + MMA Consumer (warp 0)
+        # =============================================
+        if wg_id == 1:
+            if warp_id == 3:
+                # === TMA Producer ===
+                tma_ps = PipelineState(PIPE_DEPTH, phase=1)
+
+                @T.inline
+                def tma_load(k_offset):
+                    Tx.copy_async(Asmem[tma_ps.stage, :, :],
+                                  A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=1,
+                                  mbar=tma2mma.ptr_to([tma_ps.stage]))
+                    Tx.copy_async(Bsmem[tma_ps.stage, :, :],
+                                  B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=1,
+                                  mbar=tma2mma.ptr_to([tma_ps.stage]))
+
+                if T.filter(lane_id, T.ptx.elect_sync()):
+                    while tile_scheduler.valid():
+                        for k in range(K_TILES):
+                            mma2tma.wait(tma_ps.stage, tma_ps.phase)
+                            tma_load(k * BLK_K)
+                            tma2mma.arrive(tma_ps.stage,
+                                           (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE)
+                            tma_ps.advance()
+                        tile_scheduler.next_tile()
+
+            elif warp_id == 0:
+                # === MMA Consumer ===
+                mma_ps = PipelineState(PIPE_DEPTH, phase=0)
+                ld_ps = PipelineState(1, phase=1)
+
+                if T.filter(lane_id, T.ptx.elect_sync()):
+                    while tile_scheduler.valid():
+                        # Wait for TMEM to be free from previous tile's writeback
+                        ld2mma.wait(ld_ps.stage, ld_ps.phase)
+                        ld_ps.advance()
+
+                        for k in range(K_TILES):
+                            tma2mma.wait(mma_ps.stage, mma_ps.phase)
+                            Tx.gemm_async(
+                                tmem[:, :BLK_N],
+                                Asmem[mma_ps.stage, :, :],
+                                Bsmem[mma_ps.stage, :, :],
+                                accum=(k != 0), dispatch="tcgen05", cta_group=1)
+                            mma2tma.arrive(mma_ps.stage, cta_group=1, cta_mask=0)
+                            mma_ps.advance()
+
+                        # Signal results ready for writeback
+                        mma2ld.arrive(0, cta_group=1, cta_mask=0)
+                        tile_scheduler.next_tile()
+
+        # =============================================
+        # Warpgroup 0: Writeback
+        # =============================================
+        elif wg_id == 0:
+            wb_ps = PipelineState(1, phase=0)
+            reg_f16 = T.alloc_local((BLK_N,), d_type)
+
+            while tile_scheduler.valid():
+                # Wait for MMA results
+                mma2ld.wait(wb_ps.stage, wb_ps.phase)
+                wb_ps.advance()
+
+                # Read TMEM -> registers (warpgroup scope)
+                reg = T.alloc_local((BLK_N,), acc_type)
+                reg_wg = reg.view(128, BLK_N,
+                    layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+                Tx.wg.copy_async(reg_wg[:], tmem[:, :BLK_N])
+                T.ptx.tcgen05.wait.ld()
+
+                # Signal TMEM free (all 128 threads arrive)
+                ld2mma.arrive(0, cta_id=0, pred=True)
+
+                # Cast fp32 -> fp16
+                Tx.cast(reg_f16[:], reg[:])
+
+                # Write to Dsmem + TMA store
+                Tx.copy(Dsmem[warp_id * 32 + lane_id, :], reg_f16[:])
+                T.ptx.fence.proxy_async("shared::cta")
+                T.cuda.warpgroup_sync(10)
+                if warp_id == 0:
+                    if lane_id == 0:
+                        Tx.copy_async(D[m_st:m_st+BLK_M, n_st:n_st+BLK_N],
+                                      Dsmem[:, :], dispatch="tma")
+                        T.ptx.cp_async.bulk.commit_group()
+                        T.ptx.cp_async.bulk.wait_group(0)
+                T.cuda.warpgroup_sync(10)
+
+                tile_scheduler.next_tile()
+
+        # --- Cleanup ---
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+要运行这些 kernel 中的任意一个，可以复用我们在 Step 1（{ref}`zh_chap_gemm_basics`）中展示过的 compile / run / check harness：把 `hgemm_v1` 换成 `hgemm_v7`、`hgemm_v8` 或 `hgemm_v9`，并选择例如 `M=N=K=4096` 的问题规模。注意，cluster 版本要求 `M` 和 `N` 是 cluster tile 的倍数（Step 8 为 `256×256`，Step 9 为 `512×256`），因此很小的 `128×128` 规模根本不会产生 tile。每个 step 请在全新的 Python session 中单独编译，切换 step 前重启 kernel，因为这些 kernel 会复用内部名字，而编译器会保留 session 内状态。各 step 的计时汇总在下面的 *End-to-End Result* 中。
+
+### Epilogue（Writeback）细节
+
+Step 7 的 epilogue 可以相当简单。因为只有 `BLK_N=128` 列，writeback warpgroup 能一次性把整个 TMEM tile 读入寄存器，然后发起一次 TMA store。Step 8 和 Step 9 就没有这个余裕了，所以后面会引入 chunking；但现在流程是：
+
+1. 等待 MMA：`mma2ld.wait(phase)`。本教程中的 Step 8 和 Step 9 会在这里额外加一个保守的 `fence.after_thread_sync()`；MMA-completion mbarrier 已经覆盖了顺序保证，大多数 kernel（包括 CUTLASS）都会省略它，所以 Step 7 也省略。
+2. 读取 TMEM -> registers（每个线程 128 个 fp32；warpgroup scope 下通过 `Tx.copy_async(reg_wg, tmem[:, :BLK_N])`，随后 `T.ptx.tcgen05.wait.ld()`）。
+3. 通知 MMA：`ld2mma.arrive(0, cta_id=0, pred=True)`（全部 128 个线程 arrive）；此时 TMEM 可以被下一个 tile 复用。这两个 `arrive` kwargs 在 cluster 版本中会再次出现：`cta_id` 表示要 signal *哪个 CTA 的* barrier 副本（`0` = 本 CTA，也就是 local barrier；Step 8 的 cooperative arrive 会改用 `cta_mask` 指向 CTA-0），`pred` 是逐线程谓词，用来决定该线程是否真的 arrive（这里是 `True`，所以每个 writeback 线程都会计入 arrival 总数）。
+4. 在寄存器中把 fp32 转成 fp16。
+5. 写 registers -> Dsmem，然后用 `fence.proxy_async("shared::cta") + warpgroup_sync(10)` 刷新。
+6. 通过 `cp_async.bulk.commit_group() + wait_group(0)` 做 TMA store，把 Dsmem 写回 GMEM。
+
+Step 8（`BLK_N=256`）和 Step 9（每个 consumer 的 `MMA_N=256`）不能继续保持这种一次性形式，原因是寄存器压力。每个线程读取 256 个 fp32 值，意味着 256 × 4 = 1024 字节必须同时驻留在该线程寄存器中，这有溢出到 local memory 的风险；同时还会迫使 Dsmem buffer 变大。因此这些 step 会把 writeback 拆成 `EPI_N` 列的 chunk（`EPI_N=64`）：每次 iteration 只保留 `EPI_N` 个 fp32 寄存器值，并发起一个相应更小的 TMA store，用少量额外 store 指令换取更舒服的寄存器预算。
+
+**实现说明。**
+
+- **Persistent kernel**：`bx = T.cta_id([SM_COUNT])` --- 每个 SM 一个 CTA，在 tile 上循环
+
+- **L2-friendly scheduling**：`ClusterPersistentScheduler2D` 按有利于 cache locality 的顺序安排 tile
+
+- 这种模式 --- warp specialization 加 software pipelining --- 在高性能 GEMM kernel 中很常见，包括 CUTLASS 风格的设计。
+
+### Step 7 行为异常时
+
+Step 7 是第一个让 TMA load、`tcgen05` MMA 和 writeback 同时在路上的 GEMM kernel。Step 8 和 Step 9 会反复遇到同样的失败模式：barrier 计数不匹配、role guard 放错位置、缺少 fence，或 TMA store 还没 drain 就复用了 staging buffer。这类问题的调试清单汇总在 {ref}`zh_chap_warp_spec_debug`。
+
+**Pipeline depth 调优。** Step 7 kernel 使用最小的 `PIPE_DEPTH=2`。把它推到 4 或 6，可以让 TMA producer 领先 MMA consumer 更远，从而隐藏更多内存延迟；但代价是消耗更多 SMEM，而 SMEM 是有限的。B200 每个 SM 提供 228 KB（见 {ref}`zh_chap_background` 中的 *Numbers to Keep in Mind*）。在 `BLK_M=BLK_N=128, BLK_K=64, fp16` 下，每个 pipeline stage 中 A 和 B 合计消耗 `(128*64 + 128*64) * 2 = 32 KB`，`Dsmem` writeback staging buffer 还要再加 32 KB。因此 `PIPE_DEPTH=4` 大约是 160 KB，`PIPE_DEPTH=6` 大约是 224 KB，已经贴近预算上限。想再深入，就必须重新设计 writeback staging 策略。
+
+---
+
+warp specialization 让一个 CTA 内的线程协作起来。下一步会把这种协作扩展到 CTA 边界之外，让两个 CTA 共同处理一个更大的 tile。
+
+
+(zh_chap_cta_cluster)=
+## Step 8: 2-CTA Cluster
+
+Step 7 让各个引擎开始重叠，但每个 CTA 仍然孤立地计算自己的 128×128 tile，重新加载邻居无法借用的 operand。Step 8 打破这种隔离。两个 CTA 组成一个 cluster，并获得访问彼此 shared memory 的能力；于是单个 cooperative `tcgen05` MMA 可以产生一个横跨两者的 256×256 tile，而一次 B 加载现在能喂给两倍的 MMA 工作。和前面一样，M=N=K=4096。
+
+> **这个 step 改变了什么：Scope + Layout + Dispatch**
+> - Scope：协作 scope 现在跨越 cluster 中的两个 CTA，而不是一个。
+> - Layout：operand tile 被拆分到两个 CTA 的 SMEM 中；CTA 0 拥有共享的 completion barrier（`remote_view`）。
+> - Dispatch：MMA 增加 `cta_group` / `cta_mask`，让 `tcgen05` 以 2-CTA cooperative op 运行。
+
+**主题。**
+
+- CTA cluster：多个 CTA 在一个更大的 tile 上协作
+
+- 通过 `map_shared_rank` 进行 cross-CTA SMEM 访问
+
+- 在 256x256 cluster tile 上使用 `cta_group=2` 执行 cooperative MMA
+
+- 使用 `cta_mask` 做 cross-CTA barrier signaling
+
+
+### Cluster Tile 形状
+
+整个优化建立在一个硬件能力上：使用 `cta_group=2` 时，MMA 可以读取 *两个* CTA staged 的 operand tile，而不仅是自己所在 CTA 的 tile。每个 CTA 加载 stored B 的一个 128 行切片；转置后，它会变成 128 个逻辑输出列；cooperative MMA 再把两个切片缝合回一个 operand。下图追踪两个 CTA 的 A/B 切片如何合并成单个 256×256 cluster tile：
+
+```{raw} html
+

+
+

+```
+*交互图：每个 CTA 拥有一个 A 行切片和一个 stored-B 行切片，然后通过 cluster（DSMEM）读取另一个 CTA 的 stored-B 切片。经过 `B.T` 后，两个 stored-B 切片覆盖完整的输出列范围，因此这对 CTA 产生一个 256×256 输出 tile。*
+
+**为什么 A 和 B 要跨 cluster 拆分**：要看清 256×256 tile 如何分区，先回忆本教程把 GEMM 写成 `D = A @ B.T`，其中 stored B 的形状是 `N x K`。有两个 CTA 在一个 cluster 中时，拆分方式非常自然：
+
+- **A 竖向拆分**：CTA-0 持有 A0（行 0-127），CTA-1 持有 A1（行 128-255）。堆叠后是 `[A0; A1]`（256 行）。
+- **Stored B 按行拆分**：CTA-0 加载 B 行 0-127，CTA-1 加载 B 行 128-255。因为数学上使用 `B.T`，这两个 stored row slice 会变成逻辑右操作数的两个 128 列切片。
+- 使用 `cta_group=2` 时，MMA 硬件通过 cross-CTA shared memory access 从 **两个** CTA 的 SMEM 中读取 B，因此它能看到完整的逻辑输出列范围。
+- 结果：两个 CTA 协作处理一个 256x256 输出 tile。每个 CTA 写出这个 tile 的一个 128x256 行条带。
+
+这里值得停一下，看看为什么这是真正的收益，而不只是重新洗牌。每个 CTA 仍然只加载 128×K 的 A 和 128×K 的 B，因此整个 cluster staged operand 约为单个 CTA 的 2×；但它产生的是 256×256 tile，输出 FLOP 大约是 128×128 tile 的 4×。因此每个 staged-operand byte 对应的 MMA 工作量约翻倍，因为每个 CTA 的 B 切片会通过 cooperative MMA 与另一个 CTA 的 A 切片复用。换句话说，arithmetic intensity 大约翻倍，而这正是仍偏 memory-bound 的 kernel 所需要的杠杆：End-to-End 表中约 2.2× 的加速来自让同一批字节服务更多数学计算。
+
+### Tile 地址计算
+
+现在 cluster 成了工作单元，tile scheduler 也必须按 cluster tile 计数。它返回的每个 `(m_idx, n_idx)` 都表示一个完整的 256×256 区域，cluster 内的两个 CTA 会共同拆分这个区域。把 cluster 坐标转换成每个 CTA 实际加载的 per-CTA slice，形式如下：
+
+```python
+m_st = (m_idx * CTA_GROUP + cbx) * BLK_M
+n_st = (n_idx * CTA_GROUP + cbx) * BLK_N
+```
+
+两个 CTA 处理的是 *同一个* 256×256 cluster tile；单个坐标 `cbx`（该 CTA 在 cluster 内的位置，0 或 1）会在两个轴上选出这个 CTA 的贡献。`m_st` 选择该 CTA 拥有的输出行条带，`n_st` 选择它喂给 cooperative MMA 的 stored-B 切片，writeback 随后会写出 256 列输出范围的两个 128 列半块。还要注意，`num_m_tiles = M // 256` 和 `num_n_tiles = N // 256` 计数的是 cluster tile，而不是单个 CTA tile。
+
+乍看之下，`cbx` 同时出现在 `m_st` 和 `n_st` 中，好像一个行偏移泄漏到了列上；但两个用法都是正确的，值得拆开看。在 writeback 路径上，`cbx` 只属于 M 轴：每个 CTA 拥有不同的 128 行条带（`m_st = (m_idx * CTA_GROUP + cbx) * BLK_M`，因此 CTA-0 写 `m_idx*256 .. +128` 行，CTA-1 写接下来的 128 行），但两个 CTA 都会写 cluster tile 的 *完整* 256 个输出列。这也正是 store 的列坐标来自 cluster 的 `n_idx`（`n_st_epi = n_idx * 256 + no * 128`，完全没有 `cbx`），而不是 per-CTA `n_st` 的原因。`n_st` 之所以带着 `cbx`，是因为每个 CTA 会把不同的 stored-B 行切片加载进 MMA：在那里，`cbx` 是一个 *load* offset，而不是该 CTA 的输出列偏移。
+
+### 相比 Step 7 的代码变化
+
+相对 Step 7 的 diff 有六处改动，每一处都编码了刚才描述的 cluster contract 中的一个部分：
+
+```python
+# 1. Cluster launch
+cbx, cby = T.cta_id_in_cluster([CTA_GROUP, 1])   # cbx = CTA index within cluster (0 or 1)
+
+# 2. Cooperative MMA (was cta_group=1)
+Tx.gemm_async(..., cta_group=2)
+
+# 3. Cross-CTA shared memory access
+B_remote = T.ptx.map_shared_rank(Bsmem, cta_id=1)
+
+# 4. Cross-CTA barrier
+tma2mma_cta0 = T.decl_buffer(
+    [CTA_GROUP], "uint64",
+    data=T.ptx.map_shared_rank(tma2mma.ptr_to([0]), 0),
+    scope="shared"
+)
+
+# 5. mma2tma / mma2ld arrives go from cta_mask=0 (single CTA, Step 7)
+#    to cta_mask=3 (signal both CTAs in the cluster)
+mma2tma.arrive(mma_ps.stage, cta_group=CTA_GROUP, cta_mask=3)
+mma2ld.arrive(0, cta_group=CTA_GROUP, cta_mask=3)
+
+# 6. Cluster sync replaces cta_sync at the end
+T.cuda.cluster_sync()
+```
+
+
+### Cluster-Scope 变化
+
+这六处改动都源自同一个转变：协作 scope 现在是 cluster，而不是单个 CTA。下面几点说明这种扩展在实践中意味着什么：每个 CTA 如何找到自己的位置、cluster 以谁的 barrier 作为协调点，以及究竟哪个 CTA 发出 cooperative MMA。
+
+- **Cluster CTA ID**：`cbx` 告诉每个 CTA 它在 cluster 中的位置（0 或 1）。CTA-0 处理 A 行 0-127，CTA-1 处理行 128-255。
+
+- **Remote barrier view**：在 cluster 中，每个 CTA 都有自己的 SMEM 和自己的 barrier，这带来一个自然问题：如果 CTA-1 需要等待 CTA-0 产生的东西，它实际应该碰谁的 barrier？答案是指定 CTA-0 的 barrier 作为唯一协调点，并允许 cluster 中任意 CTA 访问它们。`map_shared_rank(tma2mma.ptr_to([0]), 0)` 会返回指向 CTA-0 barrier 的 cluster-wide pointer；TIRx wrapper 是 `tma2mma.remote_view(0)`。从那以后，每次 arrive 和 wait 都指向 CTA-0 的副本。
+
+- **MMA 只从 CTA-0 dispatch**：很容易把 `cta_group=2` 理解成并行发射两个引擎，但事实不是这样。CTA-0 只发出一个 `tcgen05.mma`，然后硬件驱动一个跨越两个 CTA 的 *单个 cooperative* MMA：它从两个 SM 的 SMEM 中读取 operand，并把 accumulator 写到两个 SM 的 TMEM 中。CTA-1 完全不发出 MMA。（每个 SM 只有一个 `tcgen05` 引擎，所以 `cta_group=2` 是一次 cross-SM MMA，不是两个引擎并排运行。）这就是代码用 `if cbx == 0:` guard MMA 的原因。
+
+- **Multicast arrive**：`tcgen05.commit(..., cta_group=2, cta_mask=3)` 只由 CTA-0 发出，但会 signal 两个 CTA 的 barrier。`cta_mask=3`（二进制 `11`）表示目标包含 CTA-0 和 CTA-1。
+
+- **ld2mma init count**：`init(128 * CTA_GROUP)` --- 两个 CTA 的 writeback warpgroup（各 128 个线程）都会 arrive。
+
+
+**实现。**
+
+```python
+def hgemm_v8(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    CTA_GROUP = 2
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    MMA_M, MMA_N = 256, 256
+    K_TILES = K // BLK_K
+    PIPE_DEPTH = 4
+    WG_NUMBER = 2
+    F16_SIZE = 2  # fp16
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (PIPE_DEPTH, BLK_N, BLK_K))
+    D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, 128))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        bx = T.cta_id([SM_COUNT])
+        cbx, cby = T.cta_id_in_cluster([CTA_GROUP, 1])
+        wg_id = T.warpgroup_id([WG_NUMBER])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        # --- Allocation ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        tma2mma = TMABar(pool, PIPE_DEPTH)
+        mma2tma = TCGen05Bar(pool, PIPE_DEPTH)
+        mma2ld  = TCGen05Bar(pool, 1)
+        ld2mma  = MBarrier(pool, 1)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout)
+        Dsmem = pool.alloc((BLK_M, 128), d_type, layout=D_layout)
+
+        # --- Barrier init ---
+        tma2mma.init(1)
+        mma2tma.init(1)
+        mma2ld.init(1)
+        ld2mma.init(128 * CTA_GROUP)  # both CTAs' writeback threads
+        pool.commit()
+
+        # --- TMEM alloc (cooperative) ---
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=CTA_GROUP)
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+            (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]))
+
+        # --- Tile scheduler (cluster tiles) ---
+        tile_scheduler = ClusterPersistentScheduler2D(
+            "ts", num_m_tiles=M // 256, num_n_tiles=N // 256,
+            l2_group_size=8, num_clusters=SM_COUNT // CTA_GROUP)
+        tile_scheduler.init(bx // CTA_GROUP)
+        m_idx = T.meta_var(tile_scheduler.m_idx)
+        n_idx = T.meta_var(tile_scheduler.n_idx)
+        m_st = T.meta_var((m_idx * CTA_GROUP + cbx) * BLK_M)
+        n_st = T.meta_var((n_idx * CTA_GROUP + cbx) * BLK_N)
+
+        # --- Cross-CTA barrier view ---
+        tma2mma_cta0 = tma2mma.remote_view(0)
+
+        # =============================================
+        # Warpgroup 1: TMA Producer (warp 3) + MMA Consumer (warp 0)
+        # =============================================
+        if wg_id == 1:
+            if warp_id == 3:
+                tma_ps = PipelineState(PIPE_DEPTH, phase=1)
+
+                @T.inline
+                def tma_load(k_offset):
+                    Tx.copy_async(Asmem[tma_ps.stage, :, :],
+                                  A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=CTA_GROUP,
+                                  mbar=tma2mma_cta0.ptr_to([tma_ps.stage]))
+                    Tx.copy_async(Bsmem[tma_ps.stage, :, :],
+                                  B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=CTA_GROUP,
+                                  mbar=tma2mma_cta0.ptr_to([tma_ps.stage]))
+
+                if T.filter(lane_id, T.ptx.elect_sync()):
+                    while tile_scheduler.valid():
+                        for k in range(K_TILES):
+                            mma2tma.wait(tma_ps.stage, tma_ps.phase)
+                            tma_load(k * BLK_K)
+                            if cbx == 0:
+                                tma2mma_cta0.arrive(tma_ps.stage,
+                                    CTA_GROUP * (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE)
+                            tma_ps.advance()
+                        tile_scheduler.next_tile()
+
+            elif warp_id == 0:
+                mma_ps = PipelineState(PIPE_DEPTH, phase=0)
+                ld_ps = PipelineState(1, phase=1)
+
+                if cbx == 0:
+                    if T.filter(lane_id, T.ptx.elect_sync()):
+                        while tile_scheduler.valid():
+                            ld2mma.wait(ld_ps.stage, ld_ps.phase)
+                            ld_ps.advance()
+
+                            for k in range(K_TILES):
+                                tma2mma.wait(mma_ps.stage, mma_ps.phase)
+                                Tx.gemm_async(
+                                    tmem[:, :MMA_N],
+                                    Asmem[mma_ps.stage, :, :],
+                                    Bsmem[mma_ps.stage, :, :],
+                                    accum=(k != 0), dispatch="tcgen05", cta_group=CTA_GROUP)
+                                mma2tma.arrive(mma_ps.stage, cta_group=CTA_GROUP, cta_mask=3)
+                                mma_ps.advance()
+
+                            mma2ld.arrive(0, cta_group=CTA_GROUP, cta_mask=3)
+                            tile_scheduler.next_tile()
+
+        # =============================================
+        # Warpgroup 0: Writeback (256 columns in 2 x 128-column chunks)
+        # =============================================
+        elif wg_id == 0:
+            wb_ps = PipelineState(1, phase=0)
+            reg_f16 = T.alloc_local((128,), d_type)
+
+            while tile_scheduler.valid():
+                mma2ld.wait(wb_ps.stage, wb_ps.phase)
+                wb_ps.advance()
+                T.ptx.tcgen05.fence.after_thread_sync()
+
+                for no in T.unroll(2):  # 2 chunks of 128 columns = 256 total
+                    reg = T.alloc_local((128,), acc_type)
+                    reg_wg = reg.view(128, 128,
+                        layout=TileLayout(S[(128, 128) : (1@tid_in_wg, 1)]))
+                    Tx.wg.copy_async(reg_wg[:], tmem[:, no * 128:(no + 1) * 128])
+                    T.ptx.tcgen05.wait.ld()
+                    Tx.cast(reg_f16[:], reg[:])
+                    Tx.copy(Dsmem[warp_id * 32 + lane_id, :], reg_f16[:])
+                    T.ptx.fence.proxy_async("shared::cta")
+                    T.cuda.warpgroup_sync(10)
+                    if warp_id == 0:
+                        if lane_id == 0:
+                            n_st_epi = T.meta_var(n_idx * 256 + no * 128)
+                            Tx.copy_async(D[m_st:m_st+BLK_M, n_st_epi:n_st_epi+128],
+                                          Dsmem[:, :], dispatch="tma")
+                            T.ptx.cp_async.bulk.commit_group()
+                            T.ptx.cp_async.bulk.wait_group(0)
+                    T.cuda.warpgroup_sync(10)
+
+                ld2mma.arrive(0, cta_id=0, pred=True)
+                tile_scheduler.next_tile()
+
+        # --- Cleanup ---
+        T.cuda.cluster_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=CTA_GROUP)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=CTA_GROUP)
+
+    return kernel
+```
+
+**2 个 CTA 带来的变化。**
+
+- `CTA_GROUP = 2`, `MMA_N = BLK_N * CTA_GROUP = 256`
+
+- `ld2mma.init(128 * CTA_GROUP)` --- 两个 CTA 的 writeback WG 都会 arrive
+
+- TMA arrive 字节数包含两个 CTA：`CTA_GROUP * (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE`
+
+- `tcgen05.alloc` 和 `tcgen05.dealloc` 必须使用 `cta_group=2`
+
+- Writeback 会把 256 个输出列拆成两个 128 列 chunk --- 一次读取全部 256 个 TMEM 列会超过寄存器容量。Step 9 会进一步把 chunk 缩小到 `EPI_N=64`
+
+- 末尾用 `cluster_sync()` 替换 `cta_sync()`（确保 TMEM dealloc 前所有 CTA 都已完成）
+
+额外的 arithmetic intensity 会直接体现在墙钟时间上：Step 8 在 4096³ 下达到 **0.104 ms**，相比相同规模下 70 ms 的 Step-1 算法约快 676×（见 End-to-End 表）。这个 kernel 现在已经开始偏向 compute-bound，这正好为 Step 9 铺路：我们会加入第二个 MMA consumer，让更多 Tensor Core 工作保持在飞行中。
+
+如果 Step 8 反而比 Step 7 *更慢*，问题几乎总是某个新的 cluster contract 写偏了。优先检查三件事：TMA arrive byte count 是否为 `CTA_GROUP * (BLK_M*BLK_K + BLK_N*BLK_K) * F16_SIZE`；256×256 cluster tile 对应的 scheduler 维度是否是 `num_m_tiles=M//256, num_n_tiles=N//256`；writeback 是否确实发起两次 TMA store（每个 128 列 chunk 一次），且每次都在 Dsmem 复用前 drain 完。
+
+---
+
+Cluster 提升的是 CTA *之间* 的复用。最后一步转向内部：给 producer 再配一个 MMA consumer，让每个 CTA *内部* 的计算密度更高。
+
+
+(zh_chap_multi_consumer)=
+## Step 9: Multi-Consumer Warp Specialization
+
+到 Step 8，MMA 已经真正忙起来了；但单个 consumer warp 消化 staged B tile 的速度毕竟有限，而这块 B tile 一直待在 SMEM 中，任何愿意读取它的角色都能复用。最后一个优化正是利用这一点：加入第二个 MMA consumer，用 *不同的* A block 乘同一个 B tile。每个 CTA 的计算密度翻倍，cluster 输出从 256×256 增长到 512×256。和前面一样，M=N=K=4096。
+
+> **这个 step 改变了什么：Scope + Layout**
+> - Scope：一个 MMA consumer 变成两个，由 `warp_id` 选择。
+> - Layout：同一个 staged B tile 被两个 consumer 复用；A 增加一个 consumer 轴。
+> - Dispatch：不变。
+
+**主题。**
+
+- 多个 MMA warp（consumer）以获得更高吞吐
+
+- 多个 writeback warpgroup，各自拥有独立 barrier slot
+
+- 本教程中最高优化 GEMM 变体使用的结构
+
+
+### Multi-Consumer 结构
+
+加入第二个 consumer 意味着 kernel 现在需要安排更多独立角色：两个 MMA warp，而不是一个；再配一个第二 writeback warpgroup，用来 drain 额外的 accumulator。在 `NUM_CONSUMER=2` 且 `WG_NUMBER=3` 时，kernel 现在横跨三个 warpgroup（角色表中缩写为 WG）：
+
+| Warpgroup | Warp | 角色 |
+|-----------|------|------|
+| **WG 2** | warp 0 | MMA consumer 0：`Asmem[..., 0] x B` -> TMEM cols `[0:256]` |
+| **WG 2** | warp 1 | MMA consumer 1：`Asmem[..., 1] x B` -> TMEM cols `[256:512]` |
+| **WG 2** | warp 3 | TMA producer：每个 stage 加载 2x A block + 1x B block |
+| **WG 0** | all | consumer 0 的 writeback：读取 TMEM `[0:256]` |
+| **WG 1** | all | consumer 1 的 writeback：读取 TMEM `[256:512]` |
+
+整个安排依赖一个不对称性。每个 consumer 都用自己的 A block 去乘 *同一个* staged B tile，因此一次 B 加载现在能喂给 2× 的 MMA 工作，B 在每个有效 FLOP 上的加载成本等效减半。我们共享 B 而不是 A，是因为两个 consumer 覆盖不同的 M 行条带：它们的 A block 确实是不同数据，而 B 对两者完全相同。练习 3 会让你说服自己：这是唯一可行的共享方式。
+
+### 相比 Step 8 的变化
+
+具体来说，支持第二个 consumer 会触碰 kernel 中的几个地方，而每处变化都可以追溯到同一个事实：每个 stage 现在要喂给并 drain 两个 A block 和两个 TMEM range，而 B 保持共享。下面这些修改会额外 stage 一个 A block，给每个 consumer 自己的 barrier slot，并为更高的 512×256 cluster tile 调整 tile addressing。
+
+- `Asmem = pool.alloc((PIPE_DEPTH, NUM_CONSUMER, BLK_M, BLK_K), ...)` --- 每个 stage 有 2 个 A block，每个 consumer 一个
+
+- TMA 同时加载 `Asmem[stage, 0]` 和 `Asmem[stage, 1]`，TMA arrive bytes 现在是 `CTA_GROUP * (NUM_CONSUMER * BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE`（多了一个 A block）
+
+- MMA warp 用 `warp_id` 选择哪个 A block 和哪个 TMEM range
+
+- `mma2tma.init(NUM_CONSUMER)` --- 每个 stage 中两个 consumer 都会 signal TMA
+
+- `mma2ld` 和 `ld2mma` 都有 `depth=NUM_CONSUMER` --- 每个 consumer 使用自己的 barrier slot（MMA 侧用 `warp_id`，writeback 侧用 `wg_id`）
+
+- Tile 地址：`m_st = (m_idx * NUM_CONSUMER * CTA_GROUP + cbx) * BLK_M` --- M 方向多了一个 `NUM_CONSUMER` 因子，因为每个 cluster tile 现在在 M 方向跨越 `NUM_CONSUMER` 个 consumer。Tile scheduler 使用 `num_m_tiles = M // 256 // NUM_CONSUMER`（cluster tile 为 512x256）
+
+- Writeback 使用分块的 `EPI_N`，让每次 iteration 中活跃在寄存器里的 TMEM-readback 值更少
+
+
+**实现。**
+
+```python
+def hgemm_v9(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    CTA_GROUP = 2
+    NUM_CONSUMER = 2
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    MMA_N = BLK_N * CTA_GROUP   # 256
+    K_TILES = K // BLK_K
+    PIPE_DEPTH = 4
+    EPI_N = 64
+    WG_NUMBER = 3
+    F16_SIZE = 2  # fp16
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                 (PIPE_DEPTH, NUM_CONSUMER, BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                 (PIPE_DEPTH, BLK_N, BLK_K))
+    D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                 (NUM_CONSUMER, BLK_M, EPI_N))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        bx = T.cta_id([SM_COUNT])
+        cbx, cby = T.cta_id_in_cluster([CTA_GROUP, 1])
+        wg_id = T.warpgroup_id([WG_NUMBER])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        # --- Allocation ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        tma2mma = TMABar(pool, PIPE_DEPTH)
+        mma2tma = TCGen05Bar(pool, PIPE_DEPTH)
+        mma2ld  = TCGen05Bar(pool, NUM_CONSUMER)   # depth=2, one slot per consumer
+        ld2mma  = MBarrier(pool, NUM_CONSUMER)     # depth=2, one slot per consumer
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((PIPE_DEPTH, NUM_CONSUMER, BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout)
+        Dsmem = pool.alloc((NUM_CONSUMER, BLK_M, EPI_N), d_type, layout=D_layout)
+
+        # --- Barrier init ---
+        tma2mma.init(1)
+        mma2tma.init(NUM_CONSUMER)  # each stage expects 2 arrivals
+        mma2ld.init(1)              # each slot gets 1 arrival
+        ld2mma.init(128 * CTA_GROUP)  # both CTAs' writeback threads
+        pool.commit()
+
+        # --- TMEM alloc (cooperative) ---
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=CTA_GROUP)
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+            (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]))
+
+        # --- Tile scheduler (512x256 cluster tiles) ---
+        tile_scheduler = ClusterPersistentScheduler2D(
+            "ts", num_m_tiles=M // 256 // NUM_CONSUMER, num_n_tiles=N // 256,
+            l2_group_size=8, num_clusters=SM_COUNT // CTA_GROUP)
+        tile_scheduler.init(bx // CTA_GROUP)
+        m_idx = T.meta_var(tile_scheduler.m_idx)
+        n_idx = T.meta_var(tile_scheduler.n_idx)
+        m_st = T.meta_var((m_idx * NUM_CONSUMER * CTA_GROUP + cbx) * BLK_M)
+        n_st = T.meta_var((n_idx * CTA_GROUP + cbx) * BLK_N)
+
+        tma2mma_cta0 = tma2mma.remote_view(0)
+
+        # =============================================
+        # Warpgroup 2: TMA Producer (warp 3) + 2 MMA Consumers (warp 0, 1)
+        # =============================================
+        if wg_id == 2:
+            if warp_id == 3:
+                # === TMA Producer: loads 2 A blocks + 1 B block per stage ===
+                tma_ps = PipelineState(PIPE_DEPTH, phase=1)
+
+                @T.inline
+                def tma_load(k_offset):
+                    m_st_c1 = T.meta_var(m_st + CTA_GROUP * BLK_M)
+                    Tx.copy_async(Asmem[tma_ps.stage, 0, :, :],
+                                  A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=CTA_GROUP,
+                                  mbar=tma2mma_cta0.ptr_to([tma_ps.stage]))
+                    Tx.copy_async(Asmem[tma_ps.stage, 1, :, :],
+                                  A[m_st_c1:m_st_c1+BLK_M, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=CTA_GROUP,
+                                  mbar=tma2mma_cta0.ptr_to([tma_ps.stage]))
+                    Tx.copy_async(Bsmem[tma_ps.stage, :, :],
+                                  B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K],
+                                  dispatch="tma", cta_group=CTA_GROUP,
+                                  mbar=tma2mma_cta0.ptr_to([tma_ps.stage]))
+
+                if T.filter(lane_id, T.ptx.elect_sync()):
+                    while tile_scheduler.valid():
+                        for k in range(K_TILES):
+                            mma2tma.wait(tma_ps.stage, tma_ps.phase)
+                            tma_load(k * BLK_K)
+                            if cbx == 0:
+                                tma2mma_cta0.arrive(tma_ps.stage,
+                                    CTA_GROUP * (NUM_CONSUMER * BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE)
+                            tma_ps.advance()
+                        tile_scheduler.next_tile()
+
+            elif warp_id < NUM_CONSUMER:
+                # === MMA Consumer: warp_id selects A block and TMEM range ===
+                mma_ps = PipelineState(PIPE_DEPTH, phase=0)
+                ld_ps = PipelineState(1, phase=1)
+
+                if cbx == 0:
+                    if T.filter(lane_id, T.ptx.elect_sync()):
+                        while tile_scheduler.valid():
+                            ld2mma.wait(warp_id, ld_ps.phase)
+                            ld_ps.advance()
+
+                            for k in range(K_TILES):
+                                tma2mma.wait(mma_ps.stage, mma_ps.phase)
+                                Tx.gemm_async(
+                                    tmem[:, warp_id * MMA_N:warp_id * MMA_N + MMA_N],
+                                    Asmem[mma_ps.stage, warp_id, :, :],
+                                    Bsmem[mma_ps.stage, :, :],
+                                    accum=(k != 0), dispatch="tcgen05", cta_group=CTA_GROUP)
+                                mma2tma.arrive(mma_ps.stage, cta_group=CTA_GROUP, cta_mask=3)
+                                mma_ps.advance()
+
+                            mma2ld.arrive(warp_id, cta_group=CTA_GROUP, cta_mask=3)
+                            tile_scheduler.next_tile()
+
+        # =============================================
+        # Warpgroup 0/1: Writeback (each reads its consumer's TMEM range)
+        # =============================================
+        elif wg_id < NUM_CONSUMER:
+            wb_ps = PipelineState(1, phase=0)
+            reg_f16 = T.alloc_local((EPI_N,), d_type)
+
+            while tile_scheduler.valid():
+                mma2ld.wait(wg_id, wb_ps.phase)  # wait for THIS consumer
+                wb_ps.advance()
+                T.ptx.tcgen05.fence.after_thread_sync()
+
+                # Read TMEM in EPI_N=64 column chunks (4 iterations for 256 cols)
+                for i in T.unroll(MMA_N // EPI_N):
+                    reg = T.alloc_local((EPI_N,), acc_type)
+                    reg_wg = reg.view(128, EPI_N,
+                        layout=TileLayout(S[(128, EPI_N) : (1@tid_in_wg, 1)]))
+                    col_st = T.meta_var(wg_id * MMA_N + i * EPI_N)
+                    col_end = T.meta_var(wg_id * MMA_N + i * EPI_N + EPI_N)
+                    Tx.wg.copy_async(reg_wg[:], tmem[:, col_st:col_end])
+                    T.ptx.tcgen05.wait.ld()
+                    Tx.cast(reg_f16[:], reg[:])
+                    Tx.copy(Dsmem[wg_id, warp_id * 32 + lane_id, :], reg_f16[:])
+                    T.ptx.fence.proxy_async("shared::cta")
+                    T.cuda.warpgroup_sync(wg_id + 10)
+                    if warp_id == 0:
+                        if lane_id == 0:
+                            m_st_epi = T.meta_var(
+                                (m_idx * NUM_CONSUMER * CTA_GROUP + wg_id * CTA_GROUP + cbx) * BLK_M)
+                            n_st_epi = T.meta_var(n_idx * MMA_N + i * EPI_N)
+                            Tx.copy_async(
+                                D[m_st_epi:m_st_epi+BLK_M, n_st_epi:n_st_epi+EPI_N],
+                                Dsmem[wg_id, :, :], dispatch="tma")
+                            T.ptx.cp_async.bulk.commit_group()
+                            T.ptx.cp_async.bulk.wait_group(0)
+                    T.cuda.warpgroup_sync(wg_id + 10)
+
+                ld2mma.arrive(wg_id, cta_id=0, pred=True)
+                tile_scheduler.next_tile()
+
+        # --- Cleanup ---
+        T.cuda.cluster_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=CTA_GROUP)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=CTA_GROUP)
+
+    return kernel
+```
+
+**实现说明。**
+
+- 在这个 Step 9 设计中，`mma2ld` 和 `ld2mma` 各自都是一个带 `depth=NUM_CONSUMER` 的共享对象，而不是分开的 per-consumer 对象。slot 0 连接 MMA warp 0 和 Warpgroup 0，slot 1 连接 MMA warp 1 和 Warpgroup 1；MMA 侧用 `warp_id` 索引，writeback 侧用 `wg_id` 索引。
+
+## End-to-End 结果
+
+下表给出了从 naive baseline 到 warp-specialized cluster kernel 的实测里程碑，并列出 cuBLAS reference。参考数字来自 NVIDIA B200，M=N=K=4096，fp16，锁频，1000 次 iteration 计时 benchmark：
+
+| Step | 技术 | 时间 | 加速比 |
+|------|------|------|--------|
+| 1 | 同步加载 + MMA | 70 ms | 1× |
+| 2 | K-loop accumulation | --- | 处理大于单 tile 的 K |
+| 3 | Spatial tiling | 53.6 ms | ~1.3× |
+| 4 | TMA async load | 0.49 ms | ~142× |
+| 5 | Software pipeline | --- | 重叠加载与计算 |
+| 6 | Persistent kernel | --- | L2 cache locality |
+| 7 | Warp specialization | 0.23 ms | ~309× |
+| 8 | 2-CTA cluster | 0.104 ms | ~676× |
+| 9 | Multi-consumer | 0.094 ms | ~744× |
+| --- | cuBLAS（reference） | 0.094 ms | ~744× |
+
+表中的所有时间，包括 70 ms 的 Step 1 baseline，都是在同一个 M=N=K=4096 规模下测得的，所以整条加速链可以端到端比较。这里有必要精确说明 70 ms 指什么，因为它很容易被误读。它 *不是* {ref}`zh_chap_gemm_basics` 中那个单 tile Step-1 kernel 在 4096³ 上运行的结果；那个 kernel 只会计算一个 128×128 tile，也只在小规模下运行。这里的 70 ms 指的是一个 naive full-size baseline：它采用同样的顺序、单 tile 思路，并把它扩展到完整 4096³ 问题。{ref}`zh_chap_gemm_basics` 中的 Steps 1–3 用小规模（128×128 和 256³）引入，是为了让前几段讲解保持简单；这里的 Step 1 和 Step 3 行是它们对应的 full-size benchmark。其余破折号（Steps 2、5、6）表示这些 step 用于展示结构，但没有单独计时。
+
+请把这些数字理解为一次受控条件下的 B200 reference run，而不是 leaderboard 成绩。每个 step 中嵌入的 `{.python .input}` benchmark cell 是 smoke benchmark：适合观察趋势，不适合宣称峰值性能。
+
+几乎所有收益都来自四项技术：
+
+1. **TMA Async Data Movement**：硬件 copy engine 替代 software copy（Step 1 → Step 4 约 142×）。这个 142× 要正确理解：它反映的是从单个 128×128-tile kernel（grid 1×1）一路变成带 K-loop、spatial tiling 和大量 CTA 的完整 tiled-and-parallel kernel，*同时* 引入 TMA；它不是 TMA 单独贡献。要隔离 TMA，就需要比较两个 full-size kernel，且它们只在 copy 机制上不同。
+2. **Software Pipelining + Warp Specialization**：通过给加载和计算各自专用角色来重叠二者（Step 4 → Step 7 约 2.2×）。
+3. **CTA Clusters**：2-SM cooperative MMA 提升跨 CTA 的 B-tile 复用（本 benchmark 中 Step 7 → Step 8 约 2.2×）。
+4. **Multi-Consumer**：两个 MMA warp 带来更高计算密度（Step 8 → Step 9 约 10%）。
+
+如果把这些里程碑画出来，同样四项贡献就形成了从同步 tiled kernel 逐步逼近 cuBLAS reference 的下降轨迹。下图展示了选取的实测点：
+
+![GEMM 优化历程](../img/gemm_perf.png)
+
+注意，越往后收益越小，这背后有结构性原因，不是优化力度变弱。早期 step 处理的是 *memory* 瓶颈（TMA 替代 software copy，cluster 提升 arithmetic intensity），而 70 ms 中的大部分时间确实花在那里，所以这些 step 收益最大。到 Step 8 时，kernel 已经进入 cuBLAS 约 10% 范围内（0.104 vs 0.094 ms），并接近 *compute-bound*，这意味着剩下可隐藏的 memory stall 很少；Step 9 的 multi-consumer overlap 回收了剩余空间中的大部分。接近计算上限时，最终约 10% 的收益正是合理预期：这是一个几乎已经解决的问题所呈现的 diminishing return，而不是优化乏力的信号。
+
+本章构建的所有内容（TMA load、`tcgen05` MMA、TMEM readback，以及 warp-specialized barrier）都会直接进入下一章。Flash Attention 会复用它们，然后通过在两个 MMA 阶段之间插入 online-softmax，而不是简单重复同一个 MMA，把难度再抬高一层。
+
+
+## 练习
+
+1. 如果把 Step 7 中 TMA 和 MMA `PipelineState` 的初始 `phase` 都设为 `0`，会发生什么？画出 deadlock 场景。
+2. Step 8 使用 `cta_group=2` 时，TMA arrive byte count 是 `CTA_GROUP * (BLK_M*BLK_K + BLK_N*BLK_K) * F16_SIZE`。既然每个 CTA 加载自己的数据，为什么还要乘以 `CTA_GROUP`？
+3. Step 9 中，每个 consumer 处理不同的 M 行，但使用同一个 B tile。为什么共享 B（而不是 A）才是正确选择？
+
+**试着让你的 agent 做一遍**：粘贴 Step 7 kernel，让它追踪一个 K-tile 如何经过四个 barrier（`tma2mma`、`mma2tma`、`mma2ld`、`ld2mma`）。对每个 barrier，问谁 wait、谁 arrive、哪个 tile 变得可读，以及之后哪个 buffer 可以复用。
diff --git a/zh/chapter_gemm_async/index.md b/zh/chapter_gemm_async/index.md
new file mode 100644
index 00000000..34d63266
--- /dev/null
+++ b/zh/chapter_gemm_async/index.md
@@ -0,0 +1,765 @@
+(zh_chap_gemm_async)=
+# 用 TMA 对 GEMM 做 Pipelining
+
+:::{admonition} 概览
+:class: overview
+
+- basic GEMM 会轮流执行（copy 一个 tile、compute、copy 下一个 tile），而这些事情本可以同时运行，因此浪费时间。
+- Step 4 切换到 TMA async load，Step 5 对 SMEM 做 double-buffer 并 prefetch（PIPE_DEPTH=2）；完整 load/compute overlap 会在 Step 7 通过 warp specialization 到来，Step 6 则用 tile scheduler 把 kernel 变成 persistent。
+- 目标是在 Tensor Core 咀嚼当前 tile 的同时，load 下一个 tile。
+:::
+
+Tensor Core 是芯片上最昂贵的单元，而上一章中正确的 tiled GEMM 会让它们在大部分时钟周期里闲置。
+kernel 轮流工作：thread 把一个 tile copy 到 shared memory，Tensor Core 处理它，thread copy 下一个 tile，Tensor Core 等待。
+每个 stage 都卡在前一个 stage 上，尽管 load 下一个 tile 和 compute 当前 tile 使用的是完全不同的硬件，本可以同时运行。
+缩小这个差距不需要新的数据路径；tile、layout 和数学都已经正确。需要改变的是工作*何时*发生，以及由*谁*调度。
+本章保持 tile 数据路径完全不变，直接攻击闲置时间。
+
+我们会通过三个增量步骤到达那里，而且在开始前知道目的地会很有帮助。
+Step 4 把 bulk GMEM <-> SMEM transfer 交给 TMA，让专用 copy 硬件移动 tile，而不是 thread。
+Step 5 添加 two-stage software pipeline，让下一个 K tile 在当前 tile 仍在相乘时有地方落脚。
+Step 6 把 launch 重塑成由 tile scheduler 驱动的 persistent kernel，从而摊销 per-tile setup，
+并允许我们选择能让 operand 保持 hot 的 tile 顺序。整个过程中，SMEM、TMEM 和 register layout 都保持上一章留下的样子。
+唯一真正的新思想，是硬件单元之间的异步 handoff：让一个引擎跑在另一个前面，而不是 lockstep 前进。
+
+(zh_chap_tma_async)=
+## Step 4：TMA Async Load
+
+第一步是把 copy 本身移出 critical path。回想 Step 1-3 中 CTA 在做什么：
+它的每个 thread 都在计算地址并发出 load 指令，唯一目的只是把 tile 搬进 SMEM。
+这把 instruction bandwidth 花在管线活上，而不是数学上。Step 4 用 TMA 替换同步 `Tx.copy`：
+单个 thread 发出一条命令，TMA engine 自己完成整个 tile transfer。
+从这里开始，示例运行在完整 M=N=K=4096 尺寸上，而不是 Step 1-3 的小尺寸；
+它们的端到端时间会出现在 {ref}`zh_chap_gemm_advanced` 末尾的 *End-to-End Result* 表中。
+
+> **这一步改变什么：Dispatch**
+> - Scope：不变，一个 warpgroup。
+> - Layout：不变，同样的 SMEM/TMEM/register tile。
+> - Dispatch：GMEM → SMEM load 从同步 `Tx.copy` 移到 TMA engine。
+
+### TMA 发射模式
+
+Step 4 的唯一变化，是用 TMA load 替换同步 tile copy，因此值得仔细看这个 load 如何发射。
+源码改动只有几行，但这些行背后的执行模型在性质上不同。同步 `Tx.copy` 是 CTA thread 用自己的指令亲自做的工作；
+TMA copy 是一个 thread 发出的命令，之后所有移动由 TMA 硬件完成。把二者放在一起看很有价值。
+
+**之前（Step 3）**：全部 128 个 thread 参与 copy，随后 `cta_sync` 让 shared-memory 写入可见：
+```python
+Tx.cta.copy(Asmem[:, :], A[m_st:m_st+BLK_M, i*BLK_K:(i+1)*BLK_K])   # all 128 threads
+Tx.cta.copy(Bsmem[:, :], B[n_st:n_st+BLK_N, i*BLK_K:(i+1)*BLK_K])
+T.cuda.cta_sync()
+```
+
+**之后（Step 4）**：一个 thread 发射 TMA load，mbarrier 追踪硬件 transfer 何时完成：
+```python
+tid = warp_id * 32 + lane_id                 # 0..127 within the warpgroup
+if tid == 0:  # exactly one thread starts TMA
+    Tx.copy_async(Asmem, A[...], dispatch="tma")
+    Tx.copy_async(Bsmem, B[...], dispatch="tma")
+    T.ptx.mbarrier.arrive.expect_tx(tma_bar, byte_count)  # bytes expected from TMA
+T.ptx.mbarrier.try_wait(tma_bar, phase)                  # wait before MMA reads SMEM
+```
+
+注意，load 由 `tid == 0` gate，而不是由 `elect_sync()` gate，这个区别比看上去更重要。
+`elect.sync` 会*每个 warp* 选出一个 active lane，而一个 warpgroup 有四个 warp，
+所以 `elect_sync()` 实际会让四个 thread 进入 load protocol。
+问题在于，这个协议会向 mbarrier 宣布 expected byte count，而且必须恰好宣布一次；
+四次宣布会破坏 count，wait 将永远无法正确 release。通过 warpgroup-wide id 精确选择一个 thread，是避免这个问题的干净方式。
+
+需要诚实说明 speedup 来自哪里。Step 4 仍然在每次 TMA load 后等待，所以还没有让 load 与 compute overlap；
+那是 Step 5 的工作。这里的收益纯粹来自数据移动路径的改变：
+
+- `Tx.copy` 使用 CTA thread 计算地址并发出 load/store 指令。
+- TMA 使用一条发出的命令启动硬件 tile transfer。address generation、coalescing 和 swizzling 由 TMA descriptor 描述，并由 TMA engine 执行。
+
+因此，即使 Step 4 仍然在每次 load 上阻塞，它最终仍然更快。TMA 吸收了 bulk transfer，
+让 CTA thread 不必花 instruction bandwidth 来搬运 tile；仅这项节省就足以改变性能。
+
+### TMA Load 与 Store 同步
+
+我们已经看到 TMA copy 如何发射；故事的另一半是知道它何时完成。切换到 TMA 会同时改变两件事：
+谁启动 copy，以及代码如何知道 copy 已完成。第一点从代码中很明显；第二点容易忽略，搞错后得到的是静默 correctness bug，而不是崩溃。
+使用 `Tx.cta.copy` 时，CTA thread 一起做 copy，后面的 `cta_sync()` 足以说明它完成。
+使用 TMA 时，一个被选中的 thread 发射 `Tx.copy_async(..., dispatch="tma")`，
+engine 按自己的调度执行 transfer，并通过 mbarrier signal completion。
+
+这正是 `cta_sync()` 不再足够的原因。`cta_sync()` 只等待 CTA 自己的 thread，
+也只排序它们的 shared-memory 写入；它对 in-flight TMA transfer 一无所知，所以即使 tile 仍在到达，它也会高兴地返回。
+修复方式是显式化 completion：对于 TMA load，被选中的 thread 先告诉 mbarrier 应期待多少字节，
+然后 CTA 在任何 MMA 触碰 SMEM tile 前等待*那个* mbarrier。下图端到端追踪这个 handshake。
+
+![TMA 异步加载：同步流程](../img/tma_sync_flow.png)
+
+上图隔离了 load 侧 handshake：一个 selected thread launch TMA，mbarrier
+统计 expected bytes，MMA 在读取 SMEM 前等待 release。图中写着
+“Elected Thread”的地方，意思是启动 TMA 的 selected thread；在我们的代码里是 `tid == 0`
+thread，而不是 `elect_sync()` lane。
+
+把 load path 合起来看：selected thread 发出两个 `copy_async` 调用，然后执行 `arrive.expect_tx(total_bytes)`，
+其中 byte count 精确表示 mbarrier 应该等待多少数据。当 engine 移动了这么多字节后，
+匹配的 `mbarrier.try_wait(phase)` 才会 release；只有这时，SMEM tile 才能安全喂给 MMA。
+
+store 侧经过同样硬件，但等待方式不同，所以需要在脑中清楚区分两种协议：
+load 用 mbarrier 和 byte count 追踪 completion，而 store 用 commit group 和 wait group 追踪。
+thread 把 fp16 结果写入 `Dsmem` 并同步后，一个 selected thread 启动
+`Tx.copy_async(D[...], Dsmem, dispatch="tma")`，随后 `cp_async.bulk.commit_group()`
+和 `cp_async.bulk.wait_group(0)` 会阻塞到 store drain。这个 wait 不是可选的：
+在前一次 store 离开之前，`Dsmem` 不能被下一个 tile 复用。
+
+**可以让你的 agent 试试**：追踪一个 K tile 上 Step 4 的 load 和 store 同步。
+识别哪个 thread 启动每个 TMA command，哪个 mbarrier 或 commit group 追踪 completion，
+哪个 wait 保护 MMA 对 `Asmem` 和 `Bsmem` 的读取，哪个 wait 保护 `Dsmem` 的复用。
+为什么这里用 `elect_sync()` 选择 TMA load protocol 的 thread 是错误的？
+
+### 完整 Kernel
+
+完整 kernel 把 TMA load 和 store 折进 Step 3 结构中，并保持该结构其余部分不变。import 与之前相同：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+```
+
+它被包在 `hgemm_v4(M, N, K)` 中，这是我们贯穿使用的模式：
+wrapper 把 shape-dependent constant 和 layout 放在使用它们的 kernel 旁边。
+
+```python
+def hgemm_v4(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    K_TILES = K // BLK_K
+    F16_SIZE = 2
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K))
+    D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_N))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        bx, by = T.cta_id([M // BLK_M, N // BLK_N])
+        wg_id = T.warpgroup_id([1])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+    
+        # --- SMEM allocation (now includes Dsmem for TMA store) ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        tma_bar = pool.alloc((1,), "uint64", align=8)
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout)
+        Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout)
+        pool.commit()
+    
+        # --- Barrier + TMEM init ---
+        if warp_id == 0 and lane_id == 0:
+            T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+            T.ptx.mbarrier.init(tma_bar.ptr_to([0]), 1)
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+    
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+    
+        tmem = T.decl_buffer(
+            (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])
+        )
+    
+        m_st = T.meta_var(bx * BLK_M)
+        n_st = T.meta_var(by * BLK_N)
+        phase_tma: T.int32 = 0
+        phase_mma: T.int32 = 0
+    
+        # --- Inline helpers ---
+        @T.inline
+        def tma_load(k_st):
+            tma_config = T.meta_var({
+                "dispatch": "tma", "cta_group": 1,
+                "mbar": tma_bar.ptr_to([0])
+            })
+            Tx.copy_async(Asmem[:, :],
+                          A[m_st : m_st + BLK_M, k_st : k_st + BLK_K],
+                          **tma_config)
+            Tx.copy_async(Bsmem[:, :],
+                          B[n_st : n_st + BLK_N, k_st : k_st + BLK_K],
+                          **tma_config)
+            T.ptx.mbarrier.arrive.expect_tx(
+                tma_bar.ptr_to([0]),
+                (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE
+            )
+    
+        @T.inline
+        def mma(accum):
+            Tx.gemm_async(
+                tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :],
+                accum=accum, dispatch="tcgen05", cta_group=1
+            )
+            T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+    
+        # --- K-loop with TMA async ---
+        tid = T.meta_var(warp_id * 32 + lane_id)
+        for k in range(K_TILES):
+            k_st = T.meta_var(k * BLK_K)
+    
+            # Single thread issues TMA load
+            if tid == 0:
+                tma_load(k_st)
+    
+            # Wait for TMA to finish; the mbarrier release carries SMEM
+            # visibility to the subsequent MMA, so no extra fence is needed.
+            T.ptx.mbarrier.try_wait(tma_bar.ptr_to([0]), phase_tma)
+    
+            # Single thread issues MMA
+            if tid == 0:
+                mma(accum=k != 0)
+    
+            # Wait for MMA to finish
+            T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+            phase_tma ^= 1
+            phase_mma ^= 1
+    
+        # --- TMA Store Writeback ---
+        Dreg = T.alloc_local((BLK_N,), acc_type)
+        Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+        Dreg_wg = Dreg.view(128, BLK_N,
+                            layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+    
+        # Read TMEM -> registers (async; wait.ld then cta_sync to ensure read completes)
+        Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+        T.ptx.tcgen05.wait.ld()
+        T.cuda.cta_sync()
+        # Cast fp32 -> fp16
+        Tx.cast(Dreg_f16[:], Dreg[:])
+        # Write registers -> Dsmem, flush, then sync
+        Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:])
+        T.ptx.fence.proxy_async("shared::cta")
+        T.cuda.warpgroup_sync(10)
+        # TMA store: Dsmem -> GMEM. One selected thread starts the store and drains the
+        # store group before Dsmem is reused.
+        if tid == 0:
+            Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N],
+                          Dsmem[:, :], dispatch="tma")
+            T.ptx.cp_async.bulk.commit_group()
+            T.ptx.cp_async.bulk.wait_group(0)
+        T.cuda.warpgroup_sync(10)
+    
+        # --- Deallocate TMEM ---
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+### Kernel 中的 TMA 配置
+
+这个 kernel 中几乎所有内容都继承自 Step 3。真正承载 TMA 语义的只有五个配置点，值得逐一认识：
+
+- **TMA config**：`{"dispatch": "tma", "cta_group": 1, "mbar": tma_bar.ptr_to([0])}` 告诉 `Tx.copy_async` 使用 TMA，并通过 `tma_bar` 报告 load completion。
+
+- **Byte count**：`(BLK_M * BLK_K + BLK_N * BLK_K) * 2` 是两个 fp16 operand tile 载入的字节数。`arrive.expect_tx(...)` 把这个 count 交给 mbarrier。
+
+- **mbarrier 初始化**：`init(tma_bar.ptr_to([0]), 1)` 创建 TMA load 使用的 completion barrier。
+
+- **`@T.inline`**：`tma_load(...)` 和 `mma(...)` 是 helper function。它们会在编译期展开到 kernel body 中，并且可以使用周围 kernel 的变量。
+
+- **TMA store 同步**：epilogue 先把 fp16 row 写入 `Dsmem`。`fence.proxy_async` 和 `warpgroup_sync`
+  让这些由 thread 写入的 SMEM 值对 TMA store 路径 ready。随后 store 使用 `commit_group()` 和 `wait_group(0)`
+  等待 SMEM-to-GMEM transfer 完成。
+
+到这里，我们有了正确部件，但节奏仍然错误。Step 4 仍然在启动匹配 MMA 前完成每次 load，
+所以 load 和 multiply 从未真正同时运行；我们费力分开的两个引擎仍在轮流工作。
+下一步会保持 TMA load 和 store 路径完全不变，只重新安排 schedule，让一个 K tile 的 load 可以在另一个 K tile 上的 compute 运行时进行。
+
+(zh_chap_software_pipeline)=
+## Step 5：Software Pipeline（PIPE_DEPTH=2）
+
+既然两个引擎显然是独立的，为什么 Step 4 仍不能让 load 与 compute overlap？
+障碍原来是 storage。只有一对 SMEM tile 时，下一个 load 没地方去：
+它必须等当前 MMA 读完这一对 tile 后才能开始，因为提前开始会覆写仍在使用的数据。
+Step 5 通过 double-buffer shared memory 移除这个 storage conflict。
+single-warpgroup loop 仍然会在 launch 下一次 TMA load 前等待每个 MMA，
+但现在它有不同 stage 可以 prefetch 和 reuse。我们仍然使用完整 M=N=K=4096 尺寸。
+
+> **这一步改变什么：Layout**
+> - Scope：不变，一个 warpgroup。
+> - Layout：单个 SMEM tile pair 变成 `PIPE_DEPTH`-stage ring buffer。
+> - Dispatch：不变，仍是 TMA load 和 `tcgen05` MMA；这一步加入 prefetch 和 stage reuse，完整 load/compute overlap 会在 Step 7 到来。
+
+### Pipeline Walkthrough
+
+当 `PIPE_DEPTH=2` 时，kernel 会分配两个 SMEM stage，让 load path 和 MMA path 拥有不同 slot 可用。
+
+请把下图读成 two-stage buffer 旨在启用的 pipeline 结构，而不是这个 single-warpgroup kernel 的精确执行 trace。
+Step 5 构建 ring buffer，并 prefetch 后续 stage，但主 loop 仍然在发射下一次 TMA load 前等待当前 MMA。
+完整 load/compute overlap 会在 Step 7 到来，届时 warp specialization 会给 TMA 和 MMA 分配不同角色。
+
+![*PIPE_DEPTH=2 的流水线目标调度；这个 single-warpgroup 步骤只做预取，完整重叠会在步骤 7 随 warp 专门化到来*](../img/pipe_depth2.png)
+
+primed 之后，loop 会在两个 stage 之间交替。两次 TMA load 会预先填满两个 stage；
+之后，loop 等待当前 stage，在其上运行 MMA，等待该 MMA 读完这个 stage，
+然后把 `k + PIPE_DEPTH` 的 load launch 到刚刚变得可复用的 stage 中。
+这还不是 concurrent TMA/MMA schedule，但它建立了 Step 7 将拆分到 producer 和 consumer 角色上的 ring-buffer 结构。
+
+具体来说，代码与 Step 4 有四处不同：
+
+1. `Asmem` 和 `Bsmem` 增加一个 leading `PIPE_DEPTH` 维度，因此每个 stage 有自己的 SMEM storage。
+2. `tma_bar` 变成数组，每个 stage 一个 mbarrier。
+3. 在主 K loop 之前，kernel prefetch 前两个 stage。
+4. K loop 使用 `stage = k % PIPE_DEPTH`：等待当前 stage，在其上运行 MMA，然后为 `k + PIPE_DEPTH` 复用该 stage。
+
+### Pipeline 机制
+
+**1. Prefetch**：在主 loop 运行前，我们 load 前 `PIPE_DEPTH` 个 stage，
+让 loop 在第一次 iteration 时就总能发现已有数据在等它：
+```python
+for s in range(min(PIPE_DEPTH, K_TILES)):
+    tma_load(s, s * BLK_K)
+```
+
+**2. Main loop**：对每个 K tile，我们等待它的 stage ready，在其上运行 MMA，
+然后立即通过 launch 提前 `PIPE_DEPTH` 个 tile 的 load，让这个现在空闲的 stage 重新工作：
+```python
+stage = k % PIPE_DEPTH
+wait(tma_bar[stage], phase_tma)
+mma(stage, accum)
+wait(mma_bar[0], phase_mma)
+phase_mma ^= 1
+tma_load(stage, next_k * BLK_K)
+```
+
+**3. Phase management**：这是容易绊倒人的部分，但规则比第一眼看上去更简单。
+每个 barrier 的 phase-flip 规则直接来自该 barrier 有多少个 slot，这就是两个 barrier 以不同节奏翻转的原因。
+MMA accumulator 位于一个 TMEM slot 中，所以 `mma_bar` 是单个 barrier（`mma_bar.ptr_to([0])`），每个 iteration 都会重新访问；
+每个 iteration 都重新访问的 barrier，必须每个 iteration 翻转 phase。
+TMA barrier 的故事不同：它们形成一个 `PIPE_DEPTH` 元素数组，每个 stage 一个 barrier；
+任意给定 stage 的 barrier 只有在 ring 中绕一圈后才会再次出现。因此，`phase_tma` 只在 stage index wrap 回 0 时翻转：
+```python
+if stage == PIPE_DEPTH - 1:
+    phase_tma ^= 1
+```
+
+**可以让你的 agent 试试**：在 `PIPE_DEPTH=2` 且 `K_TILES=5` 时，让它追踪主 loop。
+对每个 `k`，列出 `stage`、传给 wait 的 `phase_tma` 和 `phase_mma` 值，以及是否发射新的 prefetch。
+`phase_tma` 究竟在哪里翻转？为什么最后两个 iteration 没有 prefetch？
+
+### 完整 Kernel
+
+完整 kernel 原样保留 Step 4 的 TMA load 和 store 路径，然后用刚才描述的 staged buffer 和 phase logic 包住它。
+import 不变：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+```
+
+它被包在 `hgemm_v5(M, N, K)` 中。`PIPE_DEPTH=2` 常量设置 pipeline stage 数量
+（这里是两个，正好是 double buffering）：
+
+```python
+PIPE_DEPTH = 2
+
+def hgemm_v5(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+    F16_SIZE = 2
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    K_TILES = K // BLK_K
+
+    # Double-buffered layouts: first dimension is pipeline stage
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                  (PIPE_DEPTH, BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                  (PIPE_DEPTH, BLK_N, BLK_K))
+    D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                  (BLK_M, BLK_N))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        bx, by = T.cta_id([M // BLK_M, N // BLK_N])
+        wg_id = T.warpgroup_id([1])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        # --- SMEM allocation ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        # Double-buffered TMA barriers (one per stage), single MMA barrier
+        tma_bar = pool.alloc((PIPE_DEPTH,), "uint64", align=8)
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout)
+        Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout)
+        pool.commit()
+
+        # Initialize barriers: PIPE_DEPTH for TMA, 1 for MMA
+        if warp_id == 0:
+            if lane_id == 0:
+                T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+                for s in range(PIPE_DEPTH):
+                    T.ptx.mbarrier.init(tma_bar.ptr_to([s]), 1)
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+            (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])
+        )
+
+        m_st = T.meta_var(bx * BLK_M)
+        n_st = T.meta_var(by * BLK_N)
+        phase_tma: T.int32 = 0
+        phase_mma: T.int32 = 0
+
+        @T.inline
+        def tma_load(stage, k_offset):
+            tma_config = T.meta_var({
+                "dispatch": "tma", "cta_group": 1,
+                "mbar": tma_bar.ptr_to([stage])
+            })
+            Tx.copy_async(Asmem[stage, :, :],
+                          A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K],
+                          **tma_config)
+            Tx.copy_async(Bsmem[stage, :, :],
+                          B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K],
+                          **tma_config)
+            T.ptx.mbarrier.arrive.expect_tx(
+                tma_bar.ptr_to([stage]),
+                (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE)
+
+        @T.inline
+        def mma(stage, accum):
+            Tx.gemm_async(tmem[:, :BLK_N], Asmem[stage, :, :], Bsmem[stage, :, :],
+                          accum=accum, dispatch="tcgen05", cta_group=1)
+            T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+
+        tid = T.meta_var(warp_id * 32 + lane_id)
+
+        # === Prefetch: load first PIPE_DEPTH stages ===
+        if tid == 0:
+            for s in range(min(PIPE_DEPTH, K_TILES)):
+                tma_load(s, s * BLK_K)
+
+        # === Main loop ===
+        for k in range(K_TILES):
+            stage = k % PIPE_DEPTH
+
+            # Wait for TMA to finish loading this stage
+            T.ptx.mbarrier.try_wait(tma_bar.ptr_to([stage]), phase_tma)
+
+            # MMA on this stage's data
+            if tid == 0:
+                mma(stage, accum=(k != 0))
+
+            T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+            phase_mma ^= 1
+
+            # Issue next prefetch load (k + PIPE_DEPTH)
+            next_k = k + PIPE_DEPTH
+            if next_k < K_TILES:
+                if tid == 0:
+                    tma_load(stage, next_k * BLK_K)
+
+            # TMA phase flips when stage wraps around
+            if stage == PIPE_DEPTH - 1:
+                phase_tma ^= 1
+
+        # === TMA Store Writeback: TMEM -> RF -> Dsmem -> TMA -> GMEM ===
+        Dreg = T.alloc_local((BLK_N,), acc_type)
+        Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+        Dreg_wg = Dreg.view(128, BLK_N,
+                            layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+        Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+        T.ptx.tcgen05.wait.ld()
+        T.cuda.cta_sync()
+        Tx.cast(Dreg_f16[:], Dreg[:])
+        Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:])
+        T.ptx.fence.proxy_async("shared::cta")
+        T.cuda.warpgroup_sync(10)
+        if tid == 0:
+            Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N],
+                          Dsmem[:, :], dispatch="tma")
+            T.ptx.cp_async.bulk.commit_group()
+            T.ptx.cp_async.bulk.wait_group(0)
+        T.cuda.warpgroup_sync(10)
+
+        # Deallocate TMEM
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+(zh_chap_persistent_kernel)=
+## Step 6：Persistent Kernel + Tile Scheduler
+
+到目前为止，一切都在优化单个 tile 内部的工作。Step 6 改变问题的尺度，开始跨 tile 优化。
+
+Step 5 为每个 128 x 128 output tile 启动一个 CTA。对于 4096 x 4096 输出，
+这意味着 1024 个独立 CTA，每个 CTA 都支付自己的 setup cost，并在自己的 tile 完成后立即消失。
+
+Step 6 改为启动固定 CTA 池，然后让每个 CTA 依次处理多个 tile。
+这带来两点收益：setup work 被摊销到多个 tile 上；tile assignment 移入 kernel 内部，
+在那里 scheduler 可以选择一种能复用 operand 的顺序。我们仍保持完整 M=N=K=4096 尺寸。
+
+> **这一步改变什么：Scope**
+> - Scope：固定的 persistent CTA 池，每个 CTA 通过 scheduler 循环处理多个 output tile。
+> - Layout：不变，同样的 per-tile SMEM/TMEM/register 路径。
+> - Dispatch：不变。
+
+### Persistent Scheduling
+
+persistent kernel 的定义性想法是：按硬件规模，而不是按问题规模来设置 grid。
+它启动 `SM_COUNT` 个 CTA，大致每个 SM 一个，不管实际有多少 output tile，目标是让每个 SM 持续保持占用。
+我们有意说“大致”：精确 1:1 residency 并不保证，因为它取决于 occupancy，也取决于硬件如何调度 CTA。
+
+在这里目标的 B200 上，`SM_COUNT=148`。这 148 个 CTA 中的每一个，都会循环处理
+`ClusterPersistentScheduler2D` 分配给它的 tile。
+
+第一项收益是摊销。TMEM allocation、barrier initialization 和 scheduler state 现在每个 CTA 只发生一次，
+并在该 CTA 处理的大约 7 个 tile 间复用，而不是在一次性 CTA 上重复 1024 次。
+
+第二项收益来自 scheduler 选择的顺序。设置 `l2_group_size=8` 会把相邻 tile 分组，
+因此共享同一 row band 的 tile 会复用相同 A row-tile，共享同一 column band 的 tile 会复用相同 B tile。
+背靠背运行这些 tile，可以让 operand 在 L2 中保持 hot，而不是从 HBM 重新获取。
+这正是 Step 3 留在桌面上的复用。
+
+```python
+bx = T.cta_id([SM_COUNT])  # 1D grid, one CTA per SM
+
+tile_scheduler = ClusterPersistentScheduler2D(
+    "ts",
+    num_m_tiles=M // BLK_M,
+    num_n_tiles=N // BLK_N,
+    l2_group_size=8,       # Group 8 nearby tiles together
+    num_clusters=SM_COUNT
+)
+tile_scheduler.init(bx)
+```
+
+跨 tile 循环带来一个容易漏掉的 correctness 后果。每个 tile 都运行自己的全新 K-loop，
+这意味着它的 barrier phase 必须从已知状态开始。Step 5 中，一个 CTA 恰好处理一个 tile，
+所以只初始化一次 `phase_tma` 和 `phase_mma` 完全没问题。
+Step 6 中，这些 initializer 必须移到 `while tile_scheduler.valid()` loop *内部*，
+让每个 tile 都以匹配自己 TMA 和 MMA 工作的 phase state 开始，而不是继承前一个 tile 偶然留下的状态：
+
+```python
+while tile_scheduler.valid():
+    phase_tma: T.int32 = 0
+    phase_mma: T.int32 = 0
+    ...
+```
+
+### 完整 Kernel
+
+从结构上说，这个 kernel 只是把 Step 5 的 pipeline 包在 tile-level outer loop 里。
+唯一新的依赖是 scheduler 本身，我们把它和其他依赖一起 import：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+from tvm.tirx.lang.tile_scheduler import ClusterPersistentScheduler2D
+```
+
+grid dimension 现在只是 `SM_COUNT`，而不是 `(M//BLK_M, N//BLK_N)`；
+`ClusterPersistentScheduler2D` 接管了给每个 CTA 分配 tile 的工作：
+
+```python
+SM_COUNT = 148  # Number of SMs on NVIDIA B200 GPU
+PIPE_DEPTH = 2
+
+def hgemm_v6(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+    F16_SIZE = 2
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    K_TILES = K // BLK_K
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                  (PIPE_DEPTH, BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                  (PIPE_DEPTH, BLK_N, BLK_K))
+    D_layout = tma_shared_layout(d_type, SwizzleMode.SWIZZLE_128B_ATOM,
+                                  (BLK_M, BLK_N))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        # 1D grid: one CTA per SM (not a 2D grid anymore!)
+        bx = T.cta_id([SM_COUNT])
+        wg_id = T.warpgroup_id([1])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        # --- SMEM allocation (same as Step 5) ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        tma_bar = pool.alloc((PIPE_DEPTH,), "uint64", align=8)
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((PIPE_DEPTH, BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((PIPE_DEPTH, BLK_N, BLK_K), b_type, layout=B_layout)
+        Dsmem = pool.alloc((BLK_M, BLK_N), d_type, layout=D_layout)
+        pool.commit()
+
+        # --- Barrier + TMEM init (same as Step 5) ---
+        if warp_id == 0 and lane_id == 0:
+            T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+            for s in range(PIPE_DEPTH):
+                T.ptx.mbarrier.init(tma_bar.ptr_to([s]), 1)
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+            (128, 512), acc_type, scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])
+        )
+
+        # Tile scheduler: assigns tiles to CTAs in L2-friendly order
+        tile_scheduler = ClusterPersistentScheduler2D(
+            "ts",
+            num_m_tiles=M // BLK_M,
+            num_n_tiles=N // BLK_N,
+            l2_group_size=8,
+            num_clusters=SM_COUNT
+        )
+        tile_scheduler.init(bx)
+
+        tid = T.meta_var(warp_id * 32 + lane_id)
+
+        @T.inline
+        def tma_load(stage, k_offset, m_st, n_st):
+            tma_config = T.meta_var({
+                "dispatch": "tma", "cta_group": 1,
+                "mbar": tma_bar.ptr_to([stage])
+            })
+            Tx.copy_async(Asmem[stage, :, :],
+                          A[m_st:m_st+BLK_M, k_offset:k_offset+BLK_K],
+                          **tma_config)
+            Tx.copy_async(Bsmem[stage, :, :],
+                          B[n_st:n_st+BLK_N, k_offset:k_offset+BLK_K],
+                          **tma_config)
+            T.ptx.mbarrier.arrive.expect_tx(
+                tma_bar.ptr_to([stage]),
+                (BLK_M * BLK_K + BLK_N * BLK_K) * F16_SIZE)
+
+        @T.inline
+        def mma(stage, accum):
+            Tx.gemm_async(tmem[:, :BLK_N], Asmem[stage, :, :], Bsmem[stage, :, :],
+                          accum=accum, dispatch="tcgen05", cta_group=1)
+            T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+
+        # === Outer loop: iterate over tiles ===
+        while tile_scheduler.valid():
+            # Get current tile position from scheduler
+            m_st = T.meta_var(tile_scheduler.m_idx * BLK_M)
+            n_st = T.meta_var(tile_scheduler.n_idx * BLK_N)
+
+            # === Inner loop: same pipeline as Step 5 ===
+            phase_tma: T.int32 = 0
+            phase_mma: T.int32 = 0
+
+            # Prefetch first PIPE_DEPTH stages
+            if tid == 0:
+                for s in range(min(PIPE_DEPTH, K_TILES)):
+                    tma_load(s, s * BLK_K, m_st, n_st)
+
+            # Main K-loop
+            for k in range(K_TILES):
+                stage = k % PIPE_DEPTH
+                T.ptx.mbarrier.try_wait(tma_bar.ptr_to([stage]), phase_tma)
+                if tid == 0:
+                    mma(stage, accum=(k != 0))
+                T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+                phase_mma ^= 1
+                next_k = k + PIPE_DEPTH
+                if next_k < K_TILES:
+                    if tid == 0:
+                        tma_load(stage, next_k * BLK_K, m_st, n_st)
+                if stage == PIPE_DEPTH - 1:
+                    phase_tma ^= 1
+
+            # === TMA Store Writeback: TMEM -> RF -> Dsmem -> TMA -> GMEM ===
+            Dreg = T.alloc_local((BLK_N,), acc_type)
+            Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+            Dreg_wg = Dreg.view(128, BLK_N,
+                                layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+            Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+            T.ptx.tcgen05.wait.ld()
+            T.cuda.cta_sync()
+            Tx.cast(Dreg_f16[:], Dreg[:])
+            Tx.copy(Dsmem[warp_id * 32 + lane_id, 0:BLK_N], Dreg_f16[:])
+            T.ptx.fence.proxy_async("shared::cta")
+            T.cuda.warpgroup_sync(10)
+            if tid == 0:
+                Tx.copy_async(D[m_st : m_st + BLK_M, n_st : n_st + BLK_N],
+                              Dsmem[:, :], dispatch="tma")
+                T.ptx.cp_async.bulk.commit_group()
+                T.ptx.cp_async.bulk.wait_group(0)
+            T.cuda.warpgroup_sync(10)
+
+            T.cuda.cta_sync()
+            tile_scheduler.next_tile()  # Move to next tile
+
+        # Deallocate TMEM
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+## 练习
+
+1. 在 Step 4 中，`arrive.expect_tx` 使用 `(BLK_M * BLK_K + BLK_N * BLK_K) * 2` 字节。如果这个 byte count 太小或太大，mbarrier 会等待什么？
+2. 在 Step 5 中，为什么每个 SMEM stage 都需要自己的 TMA barrier，而不是两个 stage 共享一个 `tma_bar`？
+3. 在 Step 6 中，`BLK_M=BLK_N=128` 时，一个 4096 x 4096 输出有多少个 output tile？当 `SM_COUNT=148` 时，每个 persistent CTA 平均处理多少个 tile？
diff --git a/zh/chapter_gemm_basics/index.md b/zh/chapter_gemm_basics/index.md
new file mode 100644
index 00000000..39d2e8b2
--- /dev/null
+++ b/zh/chapter_gemm_basics/index.md
@@ -0,0 +1,693 @@
+(zh_chap_gemm_basics)=
+# 构建 Tiled GEMM
+
+:::{admonition} 概览
+:class: overview
+
+- 从单个 output tile 开始，用 TIRx tile primitives 构建一个正确的 tiled GEMM。
+- Step 1 是 single-tile GEMM，Step 2 添加 K-loop accumulation，Step 3 针对完整矩阵跨 CTA 做 spatial tiling。
+- 正确性优先；性能是后两章的任务。
+:::
+
+GEMM 是整本书围绕的工作负载。它位于 linear layer、attention projection 和 convolution 之下，
+而这些操作主导着 GPU 时间。因此，正确 GEMM 和快速 GEMM 之间的差距，
+就是让芯片大部分闲置和把芯片打满之间的差距。
+
+这个差距太大，无法一步跨过。一个能打满硬件的 kernel 会让你同时 debug memory movement、accumulation、tiling 和 Tensor Core scheduling，
+而且没有可信对象可供比较。更安全的路径是从能产生正确答案的最小 kernel 开始，然后一次增加一个决策。
+
+本章会写出第一个正确的 tiled GEMM。前面章节从抽象层面介绍了 TIRx 的 scope / layout / dispatch 模型；
+这里我们把它应用到真实 kernel。我们从一个 128 x 128 output tile 开始，
+逐步把它扩展成能处理完整矩阵的 kernel：先加入 K 维度 accumulation，再跨多个 CTA 加入 spatial tiling。
+
+这是三章 GEMM 优化路径中的第一章，三章会端到端走完同一条路径。本章构建一个正确的 tiled kernel，并到此为止。
+下一章（{ref}`zh_chap_gemm_async`）会用 TMA 替换 thread copy，并通过 pipelining 让数据移动和计算 overlap；
+{ref}`zh_chap_gemm_advanced` 则进一步加入 warp specialization 和 CTA cluster。
+每章都建立在前一章之上，所以 kernel 会逐步积累功能，而不是从头再来。
+
+把每一步读成对同一个三项 contract 的修改会很有帮助：哪个 **scope** 运行操作，operand tile 使用哪个 **layout**，
+以及由哪条 **dispatch** 路径执行。大多数步骤都有一个主要变化，所以我们会用一个小卡片开头，点明这个变化，
+并指出为了安全复用而需要的同步细节。Step 1 建立后续路径要不断修改的 baseline。
+
+## GEMM
+
+GEMM 是 dense matrix multiply，位于 linear layer、attention projection 和许多 convolution 实现之下；
+因此，快速 GEMM kernel 几乎处处都有回报。本教程中的例子使用 $D = A B^{\top}$：
+
+- $A$ 的 shape 是 $M \times K$。
+- $B$ 的 shape 是 $N \times K$。
+- $D$ 的 shape 是 $M \times N$。
+- $D[m,n] = \sum_k A[m,k] \cdot B[n,k]$.
+
+transpose 不是我们选择额外执行的操作；它来自数据的存储方式。
+这些例子把 $B$ 保持为 $N$ 行、每行长度 $K$，这通常也是 linear-layer weight 的 layout。
+因此，沿 $K$ contract 会自然读取 $B^{\top}$，不需要任何重排。
+
+整个教程中，我们用 TFLOPS 吞吐来衡量 kernel：把每次 multiply-add 计为两个 floating-point operation，
+再除以 wall-clock time：
+
+$$\text{TFLOPS} = \frac{2 \times M \times N \times K}{t_{\text{seconds}} \times 10^{12}}$$
+
+### GEMM 数据路径
+
+本教程中的每个优化最终都落到数据住在哪里、如何移动，所以在写代码前值得先把这条路径画出来。
+从核心上说，Blackwell GEMM kernel 围绕两类活动组织：在不同内存之间移动 tile，以及在 tile 上计算。
+下图追踪一个 tile 从输入到输出过程中接触的每个内存空间：
+
+![*内存数据流*](../img/memory_dataflow.png)
+
+上图展示了 baseline 路径：后续每个优化都会修改它，但不会替换它。
+从左到右读：operand tile 先从 GMEM 移到 SMEM；随后 `tcgen05.mma` 消费 SMEM operand，
+并把 accumulator 写入 TMEM；最后 epilogue 先把 TMEM 读回寄存器，再把结果存到 GMEM。
+请记住这条链路，因为下面每一步只会改变某一跳*如何*发生，而不会改变这些跳本身。
+
+## 优化路径
+
+上面这条朴素数据路径足以得到正确答案，但会让大部分硬件闲置。
+教程剩余部分会一次加入一个 Blackwell 特性来缩小这个差距，而每个特性都通过 TIRx tile primitive 表达。
+我们将依次经过这些特性：
+
+- **TMA async movement** 通过 Blackwell 的硬件 copy path 移动 GMEM <-> SMEM tile，并用 barrier 追踪 completion。
+- **Software pipelining** 使用多个 SMEM stage，让下一个 K tile 的数据移动可以与当前 tile 上的 Tensor Core 计算 overlap。
+- **Persistent scheduling** 保持固定 CTA 池，每个 CTA 通过 tile scheduler 处理许多 output tile，而不是每个 tile 启动一个 CTA。
+- **Warp specialization** 把 producer、MMA consumer 和 writeback 角色拆到独立 warpgroup 上。
+- **CTA clusters** 让两个 CTA 协作处理一个更大的 Blackwell MMA tile。
+- **Multi-consumer execution** 使用多个 consumer warpgroup 同时计算 tile 的不同部分，提高 compute density。
+
+---
+
+(zh_chap_single_tile)=
+## Step 1：顺序 Single-Tile GEMM
+
+仍然能走完整硬件路径的最简单 GEMM，是计算单个 output tile 的 GEMM。所以我们从这里开始。
+Step 1 计算一个 128 x 128 output tile，K = 64；它足够小，不需要任何 loop，
+而数据路径的每个部分都恰好出现一次。由于没有重复，我们可以在需要推理 loop 之前，先孤立地观察每一跳。
+
+> **这一步建立什么：baseline**
+> - Scope：一个包含 128 个 thread 的单个 warpgroup 按顺序走完整条路径，一个 stage 接一个 stage。
+> - Layout：A 和 B tile 位于 SMEM，accumulator 位于 TMEM，结果通过寄存器 stage 出去。
+> - Dispatch：同步 `Tx.copy` 执行 load，`tcgen05` 执行 MMA。
+
+### Single-Tile Dataflow
+
+baseline contract 固定后，接下来要确定的是一个 tile 穿过它的顺序。
+第一个 kernel 会精确走一次核心 GEMM 数据路径，也就是数据流图中的同一条 GMEM -> SMEM -> TMEM -> registers -> GMEM 链，
+外面没有任何 loop。它分配工作内存、载入 operand、计算乘积、写回结果，并清理自身：
+
+1. **分配**：SMEM（pool allocator）、TMEM（`tcgen05.alloc`）、mbarrier
+2. **加载**：全部 128 个 thread 协作地把 A 和 B tile 从 GMEM copy 到 SMEM（同步 `Tx.copy`）
+3. **计算**：单个 elected thread 发射 `Tx.gemm_async` + `tcgen05.commit`；所有 thread 等待 mbarrier
+4. **写回**：warpgroup 读取 TMEM → registers；每个 thread 把 fp32 cast 为 fp16 并写到 GMEM
+5. **释放**：释放 TMEM
+
+### 第一个 Kernel 的四个部分
+
+完整 kernel 只有几十行，但分块阅读更容易消化。我们会分四部分阅读它：
+memory allocation、synchronous load、MMA dispatch 和 writeback，之后再把它们组装成一个 kernel。
+过程中出现的 API 名称，是第二部分介绍的 TIRx tile-primitive 词汇（{ref}`zh_chap_tirx_primer`，{ref}`zh_chap_tirx_layout_api`）。
+
+**Memory allocation。** kernel 首先为 operand 切出 shared memory，同时为 TMEM address 和 mbarrier 留出 slot：
+
+```python
+pool = T.SMEMPool()
+tmem_addr = pool.alloc((1,), "uint32")           # TMEM address (4 bytes)
+mma_bar = pool.alloc((1,), "uint64", align=8)    # mbarrier (8 bytes)
+pool.move_base_to(1024)                           # Skip to offset 1024
+Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout)  # 128×64 fp16
+Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout)  # 128×64 fp16
+pool.commit()
+```
+
+这里有两个细节值得停一下。`pool.move_base_to(1024)` 把 Asmem 和 Bsmem 推到 offset 1024，
+把低地址留给前面那些小 metadata，让体积较大的 operand tile 落在干净边界上。
+另外，`layout=A_layout` 会向 `tma_shared_layout` 请求一个 swizzled SMEM placement，
+它能被 TMA 和 `tcgen05.mma` 直接读取，正是第二部分描述的那种 layout-as-contract 义务。
+
+**Synchronous load。** buffer 到位后，operand 仍然必须到达 SMEM。在第一个版本中，我们让 CTA 自己的 thread 完成 copy：
+
+```python
+Tx.cta.copy(Asmem[:, :], A[:, :])
+Tx.cta.copy(Bsmem[:, :], B[:, :])
+T.cuda.cta_sync()
+```
+
+因为这里只存在一个 tile（M=N=128，K=64），copy 整个 A 和 B 就是全部 load。
+`Tx.cta.copy(...)` 让 CTA 协作完成这次 copy，每个 thread 负责自己的数据切片。
+后面的 `T.cuda.cta_sync()` 一举两得：它等待每个 thread 完成，同时发布它们的 shared-memory 写入；
+这样后续 MMA 读取 `Asmem` 和 `Bsmem` 时，看到的是完整 tile，而不是半填充 buffer。
+这种 thread-driven copy 也是我们首先要替换的东西；下一章（{ref}`zh_chap_gemm_async`）会把它换成 TMA。
+
+**MMA dispatch。** operand 现在位于 SMEM 中，我们可以发射 MMA，并且从单个 elected thread 发射：
+
+```python
+if warp_id == 0:
+    if T.ptx.elect_sync():
+        Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :],
+                      accum=False, dispatch="tcgen05", cta_group=1)
+        T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+```
+
+两个嵌套 guard 分两步缩小 issuer。外层 `if warp_id == 0` 只保留 warpgroup 中的 warp 0；
+内层 `if T.ptx.elect_sync():` 随后在该 warp 中选出一个 active lane。
+二者合起来只留下一个 thread 来运行 `Tx.gemm_async` 和 `tcgen05.commit`。
+
+这里值得说清楚这个单个 thread 意味着什么、不意味着什么，因为自然读法很容易误导。
+单个 issuing thread *并不*意味着单线程乘法。计算仍然是完整 tile-level MMA：
+硬件会为 SMEM operand layout 和 TMEM accumulator layout 所描述的 tile 执行 cooperative multiply。
+关键在于 `Tx.gemm_async` 是一个 *tile operation*，不是一条硬件指令。
+K = 64 的 tile 比硬件 MMA K-atom（`MMA_K = 16`）更宽，因此这个 tile op 会 lower 成沿 K 方向前进的一小段 raw `tcgen05.mma` 指令，
+warpgroup 会协作驱动其中每一条。只有一个 thread 发射 tile op 的原因是，每个底层 `tcgen05.mma`
+本身就是一个 *single-instruction* cooperative op：一次 launch 驱动 tile MMA 的那个 K-atom。
+如果全部 128 个 thread 都发射这个序列，相同工作只会被重复 launch 128 次。
+最后，`accum=False` flag 告诉 MMA 覆写 TMEM destination，而不是加到其中；
+这正是我们这里想要的，因为没有先前 partial sum 需要扩展。
+
+**Writeback。** 乘积现在位于 TMEM 中，但 caller 想要它以 fp16 回到 GMEM。
+因此 epilogue 必须通过寄存器把结果带下来，并在途中 cast：
+
+```python
+Dreg = T.alloc_local((BLK_N,), acc_type)        # per-thread fp32 register row
+Dreg_f16 = T.alloc_local((BLK_N,), d_type)      # same row, cast to fp16
+Dreg_wg = Dreg.view(128, BLK_N, layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+T.ptx.tcgen05.wait.ld()
+Tx.cast(Dreg_f16[:], Dreg[:])
+m_thr = T.meta_var(m_st + warp_id * 32 + lane_id)
+Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:])
+```
+
+MMA 会在 TMEM 中留下一个 128 x 128 fp32 accumulator tile。fp32 是有意选择的：
+GEMM 会沿 K 累加许多乘积，用更高精度保存 running sum 可以压低本来会积累的舍入误差。
+但 `D` 是 fp16，所以这些值不能直接写出。它们先落入寄存器，在那里缩窄为 fp16，然后才到达 GMEM。
+
+两个 register buffer 扮演不同角色。`Dreg` 是每个 thread 的 `BLK_N` 元素 buffer，
+而 `Dreg_wg` 是同一批寄存器在选定 layout 下的 warpgroup-wide *view*：
+
+```python
+TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)])
+```
+
+这个 layout 把 tile 的第一维映射到 warpgroup 的 thread：thread 0 拥有 row 0，thread 1 拥有 row 1，
+依此类推直到 row 127。第二维保留在每个 thread 自己的 register buffer 中，
+所以一个 thread 持有自己那一行的所有 column。warpgroup 中有 128 个 thread，tile 中有 128 行，
+因此 128 x 128 输出正好分成每个 thread 一行。
+
+在这个 view 下读出 accumulator，正是 `Tx.wg.copy_async(Dreg_wg, tmem)` 所做的事；
+它会 lower 到 Blackwell TMEM load path，即 `tcgen05.ld`。因为这个 load 是异步的，
+`T.ptx.tcgen05.wait.ld()` 必须在任何 thread 触碰 `Dreg` 前完成；否则 thread 可能读取 load 还没有填好的寄存器。
+
+wait 返回后，每个 thread 私有的 `Dreg[:]` 都持有自己那一条逻辑 output row 的 fp32 值。
+thread 在 `Dreg_f16` 中把这些值缩窄为 fp16，并算出自己负责哪个 global row：
+
+```python
+m_thr = T.meta_var(m_st + warp_id * 32 + lane_id)
+```
+
+然后写入 `D[m_thr, n_st:n_st + BLK_N]`。这些 row 干净地分布在四个 warp 上：
+warp 0 写 row 0-31，warp 1 写 row 32-63，warp 2 写 row 64-95，warp 3 写 row 96-127。
+
+### 完整 Kernel
+
+现在把四个部分重新拼成一个可运行 kernel（M=N=128，K=64）。先是 import：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+```
+
+kernel 包在后续步骤也会使用的 `hgemm_vX(M, N, K)` 风格中。
+Step 1 使用 `M=N=128, K=64` 运行，因此 launch 恰好包含一个 output tile：
+
+```python
+def hgemm_v1(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    # MMA_M/MMA_N/MMA_K document the underlying hardware MMA tile; they are not
+    # passed to gemm_async (which derives the MMA shape from the operand and
+    # accumulator tiles), so the later steps omit them.
+    MMA_M, MMA_N, MMA_K = 128, 128, 16
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        # Step 1 is a single-tile kernel: M = BLK_M and N = BLK_N, so the grid
+        # is 1x1. Starting with a 1x1 grid keeps the per-CTA tile offsets
+        # (m_st, n_st) trivially zero; Steps 3+ generalise this to larger M / N.
+        bx, by = T.cta_id([M // BLK_M, N // BLK_N])
+        wg_id = T.warpgroup_id([1])      # single warpgroup, so wg_id is always 0 (unused below)
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+    
+        # --- SMEM allocation ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout)
+        pool.commit()
+    
+        # --- Barrier + TMEM init (warp 0 only) ---
+        if warp_id == 0:
+            if lane_id == 0:
+                T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+    
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+    
+        tmem = T.decl_buffer(
+            (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])
+        )
+    
+        m_st = T.meta_var(bx * BLK_M)
+        n_st = T.meta_var(by * BLK_N)
+        phase_mma: T.int32 = 0
+    
+        # --- Load: all threads copy global -> shared (synchronous).
+        # With M=BLK_M and N=BLK_N the slices below cover the full matrices;
+        # the slice form is kept so the diff to Step 3 (multi-tile) is minimal.
+        Tx.cta.copy(Asmem[:, :], A[m_st:m_st + BLK_M, :])
+        Tx.cta.copy(Bsmem[:, :], B[n_st:n_st + BLK_N, :])
+        T.cuda.cta_sync()
+    
+        # --- Compute: single elected thread issues MMA ---
+        if warp_id == 0:
+            if T.ptx.elect_sync():
+                Tx.gemm_async(
+                    tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :],
+                    accum=False, dispatch="tcgen05", cta_group=1
+                )
+                T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+    
+        T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+    
+        # --- Writeback: TMEM -> RF -> GMEM ---
+        Dreg = T.alloc_local((BLK_N,), acc_type)
+        Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+        Dreg_wg = Dreg.view(128, BLK_N,
+                            layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+        Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+        T.ptx.tcgen05.wait.ld()
+        Tx.cast(Dreg_f16[:], Dreg[:])
+        m_thr = T.meta_var(m_st + warp_id * 32 + lane_id)
+        Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:])
+    
+        # --- Deallocate TMEM ---
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+后续每个 GEMM step 都会用同样方式编译、运行和检查自己，所以我们只在这里把这套 scaffolding 完整写一次，
+之后只展示 kernel。要运行后续 step，把下面的 `hgemm_vX` 和对应 problem size 换掉即可。
+有一点需要记住：每个全新的 Python session 只编译一个 step，尝试另一个 step 前请重启，
+因为这些例子会复用内部名称，而编译器持有 per-session state。
+
+```python
+import torch
+
+target = tvm.target.Target("cuda")
+device = torch.device('cuda')  # gpu(0)
+
+M, N, K = 128, 128, 64
+kernel = hgemm_v1(M, N, K)
+with target:
+    ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+
+torch.cuda.empty_cache()
+torch.cuda.synchronize()
+A_tensor = torch.randn(M, K, dtype=torch.float16, device=device)
+B_tensor = torch.randn(N, K, dtype=torch.float16, device=device)
+D_tensor = torch.zeros(M, N, dtype=torch.float16, device=device)
+
+# ex.mod(...) takes torch tensors directly, the same call form used in every chapter.
+ex.mod(A_tensor, B_tensor, D_tensor)
+
+D_ref = (A_tensor.float() @ B_tensor.float().T).half()
+max_err = float((D_tensor - D_ref).abs().max())
+print(f"Max error vs torch reference: {max_err:.6f}")
+# Relative tolerance, like the warp-specialization and Flash Attention cells:
+# output magnitude grows with K, so a fixed absolute bound would fail at larger K.
+torch.testing.assert_close(D_tensor, D_ref, rtol=2e-2, atol=1e-2)
+print("PASS")
+
+# Optional timing for larger kernels.
+ITERS = 10
+for _ in range(3):
+    ex.mod(A_tensor, B_tensor, D_tensor)
+torch.cuda.synchronize()
+start = torch.cuda.Event(enable_timing=True)
+end = torch.cuda.Event(enable_timing=True)
+start.record()
+for _ in range(ITERS):
+    ex.mod(A_tensor, B_tensor, D_tensor)
+end.record()
+torch.cuda.synchronize()
+ms = start.elapsed_time(end) / ITERS
+tflops = 2 * M * N * K / ms / 1e9
+print(f"Performance: {ms:.3f} ms, {tflops:.1f} TFLOPS")
+```
+
+Step 1 到 Step 3 会刻意用较小尺寸运行（这里是 128×128，Step 3 是 256³），让最初的 walkthrough 容易跟随。
+{ref}`zh_chap_gemm_advanced` 末尾的跨 step *End-to-End Result* 表格采取相反做法：
+它在统一的 M=N=K=4096 尺寸下测量每个 step，包括这个 Step 1 算法，从而让 speedup ratio 可以直接比较。
+
+### Single-Tile Kernel 的限制
+
+这个 kernel 是正确的，而这正是 Step 1 的全部目的；但它只在非常狭窄的设置中正确。
+我们有意内置了四个限制，优化路径剩余部分会一次解除一个：
+
+- 它只处理单个 K tile，因此无法在大 K 上 contract。
+- 它只处理单个 output tile，因此 M 和 N 被固定为 128。
+- 它使用同步 GMEM -> SMEM copy，而不是 TMA。
+- 它不会让数据移动和计算 overlap，因此二者从不同时运行。
+
+---
+
+(zh_chap_k_loop)=
+## Step 2：K-Loop Accumulation
+
+要移除的第一个限制是最小的那个。Step 1 只处理单个宽度为 64 的 K tile，但真实矩阵会在远大于此的 K 上 contract。
+在 Step 2 中，我们保留单个 output tile，但让 K 跨越许多宽度为 64 的 chunk。
+
+想法很直接：每个 chunk 重复一次 load -> MMA -> wait 序列，并让每次 MMA 都累加到同一个 TMEM slot。
+真正需要费心的是同步。跨 iteration 复用一个 mbarrier，会引入本章第一个真正的 correctness hazard。
+如果代码追踪了错误 phase，wait 可能在对应 MMA 实际完成*之前*返回，从而静默破坏结果。
+下面的机制会精确说明它如何出错，以及如何避免。
+
+> **这一步改变什么：Layout reuse**
+> - Scope：不变，仍然是单个 warpgroup。
+> - Layout/reuse：同一对 SMEM tile 和同一个 TMEM accumulator slot 跨 K-loop 复用。不会分配新 storage；operand tile 流经一对固定 buffer，accumulator state 保留在一个 TMEM slot 中。
+> - Synchronization：复用的 MMA barrier 必须在每个 K chunk 上推进到正确 phase，否则后续 wait 可能观察到更早的 completion。
+> - Dispatch：不变。
+
+### K-Loop 机制
+
+Step 1 只在单个宽度为 64 的 K tile 上 contract；这里我们保留它的单个 output tile，但让 K 按矩阵需求延伸。
+为了覆盖大于 64 的 K，我们以 `BLK_K=64` 为 chunk 遍历 K。每次 iteration 把下一个 A 和 B 的 K-slice load 到 SMEM，
+并发射 `Tx.gemm_async`。`accum` flag 把这些 chunk 缝合成一个 dot product：
+第一个 chunk 上，`accum=False` 初始化 TMEM accumulator；后续每个 chunk 上，
+`accum=True` 把该 chunk 的乘积加到已经位于 TMEM 中的 running sum。
+
+同步是需要谨慎的地方。每次 MMA completion 都复用同一个 mbarrier，而安全复用归结为追踪我们正在等待哪个 barrier phase。
+mbarrier 携带 1-bit phase，值为 0 或 1；每当预期 arrival 落地，它就翻转到另一个值。
+微妙之处在于 wait condition 本身：`try_wait(bar, phase)` 会阻塞到 barrier 内部 phase *不同于* `phase` 参数。
+所以我们传入的参数必须命名“我们期望离开的 phase”，而不是“我们等待抵达的 phase”：
+
+| K iteration | wait 前的本地 `phase_mma` | `try_wait` 等待什么 | wait 后的本地更新 |
+|---|---:|---|---:|
+| 0 | 0 | barrier flips to 1 | `phase_mma = 1` |
+| 1 | 1 | barrier flips to 0 | `phase_mma = 0` |
+| 2 | 0 | barrier flips to 1 | `phase_mma = 1` |
+
+单行 `phase_mma ^= 1` 正是让这张表保持正确的东西。去掉它，第二次 iteration 仍会调用 `try_wait(bar, 0)`，
+但 barrier 在第一次 MMA 后已经翻转到 phase 1，因此 wait 看到 mismatch 后会立刻返回，
+早于第二次 MMA 完成。kernel 随后会读取半计算的 accumulator，并在没有任何错误的情况下报告错误答案。
+这是一个能够完美编译和运行的 bug，也正是 phase flip 值得如此关注的原因。
+
+### 完整 Kernel
+
+下面完整 kernel 只是把 K-loop 和 phase flip 折进 Step 1。import 与之前相同：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+```
+
+它被包在 `hgemm_v2(M, N, K)` 中。grid 仍然是 `[1, 1]`，因为我们仍在计算单个 output tile；
+增长的只有它的 K extent：
+
+```python
+def hgemm_v2(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    K_TILES = K // BLK_K
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        bx, by = T.cta_id([M // BLK_M, N // BLK_N])  # still one output tile (M=N=128)
+        wg_id = T.warpgroup_id([1])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout)
+        pool.commit()
+
+        if warp_id == 0:
+            if lane_id == 0:
+                T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+        (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0],
+        layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]))
+
+        phase_mma: T.int32 = 0
+        m_st = T.meta_var(bx * BLK_M)
+        n_st = T.meta_var(by * BLK_N)
+
+        # === K-loop: iterate over K in chunks of BLK_K ===
+        for i in T.serial(K_TILES):   # serial device loop (keeps the full-K A/B parameters correctly shaped)
+            # Load the i-th K chunk
+            Tx.cta.copy(Asmem[:, :], A[:, i*BLK_K:(i+1)*BLK_K])
+            Tx.cta.copy(Bsmem[:, :], B[:, i*BLK_K:(i+1)*BLK_K])
+
+            T.cuda.cta_sync()
+
+            # MMA: accum=False for first tile, True for rest
+            if warp_id == 0:
+                if T.ptx.elect_sync():
+                    Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :],
+                                  accum=(i != 0), dispatch="tcgen05", cta_group=1)
+                    T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+
+            # Wait for MMA, then flip phase
+            T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+            phase_mma ^= 1
+
+        # === Writeback (same as Step 1) ===
+        Dreg = T.alloc_local((BLK_N,), acc_type)
+        Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+        Dreg_wg = Dreg.view(128, BLK_N,
+                            layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+
+        Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+        T.ptx.tcgen05.wait.ld()
+
+        Tx.cast(Dreg_f16[:], Dreg[:])
+        m_thr = T.meta_var(m_st + warp_id * 32 + lane_id)
+        Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:])
+
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+---
+
+(zh_chap_spatial_tiling)=
+## Step 3：Spatial Tiling（Multi-CTA）
+
+K-loop 处理了 contraction dimension，但 M 和 N 仍然固定在单个 128 x 128 tile 上。
+真实输出远大于一个 tile，因此 basic kernel 的最后一块，是一次用许多 tile 覆盖 M 和 N。
+Step 3 启动一个 2D CTA grid，每个 output tile 一个 CTA，并让 GPU 并行计算所有 tile。
+例子使用 M=N=K=256，会得到一个 2x2 tile grid，足以让 indexing 不再平凡，同时又不会被细节淹没。
+
+> **这一步改变什么：Scope**
+> - Scope：一个 2D CTA grid，每个 CTA 拥有一个 128 x 128 output tile。
+> - Layout：不变；在 CTA 内部，这与 Step 2 是同一条 SMEM/TMEM/register 路径。
+> - Dispatch：不变。
+
+### Grid Mapping
+
+grid shape 直接来自 tiling：每个 128 x 128 output tile 一个 CTA，因此总共需要 `[M // BLK_M, N // BLK_N]` 个 CTA。
+与 Step 2 相比，唯一真正的新工作是告诉每个 CTA，矩阵的哪一片是*它*要计算的片段。
+
+CTA `(bx, by)` 拥有这个 output region：
+
+```text
+D[bx * BLK_M : (bx + 1) * BLK_M,
+  by * BLK_N : (by + 1) * BLK_N]
+```
+
+为了产生它，这个 CTA 的 K-loop 会反复 load 自己的 A row band 和 B column band 中匹配的 K-slice：
+
+```text
+A[bx * BLK_M : (bx + 1) * BLK_M, k : k + BLK_K]
+B[by * BLK_N : (by + 1) * BLK_N, k : k + BLK_K]
+```
+
+indexing 直接来自 `D = A @ B.T` 约定：`bx` 选择 A 和 D 的 row，
+而 `by` 选择 B 的 row；一旦应用 transpose，这些 B row 会变成 D 的 column。
+
+每个 CTA 一个 tile 是能工作的最简单映射，但它也很浪费。同一行中的每个 CTA 都会从 GMEM 重新 load 相同 A tile，
+同一列中的每个 CTA 都会重新 load 相同 B tile，因此没有复用相邻 CTA 已经拉进来的数据。
+我们暂时保留这种浪费；persistent scheduling（{ref}`zh_chap_gemm_async` 中的 Step 6）会回到这个问题，
+并让这些共享 operand 在 L2 中保持 hot。
+
+**可以让你的 agent 试试**：在 `M=N=K=256`、`BLK_M=BLK_N=128`、`BLK_K=64` 时，
+让它追踪 CTA `(1, 0)` 和 CTA `(0, 1)`。对每个 CTA，列出 `m_st`、`n_st`、
+每个 K iteration 载入的 A/B slice，以及写入的 D region。
+由于 kernel 计算 `D = A @ B.T`，哪些 B row 会变成 D column？
+
+### 完整 Kernel
+
+这个 kernel 再一次基于 Step 2，这次只有两个变化：grid shape 和 per-CTA offset。
+内部 K-loop 和 writeback 不变。import 仍然相同：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+```
+
+grid 从 `[1, 1]` 变成 `[M // BLK_M, N // BLK_N]`，
+load 和 store 现在会根据 CTA 自己的 `m_st` 和 `n_st` 加 offset：
+
+```python
+def hgemm_v3(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    K_TILES = K // BLK_K
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        # 2D grid: one CTA per 128x128 output tile
+        bx, by = T.cta_id([M // BLK_M, N // BLK_N])
+        wg_id = T.warpgroup_id([1])
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout)
+        pool.commit()
+
+        if warp_id == 0:
+            if lane_id == 0:
+                T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+
+        tmem = T.decl_buffer(
+        (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0],
+        layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)]))
+
+        phase_mma: T.int32 = 0
+
+        # Per-CTA tile offsets
+        m_st = T.meta_var(bx * BLK_M)
+        n_st = T.meta_var(by * BLK_N)
+
+        # K-loop with offset A and B slices
+        for i in T.serial(K_TILES):   # serial device loop (keeps the full-K A/B parameters correctly shaped)
+            Tx.cta.copy(Asmem[:, :], A[m_st:m_st+BLK_M, i*BLK_K:(i+1)*BLK_K])
+            Tx.cta.copy(Bsmem[:, :], B[n_st:n_st+BLK_N, i*BLK_K:(i+1)*BLK_K])
+
+            T.cuda.cta_sync()
+
+            if warp_id == 0:
+                if T.ptx.elect_sync():
+                    Tx.gemm_async(tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :],
+                                  accum=(i != 0), dispatch="tcgen05", cta_group=1)
+                    T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+
+            T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+            phase_mma ^= 1
+
+        # Writeback to the correct output tile
+        Dreg = T.alloc_local((BLK_N,), acc_type)
+        Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+        Dreg_wg = Dreg.view(128, BLK_N,
+                            layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+
+        Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+        T.ptx.tcgen05.wait.ld()
+
+        Tx.cast(Dreg_f16[:], Dreg[:])
+        m_thr = T.meta_var(m_st + warp_id * 32 + lane_id)
+        Tx.copy(D[m_thr, n_st:n_st+BLK_N], Dreg_f16[:])
+
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+## 练习
+
+1. 在 Step 1-3 中，`Tx.copy` 会在 MMA 之前把 A 和 B tile 移入 SMEM。为什么 kernel 需要在 `Tx.gemm_async` 读取这些 SMEM tile 前执行 `T.cuda.cta_sync()`？
+2. 在 Step 2 中，如果从 K-loop 中移除 `phase_mma ^= 1`，会发生什么？kernel 会等待每个 MMA，还是后续 wait 可能过早通过？
+3. 对于 M=N=4096 且 BLK_M=BLK_N=128，Step 3 会启动多少个 CTA？哪些 operand tile 在逻辑上会被相邻 CTA 复用？Step 3 是否利用了这种复用？
diff --git a/zh/chapter_intro_tirx/index.md b/zh/chapter_intro_tirx/index.md
new file mode 100644
index 00000000..dac9ba95
--- /dev/null
+++ b/zh/chapter_intro_tirx/index.md
@@ -0,0 +1,261 @@
+(zh_chap_tirx_primer)=
+# TIRx 简介
+
+:::{admonition} 概览
+:class: overview
+
+- TIRx 是一个用于在 IR 层级编写 GPU kernel 的 Python DSL：你会直接命名硬件，但通过结构化 IR 来表达。
+- 每个 tile operation 都由三个设计元素控制：*scope*（哪些 thread）、*layout*（tile 位于哪里）和 *dispatch*（哪条硬件路径）。
+- 一个可运行的 single-MMA GEMM 会展示三者；本书其余内容就是把这些设计元素扩展到更大规模。
+:::
+
+:::{admonition} 运行示例
+:class: note
+
+这些示例需要 Blackwell GPU（`sm_100a`，例如 B200）。TIRx compiler 随 Apache TVM wheel 的
+`tvm.tirx` 模块发布；请将它与 CUDA build 的 PyTorch 一起安装：
+
+```bash
+pip install apache-tvm==0.25.0
+```
+
+用 `python -c "import tvm, tvm.tirx; print(tvm.__version__)"` 确认它能 import。
+同样的设置可以运行本书中每个可运行示例。
+:::
+
+第一部分解释了硬件是什么。要让它计算任何东西，我们还需要一种编程方式。
+
+我们可以直接写 CUDA 或 PTX，许多高速 kernel 也正是这样写的。问题在于，真正决定 kernel 行为的决策在那里很难看清：
+哪些 thread 运行某个操作、每个 data tile 位于哪里，以及由哪条硬件路径执行它。
+这些选择被埋在 intrinsic 参数、地址算术和约定之中。
+
+TIRx（Tensor IR neXt）是一个 Python DSL，它把这三个决策显式提到台面上：
+**scope**（哪些 thread 运行操作）、**layout**（operand tile 位于哪里）和 **dispatch**（哪条硬件路径执行它）。
+它仍然直接命名硬件概念，包括 thread、shared/tensor memory、barrier 和 `tcgen05` MMA。
+不同之处在于，这些选择现在是结构化 IR，编译器可以 lower、check 和 schedule。
+
+我们不会抽象地介绍这些思想，而是从一个完整 kernel 出发：一个最小 single-MMA GEMM。
+我们先让它跑起来，然后再逐行读回去，观察 scope、layout 和 dispatch 各自如何塑造它，以及 kernel 如何被编译。
+这个 kernel 依赖的 tensor layout model 会在 {ref}`zh_chap_tirx_layout_api` 中独立展开，
+完整语言特性集见 {ref}`zh_chap_language_reference`；这里我们聚焦于一个 kernel 和三个设计元素。
+
+## 第一个 Kernel：Single-MMA GEMM
+
+我们承诺的 kernel 是一个最小 GEMM，缩减到仍然能使用 Tensor Core 的最小版本。它在 K = 64 时计算
+`D = A B^T` 的单个 128 x 128 output tile。整个计算从头到尾表达为一个 `Tx.gemm_async` tile operation。
+（这个 tile operation 并不映射到单条硬件指令：因为硬件 MMA 的 K-atom 是 16，K=64 的 tile 会 lower 成沿 K 方向前进的一小段
+`tcgen05.mma` 指令序列。DSL 的意义恰恰在于我们写 tile，而不是写序列。）
+围绕这个操作，kernel 做常规杂务：分配 shared memory（SMEM）和 tensor memory（TMEM），把 A 和 B 从 global copy 到 shared memory，
+把 tile MMA 发射到 TMEM accumulator 中，通过寄存器把 accumulator 读回，并存储结果。
+虽然它很小，这个 kernel 正是 {ref}`zh_chap_gemm_basics` 中 GEMM 阶梯的 Step 1，那里会完整走读它。
+
+每个 TIRx kernel 都从同一组 import 开始，所以值得先看一次：
+
+```python
+
+import tvm
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import tma_shared_layout, SwizzleMode
+from tvm.tirx.layout import TileLayout, S, TLane, TCol, tid_in_wg
+```
+
+我们把 kernel 包在一个小 builder `hgemm_v1(M, N, K)` 中，它接收 problem shape 并返回一个 `PrimFunc`。
+对于我们选择的 shape `M=N=128, K=64`，launch 恰好只包含一个 output tile，
+这让第一个版本足够简单，可以一次读完：
+
+```python
+def hgemm_v1(M, N, K):
+    a_type = tvm.DataType("float16")
+    b_type = tvm.DataType("float16")
+    d_type = tvm.DataType("float16")
+    acc_type = tvm.DataType("float32")
+
+    BLK_M, BLK_N, BLK_K = 128, 128, 64
+    # MMA_M/MMA_N/MMA_K document the underlying hardware MMA tile; they are not
+    # passed to gemm_async (which derives the MMA shape from the operand and
+    # accumulator tiles), so the later steps omit them.
+    MMA_M, MMA_N, MMA_K = 128, 128, 16
+
+    A_layout = tma_shared_layout(a_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_M, BLK_K))
+    B_layout = tma_shared_layout(b_type, SwizzleMode.SWIZZLE_128B_ATOM, (BLK_N, BLK_K))
+
+    @T.prim_func
+    def kernel(
+        A: T.Buffer((M, K), a_type),
+        B: T.Buffer((N, K), b_type),
+        D: T.Buffer((M, N), d_type),
+    ):
+        T.device_entry()
+        # Step 1 is a single-tile kernel: M = BLK_M and N = BLK_N, so the grid
+        # is 1x1. Starting with a 1x1 grid keeps the per-CTA tile offsets
+        # (m_st, n_st) trivially zero; Steps 3+ generalise this to larger M / N.
+        bx, by = T.cta_id([M // BLK_M, N // BLK_N])
+        wg_id = T.warpgroup_id([1])      # single warpgroup, so wg_id is always 0 (unused below)
+        warp_id = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+    
+        # --- SMEM allocation ---
+        pool = T.SMEMPool()
+        tmem_addr = pool.alloc((1,), "uint32")
+        mma_bar = pool.alloc((1,), "uint64", align=8)
+        pool.move_base_to(1024)
+        Asmem = pool.alloc((BLK_M, BLK_K), a_type, layout=A_layout)
+        Bsmem = pool.alloc((BLK_N, BLK_K), b_type, layout=B_layout)
+        pool.commit()
+    
+        # --- Barrier + TMEM init (warp 0 only) ---
+        if warp_id == 0:
+            if lane_id == 0:
+                T.ptx.mbarrier.init(mma_bar.ptr_to([0]), 1)
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=512, cta_group=1)
+    
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cta_sync()
+    
+        tmem = T.decl_buffer(
+            (128, 512), "float32", scope="tmem", allocated_addr=tmem_addr[0],
+            layout=TileLayout(S[(128, 512) : (1@TLane, 1@TCol)])
+        )
+    
+        m_st = T.meta_var(bx * BLK_M)
+        n_st = T.meta_var(by * BLK_N)
+        phase_mma: T.int32 = 0
+    
+        # --- Load: all threads copy global -> shared (synchronous).
+        # With M=BLK_M and N=BLK_N the slices below cover the full matrices;
+        # the slice form is kept so the diff to Step 3 (multi-tile) is minimal.
+        Tx.cta.copy(Asmem[:, :], A[m_st:m_st + BLK_M, :])
+        Tx.cta.copy(Bsmem[:, :], B[n_st:n_st + BLK_N, :])
+        T.cuda.cta_sync()
+    
+        # --- Compute: single elected thread issues MMA ---
+        if warp_id == 0:
+            if T.ptx.elect_sync():
+                Tx.gemm_async(
+                    tmem[:, :BLK_N], Asmem[:, :], Bsmem[:, :],
+                    accum=False, dispatch="tcgen05", cta_group=1
+                )
+                T.ptx.tcgen05.commit(mma_bar.ptr_to([0]), cta_group=1)
+    
+        T.ptx.mbarrier.try_wait(mma_bar.ptr_to([0]), phase_mma)
+    
+        # --- Writeback: TMEM -> RF -> GMEM ---
+        Dreg = T.alloc_local((BLK_N,), acc_type)
+        Dreg_f16 = T.alloc_local((BLK_N,), d_type)
+        Dreg_wg = Dreg.view(128, BLK_N,
+                            layout=TileLayout(S[(128, BLK_N) : (1@tid_in_wg, 1)]))
+        Tx.wg.copy_async(Dreg_wg[:, :], tmem[:, :BLK_N])
+        T.ptx.tcgen05.wait.ld()
+        Tx.cast(Dreg_f16[:], Dreg[:])
+        m_thr = T.meta_var(m_st + warp_id * 32 + lane_id)
+        Tx.copy(D[m_thr, n_st : n_st + BLK_N], Dreg_f16[:])
+    
+        # --- Deallocate TMEM ---
+        T.cuda.cta_sync()
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+            T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=512, cta_group=1)
+
+    return kernel
+```
+
+在阅读 kernel 之前，先确保它能工作。我们编译它，并用 torch reference 检查输出。
+我们不必写出精确架构：arch（例如 `sm_100a`）会从 device 自动检测，因此 target `"cuda"` 就足够；
+`tir_pipeline="tirx"` 用来选择 TIRx lowering pipeline。编译完成后，`ex.mod(...)` 可以直接接收 torch tensor，
+中间不需要手动转换。
+
+```python
+import torch
+
+target = tvm.target.Target("cuda")
+device = torch.device('cuda')  # gpu(0)
+
+M, N, K = 128, 128, 64
+kernel = hgemm_v1(M, N, K)
+with target:
+    ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+
+torch.cuda.empty_cache()
+torch.cuda.synchronize()
+A_tensor = torch.randn(M, K, dtype=torch.float16, device=device)
+B_tensor = torch.randn(N, K, dtype=torch.float16, device=device)
+D_tensor = torch.zeros(M, N, dtype=torch.float16, device=device)
+
+# ex.mod(...) takes torch tensors directly, the same call form used in every chapter.
+ex.mod(A_tensor, B_tensor, D_tensor)
+
+D_ref = (A_tensor.float() @ B_tensor.float().T).half()
+max_err = float((D_tensor - D_ref).abs().max())
+print(f"Max error vs torch reference: {max_err:.6f}")
+torch.testing.assert_close(D_tensor, D_ref, rtol=2e-2, atol=1e-2)
+print("PASS")
+```
+
+## Scope、Layout、Dispatch
+
+现在 kernel 已经能运行，我们可以读回它，并询问每一行实际上决定了什么。
+这样看，整个 kernel 就是围绕三个设计元素作出的一组选择。其中每个操作都回答同三个问题：
+*谁*运行它，数据*位于哪里*，以及它*如何*执行；这三个答案正是 scope、layout 和 dispatch。
+本节剩余部分会逐一讨论这些设计元素；下面的交互式演示让你看到每个设计元素控制哪些代码行。
+
+```{raw} html
+
+```
+*交互：点击 Scope / Layout / Dispatch，高亮 kernel 中由该设计元素控制的代码行。*
+
+使用这个演示时，关注三个问题：
+
+- **Scope：谁运行这个操作？** `Tx.cta.copy(...)` 是 CTA-scoped，因此全部 128 个 thread 都会帮助完成 GMEM -> SMEM copy。
+  `Tx.gemm_async(...)` 由一个被选出的 thread 发射一次，因为每条 lowered `tcgen05.mma` 指令已经是一次 cooperative MMA launch。
+  `Tx.wg.copy_async(...)` 是 warpgroup-scoped，因此 warpgroup 的 128 个 thread 会逐行切分 TMEM readback。
+- **Layout：每个 tile 位于哪里？** A 和 B 使用 `tcgen05.mma` 期望的 swizzled SMEM layout。
+  accumulator 在 `TLane`/`TCol` layout 下位于 TMEM。register readback view 把 row 映射到 `tid_in_wg`，
+  因此每个 warpgroup thread 拥有一个 row fragment。
+- **Dispatch：哪条硬件路径执行它？** `Tx.gemm_async(..., dispatch="tcgen05", ...)` 选择 Blackwell Tensor Core 路径。
+  copy 操作也有 dispatch 选择：第一个 kernel 使用普通 thread copy，后续 GEMM step 会把这些 copy 换成 TMA，
+  而不改变周围的 scope 或 layout。
+
+**可以让你的 agent 试试**：从第一个 kernel 中挑三行：一个 copy、一个 MMA 和一个 TMEM readback。
+让它为每行标注 scope、layout 和 dispatch，然后检查答案是否匹配代码中的 guard、buffer layout 和 `dispatch=` 参数。
+
+## 编译如何工作
+
+我们上面已经编译过 kernel 来测试它；现在稍微仔细看一下这一步做了什么。
+流程很短：把 `PrimFunc` 包进 `IRModule`，并交给 `tvm.compile(mod, target=..., tir_pipeline="tirx")`。
+这会运行 TIRx lowering pipeline，并返回一个可以直接调用的 `Executable`。
+
+```python
+target = tvm.target.Target("cuda")
+ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+```
+
+至少从轮廓上了解 `tir_pipeline="tirx"` 启动了什么，是很有价值的。pipeline 的核心 pass `LowerTIRx`
+会根据每个 tile primitive 的 scope / layout / dispatch contract 来解析它：
+我们刚刚讨论的三个设计元素正是在这里兑现成指令。之后，常规 host/device split 和 finalize 步骤会产生可 launch 的 module。
+如果你愿意，也可以在 `with target:` block 内编译，这让 kernel 能继承外层 target context。
+
+这个流程的一个好性质是，没有什么对你隐藏：结果可以在两个层级检查。
+你可以用 `.show()` 或 `.script()` 阅读 IR 本身，也可以直接从 compiled module 读取编译器最终生成的 CUDA C。
+
+```python
+kernel.show()                          # pretty-print the TIRx (TVMScript)
+print(kernel.script())                 # ... the same, as a string
+
+# the generated CUDA C source, from the compiled Executable:
+print(ex.mod.imports[0].inspect_source())
+```
+
+这只是一个概览。完整 lowering 故事，包括所有 pass、tile-primitive dispatch 如何解析，以及 host/device split 如何完成，
+见 {ref}`zh_chap_arch`。
+
+## 接下来读什么
+
+一个 kernel 已经足够让我们认识 scope、layout 和 dispatch，并看到它们被编译和运行。
+三个设计元素中的每一个，以及这个 kernel 本身，都会通向进一步展开的章节：
+
+- {ref}`zh_chap_tirx_layout_api`：tensor layout model（`TileLayout`、named axes、swizzle），上面的 operand 和 accumulator placement 都建立在它之上。如果 layout 这个设计元素在三者中最神秘，请从这里开始。
+- {ref}`zh_chap_language_reference`：完整语言特性集，覆盖 parser utilities、data types、buffers and memory、control flow 和 thread synchronization；当你想要完整词汇表，而不只是导览时，可以读这里。
+- {ref}`zh_chap_gemm_basics`：把这个 kernel 作为 GEMM 优化路径的 Step 1，并通过 K-loop accumulation、spatial tiling、TMA 和 warp specialization 逐步扩展。如果你想看同样三个设计元素如何扩展到真实 kernel，这是自然的下一站。
diff --git a/zh/chapter_layout_generations/index.md b/zh/chapter_layout_generations/index.md
new file mode 100644
index 00000000..839c04df
--- /dev/null
+++ b/zh/chapter_layout_generations/index.md
@@ -0,0 +1,345 @@
+(zh_chap_layout_generations)=
+# 跨 GPU 世代的 Tensor Core Operand Layout
+
+:::{admonition} 概览
+:class: overview
+
+- 在 Ampere、Hopper 和 Blackwell 中，Tensor Core 执行的高层操作仍然相同：`D = A B + C`。
+- 代际之间变化的是 operand 如何到达 Tensor Core、支持哪些 tile shape 和 dtype，以及 accumulator 位于哪里。
+- Ampere 使用 warp 级 register fragment。shared memory tile 通过 `ldmatrix` 载入 fragment，accumulator 保留在寄存器中。
+- Hopper 允许 `wgmma` 通过 matrix descriptor 直接从 shared memory 读取 operand。descriptor 会指定 Tensor Core 所期望的 shared-memory swizzle format。
+- Blackwell 保留 shared-memory operand 路径，但把 accumulator 移入 TMEM。Block-scaled MMA 也通过 TMEM 暂存 scale factor。
+- 有两个内存约束贯穿所有世代：global memory coalescing 和 shared memory bank conflict。
+:::
+
+从远处看，Tensor Core 操作似乎很稳定。它把 A 和 B 的 tile 相乘，加上 accumulator C，并产生 D。
+从 Volta 开始，这个形式就没有变过。
+
+但围绕这个操作的细节并没有固定不变。某一代上很快的 kernel，到了下一代可能会变慢。
+使用错误 layout 的 kernel 也可能算出错误答案，即使逻辑数学仍然写作 `D = A B + C`。
+原因在于 Tensor Core 消费的不是抽象矩阵，而是以非常具体的硬件 layout 排列的 operand。
+
+本章沿着三代 GPU 追踪这个 layout contract。Ampere 通过 warp-level register fragment 暴露 Tensor Core。
+Hopper 把输入 operand 移到 shared memory descriptor。Blackwell 保留 shared memory operand，
+但把 accumulator 移入 TMEM。操作仍然是 matrix-multiply-accumulate，但进入和离开 Tensor Core 的路径每一代都在改变。
+
+{ref}`Data Layout ` 章节中的 layout 记法，是我们描述这些 contract 的语言。
+Blackwell 的 TMEM 细节会在 {ref}`zh_chap_tmem` 中单独介绍。
+
+## 两个从未消失的约束
+
+在 Tensor Core 参与之前，两个普通内存约束就已经在塑造 GPU kernel 的 layout。
+
+第一个是 global memory coalescing。当一个 warp 的 32 个 lane 发射 global memory load 时，
+内存系统希望这些地址落入少量连续且对齐的 memory segment。如果地址分散，warp load 就会变成多次 memory transaction。
+同样的逻辑数据移动会消耗更多带宽和更多时间。
+
+第二个是 shared memory bank conflict。shared memory 被划分成 32 个 bank。
+如果 warp 中的多个 lane 访问不同地址，但这些地址映射到同一个 bank，这些访问就无法同时被服务，硬件会把它们串行化。
+因此，一个看起来只是平坦 shared memory array、似乎无害的 layout，可能会因为 bank pattern 而变慢。
+
+swizzling 通常是修复 shared memory 侧问题的方法。逻辑 tile 保持不变，但物理地址映射会被置换，
+让访问 pattern 分散到多个 bank 上，而不是堆叠到同一个 bank。
+
+这两个约束甚至适用于完全不使用 Tensor Core 的 kernel。Tensor Core kernel 还会增加第三个约束：
+operand 必须按 Tensor Core 指令本身期望的 layout 排列。本章剩余部分讨论这个第三约束如何在 Ampere、Hopper 和 Blackwell 间变化。
+
+## Ampere：跨 Warp Lane 的 Register Fragment
+
+在 Ampere 级 GPU 上，主要 Tensor Core 指令是 warp-level 的 `mma.sync.aligned.m16n8k*` 系列。
+关键事实是这条指令从哪里读写数据：寄存器。
+
+A、B，以及 C 或 D accumulator，都是分布在一个 warp 的 32 个 lane 上的 per-thread register fragment。
+shared memory 只是 staging area。在 MMA 运行之前，operand tile 必须从 shared memory 移入该指令所期望的精确 register fragment layout。
+
+数据路径如下：
+
+```text
+SMEM -> 寄存器，使用 ldmatrix
+寄存器 -> 寄存器，使用 mma.sync
+寄存器 -> SMEM，使用普通 store
+```
+
+Ampere 的大部分 layout 故事都来自这条路径。kernel 必须以一种能高效载入的形式把 tile 存进 shared memory，
+然后使用 `ldmatrix` 产生 `mma.sync` 所需的 register fragment。
+
+## Ampere Tensor Core 期望什么
+
+Ampere Tensor Core 读取由 8×8 subtile unit 构成的 register fragment。这些 unit 是 `ldmatrix` 载入、MMA 消费的单位。
+
+以 fp16 或 bf16 输入、fp32 accumulate 的 `mma.m16n8k16` 作为具体例子。accumulator tile 的 shape 是 `16 by 8`。
+它按固定 pattern 分布在 32 个 lane 上。
+
+对于 C 或 D accumulator，lane `l` 持有的 row 是：
+
+```text
+l / 4
+l / 4 + 8
+```
+
+column 是：
+
+```text
+2 * (l % 4)
+2 * (l % 4) + 1
+```
+
+因此每个 lane 拥有四个 fp32 accumulator 值：来自两个 8-row half 的两行，交叉上两个相邻列。
+四个连续 lane 覆盖一行的八个 column。
+
+A operand 使用同样的 M 侧 row carve。K 维度分布在 `l % 4` 以及该 lane 持有的寄存器上。
+对于 fp16 或 bf16，每个 32-bit 寄存器会 pack 两个 K 值。
+
+B operand 使用匹配的 K placement，并把 N 侧分散到 lane group 和寄存器上。
+
+精确细节会随 instruction shape 和 dtype 变化，但原则固定：Tensor Core 期望某种特定的 per-lane register fragment。
+如果值不在这些寄存器的这个 pattern 中，指令就会把错误元素相乘。
+
+在 layout 记法中，m8n8 fragment 就是那种用 named lane axes 写出的 pattern，例如：
+
+```text
+S[(8, 4, 2) : (4@laneid, 1@laneid, 1@m)]
+```
+
+两个 `laneid` iter 一起描述 row 和 column 片段如何散布到 lane 上，而最后的 `m` 分量描述 per-lane register slot。
+
+## `ldmatrix`：从 Shared Memory 到 Register Fragment
+
+`ldmatrix` 是 Ampere 中连接 shared memory 和 Tensor Core register fragment 的指令。它是一个 warp-collective load。
+一条指令会把一个或多个 8×8 的 16-bit matrix 从 shared memory 移入 `mma.sync` 期望的分布式 register layout。
+
+指令形式是：
+
+```text
+ldmatrix.sync.aligned.m8n8.x1.shared.b16
+ldmatrix.sync.aligned.m8n8.x2.shared.b16
+ldmatrix.sync.aligned.m8n8.x4.shared.b16
+```
+
+并且可以带一个可选的 `.trans` qualifier。
+
+`.x1`、`.x2` 和 `.x4` 形式分别载入一个、两个或四个 8×8 matrix。row base address 由 lane 提供。
+对于 matrix `m` 和 row `r`，base address 来自 lane `m * 8 + r`。
+这意味着 `.x1` 使用 lane 0 到 7 提供 row address，`.x2` 使用 lane 0 到 15，`.x4` 使用 lane 0 到 31。
+
+结果会直接落入 MMA fragment。对于基本的 8×8 情况，lane `l` 会收到 Tensor Core 所期望的 row 和 column pair。
+如果用普通的 per-lane `ld.shared` 指令循环，就必须手工复现这种 scatter。
+`ldmatrix` 则把 shared-memory-to-fragment 的重排作为一条 warp-collective 指令完成。
+
+`.trans` 形式会在 load 时转置每个 8×8 matrix。当 operand 的存储方向与 MMA 指令期望的方向相反时，就会使用它。
+
+![ldmatrix 将一个 8x8 共享内存 tile 载入 warp 寄存器 fragment；Ampere 上的反向路径使用普通存储，专用 stmatrix 指令稍后才在 Hopper 出现](../img/ldstmatrix.svg)
+
+## 把 Ampere Fragment 写回
+
+`mma.sync` 完成后，accumulator 仍然是 register fragment。epilogue 必须把这个 fragment 移出去。
+
+在 Ampere 上，没有 `ldmatrix` 的专用反向指令。kernel 使用普通 per-thread store，
+有时在 store 之前配合 warp shuffle 或本地重排，把 accumulator 以有用 layout 写入 shared memory 或 global memory。
+
+这让 Ampere 模型保持简单，但也把许多 layout 工作暴露给 kernel。输入侧使用 `ldmatrix` 创建 fragment。
+计算指令读写 register fragment。输出侧则由从这些 fragment 发出的普通 store 处理。
+
+## Ampere 上的 Swizzle
+
+Ampere kernel 已经需要 shared memory swizzle。原因是 shared memory tile 通常以一种访问 pattern 写入，却以另一种 pattern 读取。
+
+假设一个 tile 是从 global memory 按行填充的。row-major layout 会让这种写入 coalesced 且 bank-friendly。
+但 `ldmatrix` 后面可能以一种实际沿列或跨 8×8 subtile 行走的 pattern 读取该 tile。
+如果使用朴素 row-major layout，这些读取可能堆叠到同一个 shared memory bank 上。
+
+对于一个简单的 `(8, 64)` float16 tile，一行是：
+
+```text
+64 * 2 bytes = 128 bytes
+```
+
+这刚好是一整条 shared memory bank line。沿固定 column 向下走时，每一行前进 128 字节，所以 bank index 会重复。
+八行可能坍缩到同一个 bank 上，造成 8-way conflict。
+
+改成朴素 column-major layout 并不能完整解决问题。它通常只是把 conflict 移到另一种访问上：
+row write 变差，而 column-style read 变好。
+
+XOR swizzle 通过让 physical column 依赖 row 来修复这个问题。一个简单版本是：
+
+```text
+physical_col = logical_col xor row
+```
+
+逻辑 tile 不变。shared memory 中的物理 placement 被置换，使 row-style write 和 Tensor Core read pattern 都能避免 bank conflict。
+
+在 Ampere 上，这种 swizzle 通常通过手写 shared memory index math 表达。后续世代则把它变成硬件引擎使用的 descriptor format 的一部分。
+
+![在朴素 row-major tile 上，行写入会分散到多个 bank，而列读取会在一个 bank 上碰撞；XOR swizzle 在不牺牲合并行写入的情况下，把列读取分散到多个 bank](../img/swizzle_conflict.svg)
+
+## Hopper：`wgmma`、Shared Memory Descriptor 与 Swizzle 格式
+
+Hopper 改变了 Tensor Core 路径的输入侧。它不再要求每个 operand 都用 `ldmatrix` 载入寄存器；
+Hopper 的 `wgmma` 可以直接从 shared memory 读取 operand。
+
+B operand 从 shared memory matrix descriptor 读取。A operand 可以从 shared memory descriptor 或寄存器读取，
+对应 `.ss` 和 `.rs` 两种形式。
+
+这移除了 SMEM-sourced operand 的显式 `ldmatrix` 步骤，但没有移除 layout requirement。
+Tensor Core 仍然期望 operand 以精确的 shared memory format 存储。区别在于，这个 format 现在通过 matrix descriptor 描述给硬件。
+
+## Hopper Tensor Core 期望什么
+
+Hopper shared memory matrix descriptor 是 shared memory 中 matrix tile 的一种紧凑描述。
+它告诉 `wgmma` 如何把逻辑 operand coordinate 转成 shared memory address。
+
+descriptor 包含如下字段：
+
+```text
+起始地址
+主维度偏移
+步长维度偏移
+swizzle 模式
+基址偏移
+```
+
+精确解释取决于 operand major mode。对于 K-major tile，一个 stride 沿 K 前进，另一个沿 M 前进。
+对于 MN-major tile，二者角色互换。
+
+swizzle mode 是 shared memory descriptor format 之一，例如：
+
+```text
+SWIZZLE_NONE
+SWIZZLE_32B
+SWIZZLE_64B
+SWIZZLE_128B
+```
+
+swizzle mode 决定两件事。它决定 descriptor 使用的 atom shape，也决定应用在该 atom 内部的 XOR permutation。
+例如，128-byte swizzle mode 会把 operand 看作由 8-row × 128-byte atom 组成的网格，并在每个 atom 内应用 swizzle。
+
+kernel 仍然必须正确放置字节。TMA 通常负责填充 shared memory tile，而 TMA descriptor 必须使用后续 `wgmma` descriptor 所指定的同一种 swizzle format。
+如果 TMA 写入的是 128-byte swizzled tile，那么 `wgmma` descriptor 就必须把它作为 128-byte swizzled tile 来读取。
+如果 descriptor 和数据不一致，Tensor Core 就会读到错乱的 operand。
+
+这是相对于 Ampere 的主要变化。swizzle 不再只是藏在手写 shared memory indexing 中。
+Hopper 把它提升为 first-class descriptor format。写入 tile 的 TMA load 和读取 tile 的 `wgmma` 指令，
+都可以命名同一种 format。
+
+![Hopper 共享内存 matrix descriptor 把 operand 坐标映射到 swizzled 共享内存 atom：descriptor stride 选择 atom，swizzle 选择 atom 内部的 byte 位置](../img/smem_descriptor.svg)
+
+## Hopper 输出仍然使用寄存器
+
+Hopper 改变了输入路径，但 accumulator 仍然位于寄存器中。
+
+`wgmma` 指令把 accumulator 写入 per-thread register fragment。精确的 fragment 大小和寄存器数量取决于 instruction shape，
+例如 `m64nNk16`，其中 N 会改变 accumulator register 数量。但基本思想和 Ampere 一样：epilogue 消费一个 register fragment。
+
+因此 Hopper 有一个 mixed layout model。输入 operand 可以直接来自 shared memory descriptor，swizzle 由硬件描述。
+输出 accumulator 仍然是 register layout 问题。
+
+Blackwell 改变了输出侧。
+
+## Blackwell：`tcgen05` 与 TMEM
+
+Blackwell 为数据 operand 保留 shared memory descriptor 这一思想。A 和 B 仍然在 shared memory 中按 Tensor Core 期望的 layout 准备。
+某些模式也可以从 TMEM 读取 A operand。
+
+主要变化在 accumulator。`tcgen05.mma` 会把 accumulator 写入 Tensor Memory，也就是 TMEM，
+而不是把它保留为长期存在的 register fragment。在 compute phase 中，accumulator 留在 TMEM。
+epilogue 随后使用 `tcgen05.ld` 把它载回寄存器。
+
+这把输出 layout 问题从寄存器移动到了 TMEM。kernel 必须分配 TMEM、选择正确的 TMEM layout、等待 MMA 完成，
+然后使用匹配的 `tcgen05.ld` 路径，为 epilogue 恢复 accumulator fragment。
+
+`cta_group::1` 和 `cta_group::2` 如何把 accumulator 分配到一个或两个 CTA 上，细节见 {ref}`zh_chap_tensor_cores`。
+与早期世代差异最大的 layout，是 block-scaled scale-factor layout。
+
+## TMEM 中的 Scale Factor Layout
+
+block-scaled MMA mode（如 `mxfp8` 和 `nvfp4`）会加入 scale-factor operand。除了 A 和 B，MMA 还会读取：
+
+```text
+SFA(M, SFK)
+SFB(N, SFK)
+```
+
+其中 `SFK` 是 K scale block 的数量。
+
+数据 operand A 和 B 位于 shared memory。scale factor 位于 TMEM。因此它们有不同的数据移动路径。
+
+TMA 从 global memory load 到 shared memory，并不会直接 load 到 TMEM。因此 scale factor 通常分两步移动：
+
+```text
+通过 TMA 从全局内存移到共享内存
+通过 tcgen05.cp 从共享内存移到 TMEM
+```
+
+只有在这次 copy 之后，scale factor 才进入 `tcgen05.mma` 期望读取它们的内存空间。
+
+TMEM scale-factor layout 使用 TMEM 的硬件坐标 Lane 和 Col。在 TIRx layout 记法中，这些轴写作 `TLane` 和 `TCol`。
+
+一个 128-row scale vector 会被压缩到 32-lane group 中，然后复制到 TMEM 的四个 32-lane window 上。
+在 layout 记法中，核心 pattern 是：
+
+```text
+S[(32, sf_per_mma) : (1@TLane, 1@TCol)] + R[4 : 32@TLane]
+```
+
+shard 放置 base 32-row group：
+
+```text
+TLane = r
+TCol  = s
+```
+
+replica 项会在 lane offset 0、32、64 和 96 处添加副本：
+
+```text
+TLane = r + 32 * q，其中 q ∈ {0, 1, 2, 3}
+TCol  = s
+```
+
+这就是 `warpx4` broadcast pattern。同一个紧凑 scale-factor group 会在完整的 128-lane TMEM 空间中变得可见。
+
+32-bit `TCol` cell 内部还存在 byte packing。packing 取决于 `scale_vec` mode：
+
+```text
+1X：一个 scale 值在整个 32-bit 单元中广播
+2X：打包两个 scale 值，且每个都复制一份
+4X：打包四个 K-block scale 值
+```
+
+![scale_vec byte 打包：1X 在 4-byte cell 中广播一个缩放因子；2X 打包两个缩放因子且各复制一份；4X 打包四个 K-block 缩放因子](../img/sf_scale_vec.svg)
+
+这种 packing 在 Ampere 或 Hopper 中没有直接对应物，因为那些世代没有供 `tcgen05` block-scaled MMA 使用的 TMEM scale-factor operand。
+
+在 `cta_group::2` 中，scale factor 会跟随它所缩放的数据。SFA 缩放 A，因此它按 M 在两个 CTA 之间切分，
+匹配每个 CTA 拥有的 A row。SFB 缩放 B，而 B 被计算中的两个 CTA half 共享，因此 SFB 会 multicast 到两个 CTA
+（{ref}`zh_chap_tensor_cores`）。
+
+## 反复出现的 Fragment
+
+尽管周围的内存路径在变化，一个结构会不断回归：m8n8 风格的 register fragment。
+
+在 Ampere 上，`ldmatrix` 构建这个 fragment，让 `mma.sync` 能读取它。
+
+在 Hopper 上，`wgmma` 把它的 accumulator 写成 register fragment，供 epilogue 使用。
+
+在 Blackwell 上，accumulator 在 compute 期间位于 TMEM，但在 epilogue 处理并存储它之前，
+`tcgen05.ld` 会把它载回 register fragment（{ref}`zh_chap_tmem`）。
+
+因此 fragment 并没有消失，只是角色改变了。早期世代会在整个 compute phase 中把 accumulator 保留在那里。
+Blackwell 则主要在 TMEM 和 epilogue 的边界处使用它。
+
+## 贯穿主线
+
+在 Ampere 上，kernel 显式构建 Tensor Core register fragment。shared memory swizzle 主要由 kernel 通过 index math 负责。
+
+在 Hopper 上，Tensor Core 可以通过 matrix descriptor 直接从 shared memory 读取 operand。
+swizzle 变成 TMA 和 `wgmma` 共享的 named descriptor format。
+
+在 Blackwell 上，输入侧仍然使用 shared memory operand，但 accumulator 移到 TMEM。
+block-scaled MMA 还增加了必须 stage 到 TMEM 中的 scale-factor operand。
+
+descriptor 并不会消除 layout 工作。它们只是把 contract 显式化。kernel 仍然必须确保数据移动路径、memory layout
+和 Tensor Core 指令全部一致。写入 swizzled SMEM tile 的 TMA descriptor、读取该 tile 的 MMA descriptor，
+以及附着在 buffer 上的 layout，都必须描述同一种物理排列。
+
+如果其中任何一部分不一致，硬件仍然会运行。它只是会读到错误字节，或者读取得很慢。
+这就是为什么 layout 不是 Tensor Core kernel 周围的装饰，而是 instruction interface 的一部分。
diff --git a/zh/chapter_performance/index.md b/zh/chapter_performance/index.md
new file mode 100644
index 00000000..791e63a8
--- /dev/null
+++ b/zh/chapter_performance/index.md
@@ -0,0 +1,308 @@
+(zh_chap_performance)=
+# 什么让 Kernel 变快
+
+:::{admonition} 概览
+:class: overview
+
+- Roofline 模型给出 kernel 的性能上限。这个上限由内存带宽或计算吞吐决定。
+- 算术强度决定适用哪一个上限。它表示每搬运一个字节所完成的有用算术工作量。
+- 低算术强度意味着 kernel 是 memory-bound。主要出路是搬运更少字节、更多复用数据、融合操作，或使用更小的 dtype。
+- 高算术强度意味着 kernel 可能是 compute-bound。此时主要任务就是让 Tensor Core 保持忙碌。
+- 在现代 GPU kernel 中，最主要的杠杆是 overlap。只要依赖图允许，TMA、Tensor Core、epilogue 和 store 就应该同时运行。
+:::
+
+kernel 的快慢只有相对于某个上限才有意义。像 330 TFLOP/s 这样的数字本身看起来很大，
+但如果放到一个 dense fp16 或 bf16 Tensor Core 工作能维持约 2 PFLOP/s 的 GPU 上，它的含义就完全不同了。
+如果没有上限作为参照，就很难判断一个 kernel 是已经接近硬件极限，还是仍然让芯片的大部分能力闲着。
+
+Roofline 模型给出的正是这个上限。它把 kernel 拆成两类基本活动：搬运字节，以及执行算术。
+如果 kernel 不能足够快地移动数据，内存带宽就会设定上限。如果 kernel 有足够的数据复用和足够多的算术工作，
+计算吞吐就会设定上限。
+
+本章的数字以 NVIDIA B200 作为贯穿示例。沿用 {ref}`zh_chap_background` 中的约定，我们使用便于推理的整数上限：
+dense fp16 或 bf16 Tensor Core 吞吐约为 2 PFLOP/s，HBM3e 带宽约为 8 TB/s。
+精确值取决于具体设备、时钟、功耗限制和测量设置，因此这里应把它们理解为数量级上限，而不是 datasheet 常数。
+
+## Roofline 模型
+
+每个 kernel 都会移动数据并执行算术。Roofline 模型用这两条路径中更慢的一条来约束 kernel。
+
+compute ceiling 是硬件的最大算术吞吐。对于 B200 上的 Tensor Core GEMM，相关上限就是 Tensor Core 吞吐。
+对于标量或 elementwise kernel，相关上限可能是 CUDA core 吞吐，或另一个功能单元。
+
+memory ceiling 是带宽乘以算术强度。如果一个 kernel 每搬运一个字节只做很少算术，内存带宽就会限制性能。
+如果它每字节执行很多操作，内存就不太可能是限制因素。
+
+基本的 roofline bound 是：
+
+```text
+可达到 FLOP/s <= min(峰值 FLOP/s, 内存带宽 * 算术强度)
+```
+
+算术强度是：
+
+```text
+算术强度 = 有用 FLOP 数 / 搬运的字节数
+```
+
+必须指定内存层级。对于 HBM roofline，这里的字节是 HBM 字节。对于 L2 roofline，它们是 L2 字节。
+对于 SMEM roofline，它们是 shared memory 字节。在本章中，默认 roofline 是 HBM roofline。
+
+在 roofline 图中，x 轴是算术强度，单位是 FLOP/byte。y 轴是可达到的性能。memory roof 是一条斜线：
+
+```text
+性能 = 带宽 * 算术强度
+```
+
+compute roof 是一条水平线：
+
+```text
+性能 = 峰值 FLOP/s
+```
+
+二者在 ridge point 处相交：
+
+```text
+ridge point = 峰值 FLOP/s / 带宽
+```
+
+对于这里使用的 B200 近似数字：
+
+```text
+ridge point ≈ 2000 TFLOP/s / 8 TB/s
+            ≈ 250 FLOP/byte
+```
+
+算术强度低于这个值的 kernel，在 HBM roofline 下是 memory-bound。它无法达到 Tensor Core 峰值吞吐，
+因为它无法每秒提供足够多的字节来喂饱这么多算术。
+
+算术强度高于这个值的 kernel 可能是 compute-bound。此时，内存流量不再是一阶限制。
+剩下的工作，是足够好地驱动计算单元，以接近那条水平 roof。
+
+Roofline 模型真正有用的部分不是图本身，而是它告诉程序员哪个资源正在成为约束。
+memory-bound kernel 不会因为数学指令稍微更好就变快。compute-bound kernel 也不会因为省下几个无关紧要的字节就变快。
+第一步，是知道 kernel 位于 ridge 的哪一侧。
+
+![包含示例工作负载的 B200 roofline，展示内存上限、计算上限和 ridge 点](../img/roofline.png)
+
+## 常见工作负载的算术强度
+
+算术强度通常首先是算法属性，其次才是实现细节。在编写 kernel 之前，通常就能做一个粗略估计。
+
+### Elementwise 与 Reduction
+
+elementwise kernel（如 GELU）和 reduction 风格的 kernel（如 RMSNorm）会读写大 tensor，
+但每个元素只执行少量 FLOP。
+
+它们的算术强度很低，位于 ridge point 的很左侧。这类 kernel 的最佳版本通常试图接近内存带宽 roof，
+而不是 Tensor Core compute roof。
+
+对于这些 kernel，重要问题很机械：
+
+```text
+加载和存储是否合并访问？
+每个字节是否只搬运一次？
+这个操作能否与生产者或消费者融合？
+dtype 能否更小？
+TMA 或向量化访问能否帮上忙？
+```
+
+如果没有复用，也没有 fusion 机会，memory roof 就是真正的上限。
+
+### GEMM
+
+GEMM 是相反的情况。它的算术强度会随问题规模增长，因为每个载入的 tile 都可以被复用到许多 multiply-accumulate 操作中。
+
+对于 `M = N = K` 的方阵 fp16 matmul，理想算术强度大约是：
+
+```text
+AI ≈ 2N^3 / (3 * 2N^2)
+   = N / 3 FLOP/byte
+```
+
+这个估计假设 A 和 B 各读一次，C 写一次，beta 为零，片上复用完美，并且没有额外 metadata、padding 或冗余流量。
+真实 kernel 搬运的数据会比这个理想模型更多。但这个估计仍然有用。
+
+当 `N = 4096` 时：
+
+```text
+AI ≈ 4096 / 3
+   ≈ 1365 FLOP/byte
+```
+
+这个值远在 B200 约 250 FLOP/byte 的 ridge point 右侧。因此，大型 GEMM 在 HBM roofline 下是 compute-bound。
+目标不只是减少 HBM traffic。目标是使用 Tensor Core、持续喂饱它们，并把数据移动与计算 overlap 起来，
+从而让 compute roof 变得可达。
+
+这就是为什么 GEMM 虽然有高算术强度，朴素 GEMM 仍然可能很慢。算法允许高性能，但实现可能让 Tensor Core 闲置。
+
+### 注意力
+
+Attention 位于这两个极端之间。它的算术强度取决于序列长度、head dimension、tiling、masking，
+以及中间 tensor 是否被 materialize。
+
+标准 attention 的关键问题是 score matrix。如果 kernel 把 score matrix 写入 HBM，随后又读回来，
+它就通过内存移动了一个大型中间结果。Flash Attention（{ref}`zh_chap_flash_attention`）通过把相关 tile 留在片上，
+避免这次 HBM 往返，从而提高算术强度。
+
+因此，attention 优化一部分是 roofline 问题，一部分是调度问题。算法被改写，让更少字节进入 HBM。
+随后调度 kernel，让剩余的数据移动和计算 overlap。
+
+## 当算术强度较低时
+
+如果一个 kernel 位于 ridge 左侧，它就是 memory-bound。Tensor Core 或 CUDA core 可能闲置，
+因为瓶颈是字节，而不是算术指令。
+
+有两类应对方式。
+
+第一类是提高算术强度。这条路径杠杆更高，因为它可以把 kernel 推向 compute-bound 区域。
+
+最重要的技术是 fusion。低算术强度的常见来源，是把中间 tensor 写入 HBM，并在下一个操作中立刻读回。
+融合 producer 和 consumer 可以把这个中间结果留在寄存器、SMEM 或 TMEM 中。HBM 往返就消失了。
+
+例子包括：
+
+```text
+带 elementwise epilogue 的 GEMM
+把 normalization 折叠进相邻操作
+在不 materialize 完整 score matrix 的情况下计算 attention
+```
+
+第二种技术是为了复用而 blocking。如果一个 tile 载入一次，并在被逐出之前使用很多次，每个字节就支撑了更多算术工作。
+GEMM 的高算术强度正是来自这种复用。其他工作负载只要存在对 tile 的重复使用，也可以采用同样思想。
+
+第三种技术是减少每个值占用的字节数。从 fp32 换成 fp16、fp8 或 fp4 会减少 traffic，并提高每字节 FLOP。
+当格式需要 metadata、scale factor 或额外转换工作时，真实收益会小于原始 dtype 比例。block-scaled fp8 和 fp4 就是这样的例子。
+即便如此，更小的 dtype 仍然常常是让 kernel 在 roofline 上向右移动的最直接方式之一。
+
+第二类应对方式，是接受 memory roof，并试图抵达它。有些 kernel 没有足够工作可以 fusion，也没有足够复用可以利用。
+纯 copy、简单 elementwise 操作，或对大 tensor 的单遍 reduction，可能在本质上就是 memory-bound。
+
+在这种情况下，目标不是超过 roof，而是饱和它。
+
+这意味着：
+
+```text
+每个字节只搬运一次
+避免冗余读取
+使用合并访问或向量化访问
+对规则的大块 tile 使用 TMA
+保持足够多未完成的内存请求
+算法允许时使用更小的存储 dtype
+```
+
+一旦 memory-bound kernel 达到 memory roof，进一步优化计算就没有帮助。想更快，唯一办法是改变算法，让它搬运更少字节。
+
+## 优化阶梯
+
+roofline 说明什么是可能的，但并不说明达到那个上限有多容易。
+
+大型 fp16 GEMM 理论上可能是 compute-bound。这只意味着 HBM roof 不是主要限制，
+并不意味着任何实现都能达到 Tensor Core roof。要缩小差距，需要正确的指令、layout、staging、同步和调度。
+
+第三部分的 GEMM kernel 会在 B200 上把这一点展示为一系列步骤（{ref}`zh_chap_gemm_advanced`）。
+每一步都保留相同的基本算法，但改变 tile 的计算方式或调度方式。
+
+GEMM 阶梯中第一次测得的大幅跃升，是从 thread-copy tiled 路径移动到 TMA-backed 路径。
+TMA 把规则的 GMEM -> SMEM tile 移动从 CTA thread 手中拿走，让 kernel 通过硬件管理的 bulk copy 喂给 Tensor Core。
+
+在第一次跃升之后，主要改进来自 overlap 和调度。TMA 把未来的 tile 带入 shared memory。`tcgen05.mma` 异步运行。
+epilogue 排空先前结果。software pipelining 和 warp specialization 安排这些部件，让硬件引擎同时活跃。
+
+也没有规则要求每个中间步骤本身都必须更快。像 warp specialization 这样的步骤，可能会暂时把资源花在一种结构上，
+而这种结构不会立刻改善数字。但如果它能启用更简单结构无法表达的后续 overlap，它仍然可能是正确的一步。
+
+![B200 上的 GEMM 优化旅程：从同步分块 baseline，到 TMA、warp 专门化、CTA 集群和多消费者执行的测量点](../img/gemm_perf.png)
+
+## Overlap 是主要杠杆
+
+一旦 GEMM 已经是 compute-bound，并且已经使用 Tensor Core，剩下的差距通常来自 idle time。
+
+一个简单 kernel 可能这样做：
+
+```text
+载入 tile k
+计算 tile k
+存储 tile k
+载入 tile k + 1
+计算 tile k + 1
+存储 tile k + 1
+```
+
+这种 schedule 会让硬件闲置。load 运行时，Tensor Core 在等待。Tensor Core 运行时，copy engine 可能闲置。
+store 排空时，两者都可能在等待。
+
+pipelined kernel 则试图把彼此独立的阶段一起运行：
+
+```text
+载入 tile k + 1
+计算 tile k
+存储 tile k - 1
+```
+
+这就是本书后面使用的 Blackwell kernel 结构背后的核心思想。TMA 处理异步数据移动。`tcgen05.mma` 处理异步 Tensor Core 工作。
+epilogue 和 store 处理输出侧。`mbarrier` 对象连接各个阶段，让每个 consumer 只在真正需要数据时等待。
+
+重点不是移除依赖，而是围绕依赖进行调度。tile `k` 的 MMA 必须等 tile `k` 载入后才能开始。
+tile `k` 的 epilogue 必须等 tile `k` 的 MMA 完成后才能读取 accumulator。
+但 tile `k + 1` 的 load 通常可以在 tile `k` 的 MMA 正在进行时运行，而 tile `k - 1` 的 store 通常也可以同时排空。
+
+这就是为什么后面许多章节会聚焦异步机制：
+
+```text
+用 TMA 处理全局内存到共享内存的搬运
+用 mbarrier 表示加载完成和资源交接
+用 tcgen05 执行异步 Tensor Core 计算
+用 TMEM 存放长生命周期的累加器
+用 warp specialization 分离生产者和消费者角色
+用 cluster 支持更大的协作 tile 和 multicast
+```
+
+它们是不同机制，但服务于同一个调度目标：让有用工作同时运行在不止一条硬件路径上。
+
+## Occupancy 与资源压力
+
+overlap 并不是唯一的 latency hiding 机制。更老也更通用的机制是 occupancy。
+
+occupancy 是驻留在一个 SM 上的工作量。如果一个 warp stall，scheduler 可以运行另一个已经 ready 的 warp。
+它通过保持一池可用的独立 warp 来隐藏延迟。
+
+occupancy 受每个 SM 的资源限制。主要限制包括寄存器、shared memory、warp slot 和 CTA slot。
+如果一个 kernel 每个 thread 使用很多寄存器，或每个 CTA 使用大量 shared memory，它可能 occupancy 很低，
+因为 SM 上只能放下少量 CTA 或 warp。
+
+许多现代 Tensor Core kernel 会有意以降低 occupancy 的方式消耗资源。multi-stage shared memory pipeline 会消耗 SMEM。
+大型 register fragment 会消耗寄存器。TMEM 分配会消耗 Tensor Memory 容量。warp specialization 可能会为 producer
+或 consumer 角色保留整个 warp。
+
+这种取舍是刻意的。与其通过让许多无关 warp 驻留来隐藏延迟，这些 kernel 会在较少数量的驻留 CTA 内部通过显式 overlap
+来隐藏延迟。如果 pipeline 能让 TMA、Tensor Core 和 store 保持忙碌，低 occupancy 的 kernel 仍然可以很快。
+
+没有哪种方式总是更好。有些 kernel 需要高 occupancy，因为它们具有不规则内存访问，或显式 overlap 很有限。
+另一些 kernel 需要深度 staging 和 specialization，因为那是高效喂饱 Tensor Core 的唯一方式。
+正确的问题不是 occupancy 是否很高，而是活跃的硬件单元是否保持忙碌。
+
+## 这对后续章节有什么帮助
+
+本书后续会不断回到同一套诊断问题：
+
+```text
+这个 kernel 受哪条 roof 约束？
+哪个资源正在成为瓶颈？
+什么改动能让 kernel 更接近那条 roof？
+```
+
+对于 memory-bound kernel，答案通常是更少字节和更好的带宽使用。这意味着 fusion、coalescing、vectorized access、
+适用时使用 TMA，以及更小的 dtype。
+
+对于 compute-bound GEMM，答案是先使用 Tensor Core，然后做 overlap。kernel 必须 stage operand、发射异步 MMA 工作、
+保持 pipeline 充满，并在不阻塞计算路径的情况下排空结果。
+
+对于 Flash Attention，第一步是通过把 score 和 probability tile 留在片上来提高算术强度。
+之后，它使用与 GEMM 相同的 overlap 工具：tiled data movement、shared memory staging、异步计算，以及谨慎的资源交接。
+
+这给出了一个实用的优化流程：估计算术强度，定位 roof，判断 kernel 是 memory-bound 还是 compute-bound，
+然后优化真正设定上限的资源。
+
+如果没有这一步，kernel 优化就会变成猜谜。有了它，每一次修改都有理由：要么提高算术强度，
+要么让内存路径更接近带宽峰值，要么减少 compute roof 下的 idle time。
diff --git a/zh/chapter_tensor_cores/index.md b/zh/chapter_tensor_cores/index.md
new file mode 100644
index 00000000..cae9ff26
--- /dev/null
+++ b/zh/chapter_tensor_cores/index.md
@@ -0,0 +1,247 @@
+(zh_chap_tensor_cores)=
+# Tensor Core：`tcgen05`
+
+:::{admonition} 概览
+:class: overview
+
+- `tcgen05` 是 Blackwell 的 Tensor Core 指令家族。它的 MMA 指令以协作方式执行 tile matrix-multiply-accumulate 工作，并由一个被选出的 thread commit。
+- accumulator 位于 TMEM 中，而不是寄存器中。epilogue 稍后用 `tcgen05.ld` 把它带回寄存器。
+- `cta_group::1` 和 `cta_group::2` 控制是一个 CTA 还是两个 CTA 协作完成 MMA。这个选择也会改变 M 维度如何映射到 TMEM。
+- block-scaled MMA mode（如 `mxfp8` 和 `nvfp4`）会添加 scale-factor operand。数据 operand 位于 SMEM，而 scale factor 会通过 TMEM stage。
+:::
+
+dense linear algebra 是现代 GPU 花费大部分有用工作的地方。普通 CUDA-core matrix multiply 无法接近芯片标称峰值
+（{ref}`zh_chap_background`）。快速 GEMM 和 attention kernel 通过以正确的 tile shape、layout 和 synchronization
+喂给 Tensor Core，来接近这个峰值。
+
+从 Volta 开始，基本操作在精神上没有改变。Tensor Core 消费 matrix tile，把它们相乘，并累加结果。
+代际之间变化的是操作如何发射、operand 如何布局，以及 accumulator 位于哪里。
+
+Blackwell 对最后一点做了重大改变。`tcgen05` 的 accumulator 不再作为长期存在的 register fragment 保存。
+它被写入 Tensor Memory，也就是 TMEM（{ref}`zh_chap_tmem`）。这一个变化会影响整个 kernel：
+MMA 写入 TMEM，completion 被异步追踪，epilogue 稍后从 TMEM 中 load accumulator，
+并把它变回自己用于转换和 store 的 register fragment。
+
+本章聚焦计算指令本身。TMA（{ref}`zh_chap_tma`）负责把 operand 移入 SMEM。
+TMEM 负责保存 accumulator 以及某些 scale-factor operand。`tcgen05.mma` 是位于这两次内存移动之间的 Tensor Core 操作。
+
+```{raw} html
+

+
+

+```
+*交互：`tcgen05` accumulator 行为。切换 A 或 B 的 transpose，选择输出宽度 `N`，并逐步走过 `K` 次迭代，观察 partial sum 如何在 TMEM 中累加。*
+
+## `tcgen05` MMA
+
+`tcgen05` MMA 是 Blackwell Tensor Core 的 matrix-multiply-accumulate 指令。它是一条协作指令：
+工作面向一个 warpgroup 执行，在某些模式下还可以涉及同一个 cluster 中的两个 CTA。
+这条指令不是由每个 thread 独立发射的；而是由一个被选出的 thread 代表参与组 commit 这个操作。
+
+把 MMA 拆成三个问题会更清楚。
+
+第一个问题是谁协作。普通模式使用一个 CTA，写作 `cta_group::1`。更大的模式使用 cluster 中的两个 CTA，写作 `cta_group::2`。
+两种情况下，这条指令都表示对一个 tile 的一次 Tensor Core 操作，而不是一个 thread 的标量操作。
+
+第二个问题是 operand 和结果住在哪里。数据 operand 通常位于 SMEM。某些变体也可以从 TMEM 读取 A operand。
+accumulator 写入 TMEM。operand layout 必须匹配 Tensor Core 的期望，包括数据 operand 使用的 swizzled shared-memory layout
+（{ref}`zh_chap_data_layout`）。
+
+第三个问题是如何观察 completion。`tcgen05.mma` 是异步的。发射 MMA 并不意味着 multiply-accumulate 已经完成。
+指令会在操作 commit 后返回，而 Tensor Core 继续运行。kernel 使用 commit group 和 `mbarrier` 来得知结果何时 ready
+（{ref}`zh_chap_async_barriers`）。
+
+正是这种异步行为让 overlap 成为可能。快速 kernel 不会发射 MMA 后立刻 stall 到它结束。
+它可以发射 MMA，开始准备后续 tile，并且只在真正需要结果时等待。代价是每一次 handoff 都必须显式。
+如果 epilogue 在 MMA completion barrier 触发之前读取 TMEM，那就是读得太早。
+
+## Accumulator 位于 TMEM
+
+在 Ampere 和 Hopper 上，accumulator 以寄存器形式暴露给程序。MMA 产生 per-lane register fragment，
+epilogue 直接消费这个 fragment。这很简单，但它把 accumulator 大小绑定到了每个 thread 的寄存器预算上。
+
+Blackwell 打破了这个链接。`tcgen05.mma` 把 accumulator 写入 TMEM，这是 Blackwell 上一个 CTA 作用域的内存空间。
+accumulator 可以在整个 compute phase 中留在 TMEM，epilogue 稍后使用 `tcgen05.ld` 把它载回寄存器。
+
+这改变了 kernel 的形态。register fragment 在边界处仍然重要。epilogue 仍然需要寄存器，以便转换、应用 elementwise work，
+并存储结果。但长期存在的 accumulator state 不再是 register allocation 问题，而是 TMEM allocation 和 layout 问题
+（{ref}`zh_chap_tmem`）。
+
+这就是为什么 `tcgen05` 和 TMEM 必须放在一起理解。MMA 指令决定计算哪个 tile。
+TMEM 决定 accumulator 落在哪里。epilogue 必须使用匹配的 load 路径，以自己期望的 register layout 恢复 accumulator。
+
+## `cta_group::1` and `cta_group::2`
+
+`tcgen05` MMA 可以运行在 `cta_group::1` 或 `cta_group::2` 模式。
+
+在 `cta_group::1` 中，一个 CTA 拥有这次 MMA。它的 operand 位于该 CTA 的 SMEM 中，accumulator 写入该 CTA 的 TMEM。
+
+在 `cta_group::2` 中，cluster 中的两个 CTA 协作处理一个 MMA tile。每个 CTA 都有自己的 SMEM 和 TMEM。
+accumulator 并不是存储在一个跨越两个 CTA 的物理 TMEM 区域中，而是切分到两个 CTA 上，每个 CTA 持有自己的部分。
+偶数 CTA 发射指令，并为这一对 CTA commit completion barrier。
+
+这个选择很重要，因为它改变逻辑 accumulator tile `C(M, N)` 如何映射到 TMEM。
+TMEM 有 128 个硬件 Lane row，以及最多 512 个硬件 Col column。在 TIRx layout 记法中，这些轴写作 `TLane` 和 `TCol`。
+MMA mode 决定 `C` 的 row 和 column 如何放到这些 TMEM 轴上。
+
+有四种有用情况值得记住。
+
+下图沿用演示中的颜色约定：紫色表示 SMEM operand，橙色表示 TMEM accumulator state，绿色表示 Tensor Core MMA 路径。
+CTA 身份通过标签和位置表示，而不是通过改变这些硬件颜色表示。
+
+### `cta_group::1`, `M = 128`
+
+这是最简单的情况。一个 CTA 计算 128-row tile。TMEM 也有 128 个 Lane row。
+因此映射是直接的：accumulator 的 row `m` 映射到 Lane `m`，N 维度映射到 TMEM column。
+
+结果填满 128 个 Lane row 和 N 个 Col column。这是 baseline 图景。CTA 在 SMEM 中拥有 A 和 B，
+并在自己的 TMEM 中拥有完整 accumulator tile。
+
+![cta_group::1, M=128：行 m 直接映射到 TMEM lane m](../img/mma_cg1_m128.svg)
+
+### `cta_group::1`, `M = 64`
+
+当 `M = 64` 时，accumulator 只有 64 行，但 TMEM 仍然有 128 个 Lane row。
+硬件并不会简单地把 row 0 到 63 pack 到 lane 0 到 63。相反，它会把它们以四段 16-row run 的形式分散到 128 个 lane 中。
+
+row 0 到 15 去 lane 0 到 15。row 16 到 31 去 lane 32 到 47。
+row 32 到 47 去 lane 64 到 79。row 48 到 63 去 lane 96 到 111。
+
+这会在 lane 16 到 31、48 到 63、80 到 95、112 到 127 留下空隙。这些空隙是有意的。
+通过不同的 lane alignment，另一个独立的 `M = 64` MMA 可以占用互补 lane。
+这让两个较小的 M tile 能共享 128-lane TMEM 结构，而不会彼此踩踏。
+
+N 维度仍然映射到 TMEM column。不寻常的部分只在于 M row 在 Lane 上的 placement。
+
+![cta_group::1, M=64：四段 16 行 run，lane stride 为 32，为另一个对齐的 M=64 tile 留出空间](../img/mma_cg1_m64.svg)
+
+### `cta_group::2`, `M = 256`
+
+当 M 维度大到一个 CTA 无法自然持有时，MMA 可以使用 `cta_group::2`。对于 `M = 256`，切分很直接：
+CTA 0 持有 row 0 到 127，CTA 1 持有 row 128 到 255。
+
+每个 CTA 使用自己的 TMEM Lane row 0 到 127，以及完整 N column。物理上，这是两个独立的 128-row TMEM region，
+每个 CTA 一个。逻辑上，它们形成一个 256×N accumulator tile。
+
+每个 CTA 也提供 A 中对应自己 M row 的部分。B 按该模式的要求对两个 CTA 可用。
+偶数 CTA 负责发射 MMA，并为这对 CTA commit completion barrier。
+
+这是 {ref}`zh_chap_gemm_advanced` 中 two-CTA cluster GEMM 使用的模式。
+
+![cta_group::2, M=256：M 连续切分到两个 CTA 上，每个 CTA 128 行](../img/mma_cg2_m256.svg)
+
+### `cta_group::2`, `M = 128`
+
+`cta_group::2`、`M = 128` 模式仍然使用两个 CTA，但 M 维度更短。由于总共只有 128 行，
+每个 CTA 接收 64 个 M row。
+
+剩余的 lane 容量用于 pack N 维度。在每个 CTA 内，N 的一半占据 lane 0 到 63，
+另一半占据 lane 64 到 127。这样即使每个 CTA 只拥有 64 个 M row，也能使用全部 128 个 Lane row。
+
+因此这个切分有两部分。M 在 CTA pair 之间切分，每个 CTA 64 行。
+随后 N 在每个 CTA 内部跨 TMEM Lane row 的 lower half 和 upper half 切分。
+
+![cta_group::2, M=128：每个 CTA 64 个 M 行，N 的两半堆叠在下/上半 lane 上](../img/mma_cg2_m128.svg)
+
+在这些模式中，原则相同。`tcgen05.mma` 计算一个逻辑 accumulator tile，但这个 tile 必须放入物理的
+128 Lane × 最多 512 Col 的 TMEM 空间。mode 和 M shape 决定这种 placement。
+kernel 后续在把 accumulator 读回时，必须使用同一种映射。
+
+对于这里的 kernel，TMEM 中的 accumulator 通常是 f32。这是常见的高精度路径。
+它不是唯一可能的 accumulator type。`.kind::f16` 路径可以用 f16 accumulate。
+
+## Operand Placement
+
+对于 dense MMA mode，A 和 B 会在 MMA 运行前准备在 SMEM 中。TMA 负责把 global memory tile 移入 SMEM。
+kernel 会把这些 SMEM tile 安排成 Tensor Core 期望的 layout，包括任何必需的 swizzle。
+
+accumulator C 写入 TMEM。这是与早期世代的主要区别。epilogue 不会直接把 accumulator 作为 MMA 指令输出接收。
+它必须用 `tcgen05.ld` 从 TMEM 显式 load。
+
+在 `cta_group::1` 中，一个 CTA 提供 operand 并拥有 accumulator。在 `cta_group::2` 中，
+每个 CTA 从自己的 SMEM 提供自己一侧的 operand，并拥有 accumulator 中属于自己的 TMEM 部分。
+当 A 按 M 切分时，每个 CTA 保留自己 M slice 的 A row。B 按 mode 共享，因为两个 M slice 都要乘以同一个 N×K tile。
+
+阅读 kernel 时，这种分离很重要。SMEM placement 回答 Tensor Core 如何读取 A 和 B。
+TMEM placement 回答 accumulator 去哪里。这两个 layout 由 MMA mode 联系起来，但它们不是同一个内存空间，不能互换看待。
+
+## Block-Scaled MMA
+
+dense mode 直接从 SMEM 读取数据 operand，并累加到 TMEM。Block-scaled MMA 增加两个 operand：A 和 B 的 scale-factor tensor。
+
+这用于 `mxfp8` 和 `nvfp4` 这样的极低精度格式。低精度格式很高效，但动态范围很小。
+单个 global scale 通常过于粗糙。如果 scale 按最大值选择，小值会损失精度；如果 scale 按小值选择，大值可能 clip。
+
+block scaling 通过给小 K block 分配 scale factor 来修复这个问题。一组连续 K 元素共享一个 scale。
+MMA 在概念上用对应 scale 对每个 block 做 dequantize，然后用 accumulator type 累加乘积。
+
+对于 A 和 B，这会引入两个 scale-factor tensor：
+
+```text
+SFA(M, SFK)
+SFB(N, SFK)
+```
+
+其中 `SFK = K / B`，而 `B` 是沿 K 的 block size。
+
+精确 block size 取决于格式。重要的是，scale axis 以更粗粒度跟随 K。
+每个 scale factor 描述的是一块 K 值，而不是单个元素，也不是整个矩阵。
+
+数学形式是：
+
+```text
+acc += (Aq * scale_a) * (Bq * scale_b)
+```
+
+其中 `Aq` 和 `Bq` 是 quantized low-precision value，scale 在 accumulate 前恢复它们的近似幅度。
+
+scale dtype 也很重要。使用 `e8m0` scale 时，每个 scale 实际上是 2 的幂。
+使用 `nvfp4` 所采用的 `e4m3` scale 时，scale 是一个小浮点值，可以表示 2 的幂之间的值。
+
+## Scale Factor 位于哪里
+
+block-scaled `tcgen05.mma` 与 dense MMA 有一条重要 placement rule 不同：scale factor 从 TMEM 读取。
+
+数据 operand A 和 B 仍然 stage 在 SMEM 中。scale factor SFA 和 SFB 通过 TMEM stage。
+由于 TMA load 到 SMEM，scale factor 通常需要额外一步。kernel 先把它们 load 到 SMEM，
+再用 `tcgen05.cp` 从 SMEM copy 到 TMEM。只有当 scale factor 位于 TMEM 中时，block-scaled MMA 才能读取它们。
+
+这给 scale factor 带来了不同于数据 operand 的移动路径：
+
+```text
+A, B:     从全局内存到 SMEM，随后 MMA 读取 SMEM
+SFA, SFB: 从全局内存到 SMEM，随后 tcgen05.cp 将 SMEM 复制到 TMEM，最后 MMA 读取 TMEM
+```
+
+scale factor 的 TMEM layout 很紧凑。一个 128-row scale vector 可以 pack 到 32 个 Lane row 中：
+lane position 基于 `r % 32`，column 方向基于 `r / 32`。
+数据随后可以 broadcast 到读取完整 128 Lane 空间的四个 warp 上（{ref}`zh_chap_layout_generations`）。
+
+这是为什么 TMEM layout 必须显式的好例子。accumulator layout 和 scale-factor layout 都在 TMEM 中，
+但它们不是同一个 layout。accumulator 使用 MMA output mapping，scale factor 使用 block-scaled MMA 期望的 compact layout。
+
+## `cta_group::2` 中的 Scale Factor
+
+在 two-CTA 情况下，scale factor 跟随它缩放的数据。
+
+SFA 缩放 A。由于 A 按 M 在 CTA pair 之间切分，SFA 也按 M 切分。每个 CTA 持有与自己 A row 对应的 SFA row。
+
+SFB 缩放 B。由于两个 CTA 都乘以同一个 B tile，SFB 必须对两个 CTA 可见。实践中，这意味着 SFB 会 multicast 到 CTA pair。
+
+这就是 block-scaled cluster GEMM 中常见 load pattern 的来源。SFA 按 CTA load，使用该 CTA 自己 M slice 的 mask。
+SFB 会 broadcast 到这一对 CTA，因为两个 CTA 都需要同一组 N-side scale factor。
+
+![块缩放 MMA 放置：A 和 B 在 SMEM 中打包；SFA、SFB 和 C 位于 TMEM，其中 SFA 按 M 跨 CTA 切分，SFB 多播到 CTA 对](../img/mma_block_scaled.svg)
+
+## 保持 MMA Contract 匹配
+
+一个 Blackwell GEMM tile 会经过几条专门化路径。
+
+TMA 把 A 和 B 从 global memory 带入 SMEM。对于 block-scaled mode，它也会把 scale factor 带入 SMEM。
+需要时，`tcgen05.cp` 把这些 scale factor 移入 TMEM。`tcgen05.mma` 读取 operand，在 Tensor Core 上异步运行，
+并累加到 TMEM 中。completion barrier 告诉 kernel 这个 accumulator 何时 ready。
+epilogue 随后用 `tcgen05.ld` 把 accumulator 从 TMEM 载回寄存器，并存储最终输出。
+
+跨越这些路径，kernel 必须保持三个 contract 匹配：SMEM operand layout、TMEM accumulator 或 scale-factor layout，
+以及让下一个 consumer 可以安全运行的异步 completion signal。
diff --git a/zh/chapter_tirx_layout_api/index.md b/zh/chapter_tirx_layout_api/index.md
new file mode 100644
index 00000000..e379790e
--- /dev/null
+++ b/zh/chapter_tirx_layout_api/index.md
@@ -0,0 +1,783 @@
+(zh_chap_tirx_layout_api)=
+# TIRx Layout API
+
+:::{admonition} 概览
+:class: overview
+
+- TIRx layout API 会把 {ref}`zh_chap_data_layout` 中的 layout 记法变成编译器对象。主要对象是 `TileLayout`、`SwizzleLayout` 和 `ComposeLayout`。
+- `TileLayout` 描述 named hardware axes 上的仿射 placement。它由 shard spec `S[...]`、replica spec `R[...]` 和可选 offset 构成。
+- 一个 layout 会把一个逻辑坐标映射到一个或多个物理坐标。`layout.apply()` 会求值这个映射。
+- `SwizzleLayout` 描述用于避免 bank conflict 的、基于 XOR 的 shared-memory swizzle。`ComposeLayout` 会把 swizzle 叠加到 tile layout 上。
+- `tmem_datapath_layout`、`tcgen05_atom_layout` 和 `wg_local_layout` 等现成 constructor 覆盖了 kernel 中反复出现的硬件 layout。
+:::
+
+{ref}`zh_chap_data_layout` 引入了本书通用的记法：一个 tile shape、一组位于 named axes 上的 stride，
+以及一个可选 replication term，用来表示被复制而不是被 partition 的值。本章会把这个记法转成编译器使用的 API。
+
+目标是让页面上的记法和 kernel 中的代码看起来几乎一样。当你写下这样的 layout：
+
+```python
+S[(128, 256) : (1@TLane, 1@TCol)]
+```
+
+你不只是在写解释。你正在构造一个可以附着到 buffer 上的 `TileLayout` 对象。
+之后，每个接触这个 buffer 的 tile operation 都可以从 layout 中读取它的 placement。
+placement 写一次、检查一次，然后由编译器复用。
+
+layout 可以在从 pool 分配时附着，也可以在声明 buffer 时附着：
+
+```python
+pool.alloc(shape, dtype, layout=layout)
+
+T.decl_buffer(shape, dtype, scope=scope, layout=layout)
+```
+
+从那一刻起，buffer 就携带自己的物理 placement。tile operation 不需要重复说明每个元素住在哪里。
+
+layout object 位于同一个模块中：
+
+```python
+from tvm.tirx.layout import (
+    TileLayout,
+    SwizzleLayout,
+    ComposeLayout,
+    S,
+    R,
+    laneid,
+    warpid,
+    tid_in_wg,
+    TLane,
+    TCol,
+    m,
+    tcgen05_atom_layout,
+    tmem_datapath_layout,
+)
+```
+
+这个 API 背后有一个核心思想：layout 不必把逻辑索引映射到单个物理地址。
+它会把逻辑索引映射到 named axes 上的一组物理坐标。通常情况下，这个集合只有一个元素。
+当存在 replication 时，同一个逻辑元素会有多个物理 placement。
+
+这就是为什么 layout model 有三部分：shard、replica 和 offset。shard 放置元素。
+replica 把它复制到额外坐标。offset 平移整个 placement。
+
+## 通过例子理解 Layout
+
+下面的例子展示了 API 的基本形状。
+
+TMEM 中的 accumulator 可以写成 TMEM axes 上的直接 placement：
+
+```python
+acc = TileLayout(S[(128, 256) : (1@TLane, 1@TCol)])
+```
+
+这里，逻辑 row 映射到 `TLane`，逻辑 column 映射到 `TCol`。
+在 {ref}`zh_chap_tmem` 中，硬件坐标称为 Lane 和 Col。在 TIRx layout 记法中，这些硬件轴写作 `TLane` 和 `TCol`。
+
+block-scaled MMA 的 scale-factor layout 使用 replication：
+
+```python
+scale_factor_layout = TileLayout(
+    S[(32, sf_per_mma) : (1@TLane, 1@TCol)] + R[4 : 32@TLane]
+)
+```
+
+shard 会在 TMEM 中放置一个 32-row group。replica 以 32 lane 的 stride 把这个 group 重复四次，
+因此这个 32-row group 在完整 128-lane TMEM 空间中可见。
+
+tensor-core register fragment 可以分布在 lane 和 warp 上：
+
+```python
+frag = TileLayout(
+    S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)]
+)
+```
+
+同一个物理轴可以出现多次。在这个例子中，两个不同 iter 都贡献到 `laneid`。
+没有显式轴的 stride 会使用默认 memory axis `m`。
+
+在真实 kernel 中，常见硬件 layout 通常来自 constructor：
+
+```python
+acc = tmem_datapath_layout("D", 128, 256)
+
+ld = tcgen05_atom_layout("32x32b", (128, 64), "float32")
+```
+
+这些 constructor 返回普通 `TileLayout` 对象。它们只是便利函数，不是另一套机制。
+你可以检查返回的 layout，把它与其他 layout compose，或者在 shape 比较特殊时手写底层 `S[...]` 和 `R[...]` 形式。
+
+## 交互式演示
+
+进入机制之前，有一个能动手戳的具体对象会很有帮助。下面的演示允许你选择 preset layout，
+编辑 logical shape 和 `S` 或 `R` 项，选择 dtype 和 swizzle mode，并点击一个元素，查看哪个或哪些物理坐标拥有它。
+
+```{raw} html
+

+  ▶ Open the demo full screen ↗
+

+
+
+```
+
+这个演示很有用，因为 API 的大部分内容就是演示所展示过程的精确版本。
+一个逻辑元素进入 layout。layout 把它 flatten，跨自己的 iter split，在 named axes 上累加坐标，
+然后在需要时应用 replication。
+
+## TileLayout
+
+`TileLayout` 是主要的 affine layout object。它通常用正文中相同的记法书写：
+
+```python
+TileLayout(S[shape : strides])
+```
+
+`S` 项是 shard spec。你可以这样读它：取一个这种 shape 的逻辑 tile，并用这些 named axes 上的 strides 放置它。
+
+当一个值需要出现在多个位置时，shard spec 会用 replica spec 扩展：
+
+```python
+TileLayout(S[shape : strides] + R[replica_shape : replica_stride])
+```
+
+也可以加入可选 offset：
+
+```python
+TileLayout(S[shape : strides] + R[replica_shape : replica_stride] + offset)
+```
+
+在表面之下，这些部分由 iter 表示。一个 iter 是三元组：
+
+```text
+(extent, stride, axis)
+```
+
+它描述了在一个 named axis 上的 strided walk。extent 表示这个 iter 有多少个位置。
+stride 表示每一步移动多远。axis 表示哪个硬件坐标正在改变。
+
+一个 layout 有三部分。
+
+### Shard
+
+shard，也就是 `D`，是由 `S[...]` 构建的部分。它把逻辑索引 partition 到一个或多个 iter 上，并产生 base physical coordinate。
+
+For example:
+
+```python
+S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)]
+```
+
+有四个 shard iter。它们的 extent 是 `8`、`2`、`4` 和 `2`。
+它们的 stride 分别把数据放到 `laneid`、`warpid`、再次 `laneid`，以及默认 memory axis `m` 上。
+
+这推广了普通 shape-and-stride 规则。区别在于，stride 附着到 named hardware axes 上，而不是附着到单个 flat address 上。
+
+### Replica
+
+replica，也就是 `R`，描述同一个逻辑元素的额外物理副本。replica iter 与逻辑索引无关。
+它们枚举硬件空间中的额外 offset。
+
+For example:
+
+```python
+R[2 : 4@warpid]
+```
+
+会在 `warpid` 轴上创建两个相隔四个 warp 的副本。
+
+replication 不是为了方便而设的技巧，它描述的是真实硬件行为。有些数据会跨 warp、lane 或 memory region broadcast。
+logical-to-physical mapping 自然支持这一点，因为一个逻辑元素可以映射到一组物理坐标。
+
+### Offset
+
+offset，也就是 `O`，是加到每个结果上的固定坐标。
+
+For example:
+
+```python
+5@warpid
+```
+
+会把整个 placement 在 `warpid` 轴上平移五个单位。
+
+offset 用于把 tile 放到选定 base coordinate、为独占使用保留一段区域，
+或描述同一资源中位于另一个 tile 之后开始的 tile。
+
+### 把这些部分组合起来
+
+layout 会按顺序应用这三部分。
+
+首先，shard 计算 base coordinate。然后，replica 把这个 coordinate fan out 成零个或多个额外副本。
+最后，offset 平移每个 coordinate。
+
+对于逻辑坐标 `x`，结果是：
+
+```text
+L(x) = { D(x) + r + O | r in R }
+```
+
+如果没有 replica，`R` 只包含 zero offset，因此结果是 singleton set。如果存在 replica，
+结果会为每个 replica position 包含一个 coordinate。
+
+在 TIRx 语法中，一个完整 layout 可以这样写：
+
+```python
+layout = TileLayout(
+    S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)]
+    + R[2 : 4@warpid]
+    + 5@warpid
+)
+```
+
+从左到右读，shard 放置逻辑 tile，replica 在相隔四个 warp ID 的地方创建第二份副本，
+offset 把整个 placement 平移到从 `warpid = 5` 开始。
+
+如果 iter 已经被构造成对象，同一个 layout 可以直接构造：
+
+```python
+TileLayout.from_iters(shard, replica, offset)
+```
+
+大多数用户代码使用 `S[...]` 和 `R[...]` 记法，因为它更接近数学形式。
+
+## Named Axes
+
+layout 中的 axis 不是匿名维度。每个 axis 都命名一个真实硬件坐标，或 compiler-level placement coordinate。
+
+Examples include:
+
+```text
+bx, by, bz
+cbx, cby, cbz
+tx
+warpid
+laneid
+wgid
+tid_in_wg
+wid_in_wg
+m
+P, F
+Bank
+TLane, TCol
+```
+
+`bx`、`by` 和 `bz` 等 grid axes 会把工作放置到 CTA 之间。
+`cbx`、`cby` 和 `cbz` 等 cluster axes 会把工作放置在 CTA cluster 内部。
+`tx`、`warpid`、`laneid`、`tid_in_wg` 和 `wid_in_wg` 等 thread axes 描述 CTA 或 warpgroup 内部的 ownership。
+`m` 轴是默认 linear memory axis。`P` 和 `F` 用于二维 scratchpad-style placement。
+`Bank` 命名 shared memory bank。`TLane` 和 `TCol` 是 TMEM Lane 和 Col 坐标在 TIRx layout 中的名字。
+
+axis name 是 layout 的一部分。这很重要，因为整数值相同的两个坐标可能表示不同硬件事物。
+`1@tx` 不等于 `1@tid_in_wg`。`1@laneid` 不等于 `1@TLane`。layout 会让这些含义保持显式。
+
+## Forward Mapping
+
+求值一个 layout，就是拿一个逻辑坐标并计算它物理上落在哪里。API 方法是：
+
+```python
+layout.apply(*coord)
+```
+
+对于没有 replication 的 layout，结果是一个 coordinate dictionary。带 replication 时，结果是一组 coordinate dictionary。
+coordinate dictionary 会把 axis name 映射到整数位置，例如：
+
+```python
+{"laneid": 7, "warpid": 2, "m": 1}
+```
+
+求值规则有四步。
+
+首先，按 row-major 顺序 flatten 逻辑坐标。对于逻辑坐标：
+
+```text
+x = (x0, x1, ..., xr-1)
+```
+
+inside a logical shape:
+
+```text
+(S0, S1, ..., Sr-1)
+```
+
+flat index 是：
+
+```text
+flat = x0 * S1 * S2 * ... * Sr-1
+     + x1 * S2 * ... * Sr-1
+     + ...
+     + xr-2 * Sr-1
+     + xr-1
+```
+
+第二，把这个 flat index 跨 shard extents split。如果 shard extents 是：
+
+```text
+(e0, e1, ..., en-1)
+```
+
+那么 split 会产生 components：
+
+```text
+c0, c1, ..., cn-1
+```
+
+使用 shard extents 上相同的 row-major 顺序。
+
+第三，使用 stride 把每个 component 累加到它的 axis 上。如果 shard iter `k` 的 extent 为 `ek`、
+stride 为 `sk`、axis 为 `ak`，那么 component `ck` 贡献：
+
+```text
+ck * sk @ ak
+```
+
+同一个 axis 上的所有贡献会加在一起。随后再加上 offset。
+
+第四，应用 replica iter。每个 replica iter 都贡献一个与逻辑坐标无关的额外 offset。
+如果有多个 replica iter，layout 会枚举所有组合。
+
+这条规则的一个有用后果是，layout 不需要 hard-code 输入 shape。它需要的是逻辑 tile 的总元素数等于 shard extents 的乘积。
+一旦满足这一点，flatten 和 splitting 就定义了映射。
+
+## 案例研究：Tensor Core Register Tile
+
+考虑一个逻辑 `(8, 16)` tile，它分布在两个 warp 上，每个 warp 有 32 个 lane。
+每个 lane 拥有一个小 register fragment。register slot 由默认 memory axis `m` 表示。
+
+```python
+layout = TileLayout(
+    S[(8, 2, 4, 2) : (4@laneid, 1@warpid, 1@laneid, 1)]
+    + R[2 : 4@warpid]
+    + 5@warpid
+)
+```
+
+取 `(8, 16)` tile 中的一个逻辑元素 `(i, j)`。
+
+row-major flat index 是：
+
+```text
+flat = 16 * i + j
+```
+
+按 shard extents `(8, 2, 4, 2)` split 得到：
+
+```text
+c0 = i
+c1 = floor(j / 8)
+c2 = floor(j / 2) mod 4
+c3 = j mod 2
+```
+
+shard contribution 是：
+
+```text
+laneid = 4 * c0 + c2
+warpid = c1
+m      = c3
+```
+
+加上 offset `5@warpid` 后，它变成：
+
+```text
+laneid = 4 * i + floor(j / 2) mod 4
+warpid = floor(j / 8) + 5
+m      = j mod 2
+```
+
+replica 项：
+
+```python
+R[2 : 4@warpid]
+```
+
+会给 `warpid` 加上 `0` 或 `4`。因此完整映射是：
+
+```text
+laneid = 4 * i + floor(j / 2) mod 4
+warpid = floor(j / 8) + 5 + 4 * r，其中 r ∈ {0, 1}
+m      = j mod 2
+```
+
+shard 把 tile 放到 warp 5 和 6 上。replica 随后把它复制到 warp 9 和 10。
+因此，同一个逻辑元素会出现在两个 warp 位置。
+
+这个例子说明了为什么模型使用一组物理坐标。replication 并不适合自然地表示为从物理坐标到逻辑坐标的函数；
+它更自然地表示为从一个逻辑坐标到多个物理坐标的函数。
+
+## 案例研究：Blackwell Tensor Memory
+
+同一个 layout model 也适用于 memory placement。axis 不必是 thread axis，也可以是 memory axis。
+
+TMEM 通过硬件 Lane 和 Col 坐标寻址。在 TIRx layout 记法中，这些轴写作 `TLane` 和 `TCol`。
+
+考虑这个 layout：
+
+```python
+layout = TileLayout(
+    S[(2, 128, 112) : (112@TCol, 1@TLane, 1@TCol)]
+)
+```
+
+如果逻辑 tile shape 是 `(2, 128, 112)`，split components 就是逻辑坐标本身。
+对于元素 `(a, l, c)`，映射是：
+
+```text
+TLane = l
+TCol  = 112 * a + c
+```
+
+extent-128 iter 以 stride `1@TLane` 填满 128 个 TMEM Lane row。
+extent-2 iter 以 stride `112@TCol`、extent-112 iter 以 stride `1@TCol`，二者共同覆盖 224 个 column：
+
+```text
+TCol in [0, 224)
+```
+
+224-column span 是有意的。TMEM layout 不必是 2 的幂。block-scaled FP8 GEMM 可能会选择 224-column accumulator，
+因为完整 256-column tile 将无法为两个 accumulator stage 加 scale factor 留出足够 TMEM 容量。
+layout API 可以直接表达这个 shape。
+
+## Scale Factor Layout
+
+上面的 accumulator layout 是纯 placement。每个逻辑 accumulator 元素映射到一个 TMEM coordinate。
+block-scaled MMA 的 scale factor 不同，因为同一个物理 group 可能需要在多个 warp window 中可见。
+这正是 replication 变得有用的地方。
+
+一个紧凑 scale-factor layout 可以写作：
+
+```python
+scale = TileLayout(
+    S[(32, sf_per_mma) : (1@TLane, 1@TCol)]
+    + R[4 : 32@TLane]
+)
+```
+
+shard 会在 TMEM 中放置一个 32-row scale-factor group：
+
+```text
+TLane = r
+TCol  = s
+```
+
+对于逻辑 scale coordinate `(r, s)`。
+
+replica 项创建四个相隔 32 lane 的副本：
+
+```text
+TLane = r + 32 * q，其中 q ∈ {0, 1, 2, 3}
+TCol  = s
+```
+
+因此，这个 32-row group 在 TMEM lane 0 到 31、32 到 63、64 到 95、96 到 127 上都可见。
+这就是 `warpx4` 广播模式（{ref}`zh_chap_layout_generations`）。
+四个 warp-sized TMEM lane window 中的每一个都会看到同一个 scale-factor group。
+
+在完整 block-scaled MMA layout 中，这个 atom 会与 M row 和 K scale-factor group 上的 outer iter 结合。
+根据 scale-factor dtype，多个 scale factor 还可能被 pack 到一个 32-bit `TCol` cell 中。
+例如，fp8 scale factor 可以把四个值 pack 到一个 32-bit column cell 中。
+可选的 stride-zero reuse 和 pipeline-depth iter 随后可以描述跨多个 MMA 的 scale reuse 以及 double buffering。
+
+重要的是，同一个 `TileLayout` model 描述了这两种情况。accumulator 是 TMEM 中的单一 placement。
+scale factor 是同一 TMEM address space 中的复制式 placement。
+
+## 现成 Layout
+
+大多数 kernel 不会手写每一种硬件 layout。TIRx 为经常出现的 layout 提供 constructor。
+
+```python
+tmem_datapath_layout(datapath, rows, cols)
+```
+
+返回由 `tcgen05.mma` 写入的 TMEM accumulator layout。`datapath` 参数选择 row placement pattern。
+例如，`"D"` 对应 `M = 128` 的 identity-style placement，而 `"F"` 对应 `M = 64` 的 scattered placement。
+
+```python
+tcgen05_atom_layout(instr_shape, tensor_shape, dtype)
+```
+
+返回由 `tcgen05.ld` 或 `tcgen05.st` atom 移动的 register tile layout。
+instruction shape 的例子包括 `.32x32b`、`.16x64b`、`.16x128b` 以及相关形式。
+在 DSL 层面，这是一个 warpgroup-distributed tile。lowering 期间，它会变成四条 warp-collective `tcgen05.ld`
+或 `tcgen05.st` 指令，每个 warp 一条，每个 warp 处理自己的 32 个 TMEM lane。
+
+```python
+wg_local_layout(cols, rows=128)
+```
+
+返回 warpgroup-local register tile，通常在 `tid_in_wg` 上每个 thread 一行。
+
+这些 helper 用来避免手写常见硬件映射。它们并不隐藏模型。
+每个 helper 都返回普通 `TileLayout`，由上面描述的同一组 `S` 和 `R` 部分构成。
+
+## SwizzleLayout and ComposeLayout
+
+`TileLayout` 是 affine 的。它可以表达 named axes 上的 stride、replication 和 offset。
+这足以覆盖许多 placement，包括 thread fragment、TMEM tile 和 compact scale-factor layout。
+
+shared memory swizzle 需要别的东西。用于避免 bank conflict 的 swizzle 不是 affine stride pattern，
+而是对线性 shared-memory address 做基于 XOR 的 permutation。
+
+因此，TIRx 把 swizzling 保留为单独的 layout object：
+
+```python
+SwizzleLayout(...)
+```
+
+并把它与 tile layout compose：
+
+```python
+ComposeLayout(swizzle, tile)
+```
+
+tile layout 先产生一个 linear memory address。随后 swizzle 对这个地址做 permutation。
+把这两层分开，比强行把 XOR permutation 塞进 affine layout model 更清晰。
+
+## 为什么需要 Swizzle
+
+shared memory 被划分成 32 个 bank，每个 bank word 持有 4 字节。
+当一次访问中的多个 lane 触及同一个 bank 中的不同地址时，该访问会因为 bank conflict 被串行化。
+
+朴素 row-major tile 会结构性地产生这种 conflict。考虑一个具有 row-major layout 的 `(8, 64)` float16 tile：
+
+```python
+TileLayout(S[(8, 64) : (64@m, 1@m)])
+```
+
+逻辑元素 `(i, j)` 的 linear element address 是：
+
+```text
+m = 64 * i + j
+```
+
+每行是 64 个 float16 值，也就是 128 字节。这正好是一整条 shared memory bank line。
+如果一个 warp 以固定 `j` 向下读一列，每一步 row 都会前进一整条 128-byte line。
+bank index 会重复，因此 column read 会跨多个 row 坍缩到同一个 bank 上。
+
+swizzle 通过让低地址位依赖更高的 row bit 来改变这一点。
+原本会反复落在同一个 bank 上的一列，会被分散到不同 bank 上。
+
+## Swizzle Transform
+
+`SwizzleLayout` 由三个整数参数控制：
+
+```text
+per_element = M
+swizzle_len = B
+atom_len    = S
+```
+
+输入是一个 linear element address `m`。
+
+`m` 的低 `M` 位保持不变。这会保留一个小的 contiguous element group。
+更高位会右移到一个临时值中：
+
+```text
+x = m >> M
+```
+
+然后，`x` 中位置 `[S, S + B)` 的 bit group 会 XOR 到 `x` 中位置 `[0, B)` 的 bit group 上。
+swizzled address 随后通过把未改变的低 `M` 位放回去形成。
+
+等价地：
+
+```text
+mask = (1 << B) - 1
+
+low  = m & ((1 << M) - 1)
+x    = m >> M
+x2   = x ^ ((x >> S) & mask)
+
+addr = (x2 << M) | low
+```
+
+为了让 layout well formed，`S` 必须至少为 `B`。
+
+这个 transform 的目的不是改变 tile 中有哪些逻辑元素，而是改变这些元素在 shared memory 中落在哪里。
+MMA 仍然读取同一个逻辑 tile。swizzle 让物理 bank pattern 更好。
+
+## 选择 Swizzle 参数
+
+正常使用中，swizzle 参数由 dtype 和 shared-memory swizzle mode 选择。
+常见 mode 是 32-byte、64-byte 和 128-byte swizzle。
+
+`per_element` 参数的选择会让一个小的 vector-sized group 保持 contiguous。对于 float16，一个 16-byte vector 包含 8 个元素，因此：
+
+```text
+M = log2(8) = 3
+```
+
+使用 128-byte swizzle 时，layout 使用：
+
+```python
+SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
+```
+
+这会保持 16-byte vector group 完整，同时仍然足够置换更大的 shared-memory address pattern，以打破 column bank conflict。
+
+大多数代码不应该手工推导这些参数。dtype 和 descriptor mode 通常会决定它们。
+对程序员来说，重要的是 TIRx layout 中的 swizzle、TMA descriptor 和 MMA expectation 三者匹配。
+
+因此，一个 swizzled shared memory allocation 看起来像这样：
+
+```python
+tile = TileLayout(S[(8, 64) : (64@m, 1@m)])
+swizzle = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
+
+layout = ComposeLayout(swizzle, tile)
+```
+
+composed layout 会被附着到 shared memory buffer 上。
+
+## 元素的 Bank 与 Line
+
+要判断 swizzle 是否有帮助，可以把 swizzled element address 转回 shared memory bank。
+
+令 `addr` 为 swizzled element address，`b` 为元素大小（字节）。byte address 是：
+
+```text
+byte = addr * b
+```
+
+bank 是：
+
+```text
+bank = floor(byte / 4) mod 32
+```
+
+128-byte bank line 是：
+
+```text
+line = floor(byte / 128)
+```
+
+对于 float16，`b = 2`，因此 bank 公式变成：
+
+```text
+bank = floor(addr / 2) mod 32
+```
+
+这是下面 worked example 中使用的公式。
+
+## Worked Example：`(8, 64)` float16 Tile 上的 128B Swizzle
+
+回到 row-major float16 tile：
+
+```text
+m = 64 * i + j
+```
+
+使用：
+
+```python
+SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
+```
+
+transform 变成：
+
+```text
+x    = m >> 3
+addr = ((x ^ ((x >> 3) & 7)) << 3) | (m & 7)
+```
+
+由于：
+
+```text
+m = 64 * i + j
+```
+
+我们可以写成：
+
+```text
+q = floor(j / 8)
+r = j mod 8
+```
+
+swizzled address 是：
+
+```text
+addr = 64 * i + 8 * (q xor i) + r
+```
+
+现在看 column `j = 0`。此时 `q = 0` 且 `r = 0`，所以：
+
+```text
+addr = 72 * i
+```
+
+对于 float16，bank 是：
+
+```text
+bank = floor(addr / 2) mod 32
+```
+
+因此八个 row 映射到：
+
+```text
+i = 0: bank 0
+i = 1: bank 4
+i = 2: bank 8
+i = 3: bank 12
+i = 4: bank 16
+i = 5: bank 20
+i = 6: bank 24
+i = 7: bank 28
+```
+
+这一列现在触及八个不同 bank。conflict 消失了。
+
+如果没有 swizzling，同一列的 address 是：
+
+```text
+m = 64 * i
+```
+
+因此：
+
+```text
+bank = floor(64 * i / 2) mod 32 = 0
+```
+
+每一行都落在 bank 0 上，因此访问会被串行化。swizzle 只改变物理 placement，
+但这已经足够把 column access 变成 conflict-free。
+
+这个保证依赖于按设计方式使用 swizzle。dtype、swizzle width 和 access shape 必须匹配 TMA 与 MMA descriptor mode。
+128-byte float16 swizzle 是围绕相关 16-byte row chunk 和 Tensor Core access pattern 设计的。
+它并不承诺任意 shared memory access 都会变成 conflict-free。
+本章开头的演示会让这一点可见：选择 dtype 和 swizzle mode，观察没有 swizzle 时一列如何坍缩到一个 bank 上，
+再观察应用匹配 swizzle 后它如何散布到 bank 视图中。
+
+## 设计理由
+
+layout API 遵循三项设计选择。
+
+第一，它支持一般 shape。硬件 tile 不总是 2 的幂。global tensor、shared memory stage、TMEM accumulator
+和 scale-factor buffer 的 shape，往往来自容量限制或算法选择。layout model 把这些 shape 视为普通情况。
+
+第二，映射方向是从逻辑坐标到物理坐标。这个方向很重要，因为 replication 很常见。
+一个逻辑元素可能住在多个物理位置。logical-to-physical map 会直接把它表示为一组坐标。
+
+第三，hardware axes 是显式的。layout 不使用匿名维度，也不依赖稍后的上下文来解释它们。
+`tx`、`tid_in_wg`、`laneid`、`warpid`、`TLane` 和 `TCol` 之间的差异，会写进 layout 本身。
+
+legality 和 feasibility check 并不只是 layout object 的职责。layout 可以说明数据放在哪里。
+更高层的 tile primitive 会决定某个给定操作能否合法且高效地使用这个 placement。
+这种分离让 layout API 保持小巧，同时仍然给编译器足够信息来 dispatch 真实硬件操作。
diff --git a/zh/chapter_tma/index.md b/zh/chapter_tma/index.md
new file mode 100644
index 00000000..f8a4ca92
--- /dev/null
+++ b/zh/chapter_tma/index.md
@@ -0,0 +1,192 @@
+(zh_chap_tma)=
+# 异步数据移动：TMA
+
+:::{admonition} 概览
+:class: overview
+
+- TMA 是一个硬件引擎，用于在 global memory 和 shared memory 之间异步复制 tile。一个 thread 发射 copy，引擎负责移动字节。
+- TMA copy 由 tensor-map descriptor 描述。descriptor 会告诉引擎 global tensor 的 shape、strides、tile coordinates，以及 shared-memory swizzle mode。
+- 在 load 路径上，TMA 可以在写入 shared memory 时对 tile 做 swizzle，让 tile 直接落到 Tensor Core 所期望的 layout 中。
+- TMA load 通过带 byte-count tracking 的 `mbarrier` 完成。TMA store 使用 commit group 和 wait group。
+:::
+
+只有当 Tensor Core 有准备好的数据可消费时，它才有帮助。在 GEMM 或 attention kernel 中，一旦 pipeline 填满，
+数学部分可能是 compute-bound（{ref}`zh_chap_performance`），但只有下一个 operand tile 按时到达，pipeline 才能保持填满。
+
+移动 tile 的旧方法是让 thread 自己复制。每个 thread 计算地址，从 global memory 发出 load，并把值存进 shared memory。
+这可行，但它会把 warp 指令花在地址算术和 copy bookkeeping 上，而不是计算上。
+它还会让 copy 路径出现在同一批本应喂给 Tensor Core 的 warp 的指令流中。
+
+Tensor Memory Accelerator，即 TMA，会把这项工作移入硬件 copy engine。一个 thread 发射一次 tile copy。
+随后 copy engine 在 global memory 和 shared memory 之间异步移动一个矩形 tile。
+当引擎正在搬运字节时，CTA 的其余部分可以继续做其他工作。
+
+TMA 也处理一部分 layout 问题。Tensor Core 不只是需要 shared memory 中有正确的值；
+它还需要这些值处在正确的 shared-memory layout 中。在 load 路径上，TMA 可以在写入 tile 时应用 shared-memory swizzle。
+这让 tile 可以直接落到后续 MMA 期望的 layout 中。
+
+```{raw} html
+

+
+

+```
+*交互：TMA 将 tile 从 global memory 复制到 shared memory。切换 swizzle mode，并悬停 source cell，查看它落到 shared memory 中的哪里。*
+
+## 一个 Thread 发射，硬件移动 Tile
+
+TMA copy 从一个 issuing thread 开始。这个 thread 不会循环遍历 tile 中的所有元素。
+它把 copy 的描述交给硬件，然后由 TMA engine 执行传输。
+
+主要输入是 tensor-map descriptor。descriptor 描述 global tensor，以及应如何从中读取一个 tile。
+它记录 tensor shape、strides、element size、tile shape 和 swizzle mode 等信息。
+issuing thread 还会提供 tile 应该落到的 shared-memory address。
+
+指令发射后，copy 会异步运行。issuing thread 可以继续执行，CTA 中的其他 thread 也可以继续执行。
+传输现在由 TMA engine 负责，而不是由普通 load/store 指令循环负责。
+
+这给了 kernel 两种不同方式来表达同一个逻辑操作：“复制这个 tile”。
+
+一种路径是 thread copy。thread 协作地从 global memory load，并 store 到 shared memory。
+这让 kernel 能直接控制每一次访问，但会消耗 thread 指令和用于地址计算的寄存器。
+
+另一种路径是 TMA copy。一个 thread 发射传输，由硬件 copy engine 执行矩形 copy。
+对于大型规则 tile，尤其是 Tensor Core kernel 使用的 operand tile，这是自然的路径。
+
+这两条路径有不同的同步规则和性能行为。在二者之间选择，是一个 dispatch decision。
+layout 告诉 kernel 它想要哪种内存排列。scope 告诉它哪些 thread 或 CTA 参与其中。
+dispatch 则决定这个 copy 是由普通 thread code 实现，还是由 TMA 实现。
+
+## Swizzled Layout
+
+移动 tile 还不够。tile 还必须以 Tensor Core 能高效读取的 layout 放入 shared memory。
+
+这正是 TMA swizzling 的用武之地。当 TMA 把 tile 写入 shared memory 时，它可以置换 shared-memory address pattern。
+global memory tile 仍然是一个逻辑矩形，但 shared memory 中的 destination layout 可以是 swizzled 的。
+
+swizzle mode 是 TMA descriptor 的一部分。一旦 descriptor 设置好，issuing thread 就不必手工应用 swizzle。
+引擎会在字节落入 shared memory 时应用它。
+
+重要要求是一致性。TMA descriptor、shared-memory tile layout 和后续 MMA 指令必须都描述同一个 layout（{ref}`zh_chap_data_layout`）。
+如果 TMA 用一种 swizzle 写入 tile，而 MMA 却按另一种 swizzle 去读，硬件仍然会精确执行它被要求做的事；
+只是这些字节对计算来说会排列错误。
+
+在这一点上，layout 记法就不再只是 bookkeeping device。DSL 使用的 layout 必须匹配 TMA descriptor 和 Tensor Core 指令使用的硬件 layout。
+例如，如果 kernel 说某个 operand tile 存储在 128-byte swizzled layout 中，TMA descriptor 就必须使用匹配的 swizzle mode，
+而 MMA dispatch 也必须期望同一个 shared-memory arrangement。上面的演示允许你在 no swizzle 和 128-byte swizzle 之间切换；
+悬停某个 source element，可以查看应用 swizzle 后它落在哪里。
+
+理解 swizzle 的一种有用方式是：TMA 并没有改变逻辑 tile。它改变的是逻辑元素在 shared memory 中的物理落点。
+后续 MMA 消费的仍然是同一个逻辑 A 或 B tile。swizzle 只决定这个 tile 如何排列在 shared memory bank 上。
+
+## 用于 Tiling 和 Swizzling 的 3D TMA
+
+普通 TMA copy 移动的是平坦 2D tile，但 Tensor Core 想要的 shared-memory layout 通常会被 *tiled* 成 swizzle atom
+（来自 {ref}`zh_chap_data_layout` 的 8 x 128-byte atom）。TMA 用额外的 descriptor 维度来处理这一点。
+**3D TMA** 把 shared-memory box 描述为 `(group, row, col)`，其中 group 维度跨 atom 行走，内部两个维度则在一个 atom 内寻址。
+一次 3D copy 随后既会按 atom 布置 tile（tiling），又会在每个 atom 内应用 swizzle，
+因此数据到达时已经处在 MMA 期望的 layout 中，不需要单独的 tiling 或 swizzling pass。
+
+```{raw} html
+

+
+

+```
+*交互：一个 3D TMA copy，以 (group, row, col) 寻址，并 tiled 到 swizzled shared memory 中。*
+
+选择 swizzle *format* 与这种 tiling 绑定在一起。更宽的 swizzle 会把一列分散到更多 bank 上，
+所以能适配时默认选择 128-byte swizzle；但一个 N-byte atom 需要 tile 的 contiguous dimension 能填满它。
+因此，如果某个 tile 因 shape 约束而较小，就不能使用 128-byte swizzle，必须降到 64-byte 或 32-byte：
+经验法则是选择 tile 能填满的最大 swizzle（{ref}`zh_chap_data_layout`）。下面的演示直接展示了这个约束：
+16 x 16 tile 上的 128-byte swizzle，只有当 tile 被切成匹配 atom 的 16 x 8 group 后，才会 conflict-free。
+
+```{raw} html
+

+
+

+
+```
+*交互：16 x 16 tile 上的 128-byte swizzle；一旦 tiled 成 16 x 8 group，就会 conflict-free。*
+
+## Completion：Load
+
+copy 是异步的，所以仅仅发射还不够。consumer 不能只因为 TMA 指令已经发射，就去读取 shared-memory tile。
+只有当引擎已经完成字节写入后，tile 才能安全读取。
+
+对于 TMA load，完成信号是一个 `mbarrier`（{ref}`zh_chap_async_barriers`）。
+
+通常序列如下：
+
+1. 为 pipeline stage 初始化或复用一个 `mbarrier`；
+2. 告诉 barrier 这次 TMA transfer 预计写入多少字节；
+3. 发射 TMA load；
+4. 让 TMA engine 在字节到达时更新 barrier；
+5. 在读取 shared-memory tile 之前，让 consumer 等待对应 barrier phase。
+
+byte count 通过如下操作设置：
+
+```text
+mbarrier.arrive.expect_tx(bytes)
+```
+
+这做了两件事。它记录预期 transfer size，同时也执行 issuing thread 在 barrier 上的 arrival。
+barrier 不会仅仅因为这个调用发生就完成。它仍然等待 TMA engine 报告预期字节已经到达。
+
+随着 transfer 进行，引擎会对 barrier 执行 complete-tx update。只有两个条件都满足时，barrier phase 才会翻转：
+arrival count 已满足，并且 pending byte count 到达零。
+
+consumer 随后等待这个 barrier。一旦对预期 phase 的 wait 完成，shared-memory tile 就 ready 了。
+此时 MMA 路径可以安全读取它。
+
+![TMA 加载同步流程](../img/tma_sync_flow.png)
+
+这是其他异步 producer-consumer handoff 使用的同一个 barrier 模型。producer 是 TMA engine。
+consumer 是 MMA 路径，或任何读取 shared-memory tile 的其他代码。barrier 是二者之间的显式 handoff。
+
+## Completion：Store
+
+TMA store 按相反方向移动数据：从 shared memory 到 global memory。它们同样是异步的，但 completion mechanism 不同。
+
+TMA load 通常会喂给同一个 kernel 内部的 consumer。MMA 路径需要知道 shared-memory tile 何时 ready。
+这就是 load 路径使用 `mbarrier` 的原因。
+
+TMA store 通常把最终数据写出到 global memory。通常没有立即的 in-kernel consumer 在等待被存储的结果。
+kernel 主要需要知道的是：什么时候可以安全复用 shared-memory buffer，或结束 store sequence。
+
+为此，TMA store 使用 commit group 和 wait group。kernel 发射一个或多个 store，commit 这个 group，
+稍后等待这个 group drain。wait 完成后，从 kernel 的角度看，该 group 中的 store 已经完成，
+store 使用的 shared-memory region 可以安全复用。
+
+所以规则很简单：
+
+```text
+TMA 加载：通过带字节计数跟踪的 mbarrier 等待
+TMA 存储：通过 commit group 和 wait group 等待
+```
+
+这两种机制在不同 handoff point 服务于同一个目的。load 需要让 shared-memory tile 对后续 consumer 可见。
+store 需要确保 outgoing transfer 已完成，然后 kernel 才能复用 source storage，或依赖 store 已经 drain。
+
+## 为什么 TMA 对 Pipelining 很重要
+
+当 TMA 成为 pipeline 的一部分时，它最有用。kernel 可以在 Tensor Core 计算当前 tile 的同时，发射未来 tile 的 load。
+load 在后台运行，compute 在前台运行。当未来 tile 变成当前 tile 时，barrier 把二者连接起来。
+
+典型 GEMM loop 会反复使用这种结构。shared memory 的一个 stage 保存当前被 MMA 消费的 tile。
+另一个 stage 正在被 TMA 填充。随着 loop 前进，这些角色会轮换。MMA 读取某个 stage 之前，会等待该 stage 的 load barrier。
+TMA 覆写某个 stage 之前，kernel 会确保前一个 consumer 已经用完它。
+
+这就是为什么 TMA 和 `mbarrier` 通常一起出现在 Blackwell 和 Hopper 风格的 kernel 中。
+TMA 给 kernel 一个异步 copy engine；barrier 给 kernel 一种精确方式，知道复制的字节何时 ready。
diff --git a/zh/chapter_tmem/index.md b/zh/chapter_tmem/index.md
new file mode 100644
index 00000000..204a1f83
--- /dev/null
+++ b/zh/chapter_tmem/index.md
@@ -0,0 +1,115 @@
+(zh_chap_tmem)=
+# 特殊内存：TMEM
+
+:::{admonition} 概览
+:class: overview
+
+- TMEM 是 Blackwell 专有、供 `tcgen05` 使用的内存空间。它是每个 SM 上的二维 scratchpad，拥有 128 个 Lane row 和最多 512 个 Col column。
+- `tcgen05.mma` 把 accumulator 写入 TMEM。block-scaled MMA 也使用 TMEM 存放 scale factor。
+- TMEM 通过 Lane 和 Col 寻址。在 TIRx layout 记法中，这两个硬件轴写作 `TLane` 和 `TCol`。
+- TMEM 不像寄存器那样自动分配。kernel 必须以 32 column 为单位显式分配并释放它。
+- 普通 shared-memory load/store 不能访问 TMEM。数据通过专用异步 `tcgen05` 指令在 TMEM、寄存器和 shared memory 之间移动。
+:::
+
+在 Hopper 及更早 GPU 上，Tensor Core（{ref}`zh_chap_tensor_cores`）accumulator 位于寄存器中。
+这个模型很容易推理：MMA 指令产生一个 register fragment，kernel 在 compute phase 中保持这个 fragment live，
+epilogue 稍后读取它、转换它，并存储结果。
+
+问题在于 register pressure。寄存器是固定的 per-thread 资源。随着 MMA tile 变大，accumulator fragment 也会变大。
+到某个点，accumulator 会开始挤占 thread 还需要保存的其他值。更大的 tile 有利于 Tensor Core 吞吐，
+但把整个 accumulator 保存在寄存器中，会让这些更大 tile 更难使用。
+
+Blackwell 改变了数据路径的这一部分。`tcgen05` 的 accumulator 不必在整个 compute phase 中停留在寄存器里。
+相反，`tcgen05.mma` 会把 accumulator 写入 Tensor Memory，也就是 TMEM。
+TMEM 是早期 NVIDIA GPU 没有的内存空间。它是 SM 上的二维 scratchpad，形状为 128 个 Lane row × 最多 512 个 Col column，
+作用域是使用它的 CTA。
+
+这个额外内存空间让 Blackwell 能支持更大的 Tensor Core tile，而不必强迫整个 accumulator 进入 per-thread register。
+但 TMEM 并不像寄存器那样自动。编译器不会把它当作普通 register storage 简单发放。
+kernel 必须分配 TMEM，用正确 layout 寻址它，用正确指令把数据移入移出，并在 CTA 完成后释放它。
+
+## 二维地址空间
+
+TMEM 不是平坦 byte array，而是一个二维地址空间。硬件把它的两个坐标命名为 Lane 和 Col。
+它有 128 个 Lane row，以及最多 512 个 Col column。每个 Col 是一个 32-bit column。
+
+这个形状很重要，因为 `tcgen05.mma` 会使用这个二维结构把 accumulator 写入 TMEM。
+一个 TMEM location 由 Lane coordinate 和 Col coordinate 描述，而不是由单个 shared-memory-style byte offset 描述。
+
+当 kernel 在 TIRx 中声明 TMEM buffer 时，它会给这个 buffer 一个覆盖这两个硬件坐标的 layout。
+在 layout 记法（{ref}`zh_chap_data_layout`）中，我们把 TMEM Lane 轴写作 `TLane`，把 TMEM Col 轴写作 `TCol`。
+这些名字并不是为了替代官方硬件术语，而是 layout axis name，用来在 DSL 中显式标出 TMEM 维度。
+
+例如，一个 accumulator tile 可以写成：
+
+```text
+S[(128, N) : (1@TLane, 1@TCol)]
+```
+
+这表示 tile 沿硬件 Lane 维度有 128 行，沿硬件 Col 维度有 `N` 列。
+在 layout 记法中，这两个维度表现为 `TLane` 和 `TCol`。这个 layout 是直接映射：
+相邻 row 沿 `TLane` 移动，相邻 column 沿 `TCol` 移动。下图展示了这个网格：
+hardware Lane 沿 128 行向下，hardware Col 跨 column 横向展开。
+
+![作为二维网格的 TMEM：TLane 行 × TCol 列](../img/tmem_grid.png)
+
+重点是：TMEM 是 tile layout 故事的一部分。它不只是 Tensor Core 背后的隐藏 backing store。
+kernel 必须命名这块内存，从中分配 column，并使用匹配 `tcgen05` 指令读写方式的 layout。
+
+## 分配
+
+kernel 使用 TMEM 之前，必须先在其中预留空间。这不同于寄存器。寄存器由编译器分配，而 TMEM 由 kernel 显式分配。
+
+分配按 CTA 完成。CTA 中的一个 warp 请求一段 TMEM column。请求以 32 column 为单位，
+请求的 column 数会根据硬件分配规则向上取整。分配后，CTA 收到一个 base TMEM address。
+后续 `tcgen05` 指令使用这个 base address 访问预留区域。
+
+把 TMEM 看作一种有预算的 CTA 资源很有用，类似 shared memory。CTA 拥有自己分配到的 TMEM column。
+kernel 决定 accumulator、scale factor 或 temporary staging 需要多少 column。CTA 完成后，必须释放这次分配。
+
+这让 TMEM 成为 kernel resource planning 的一部分。更大的 accumulator tile 可能提升 Tensor Core 吞吐，
+但它会消耗更多 TMEM column。block-scaled MMA 可能需要额外 TMEM 空间来存放 scale factor。
+kernel 必须让这些用途适配可用 TMEM 预算，就像它必须让 shared-memory buffer 适配 SMEM 预算一样。
+
+## 读写 TMEM
+
+普通 `ld.shared` 和 `st.shared` 指令不能访问 TMEM。TMEM 是独立地址空间，因此数据通过专用 `tcgen05` 指令移动。
+
+主要有三条路径。
+
+第一条路径是 `tcgen05.ld`，它把数据从 TMEM load 到寄存器。这是 epilogue 在 MMA phase 之后使用的路径。
+accumulator 已经在 TMEM 中产生，但 epilogue 通常需要一个 register fragment，以便 cast、应用 elementwise operation，
+并存储最终结果。
+
+在 DSL 层面，TMEM load 分布在一个 warpgroup 上。它 lower 成四个 warp-level `tcgen05.ld` 操作，每个 warp 一个。
+每个 warp 处理 128 个 TMEM Lane row 中的 32 个，因此四个 warp 合起来覆盖完整 Lane 维度。
+在 layout 记法中，这个完整维度就是 `TLane` 轴。
+
+这条指令本身来自一组 load shape，例如 `.16x64b`、`.16x128b`、`.16x256b`、`.32x32b` 和 `.16x32bx2`，
+repeat factor 从 `.x1` 到 `.x128`。所选 shape 决定读取多少 TMEM column，以及每个 thread 接收多少寄存器。
+
+重要结果是 register fragment layout。对于常见 epilogue 路径，lane `l` 会接收来自 TMEM row `l / 4` 和两列的值。
+这产生了与早期世代直接从 MMA 暴露出的 per-lane accumulator fragment 同类的结构（{ref}`zh_chap_layout_generations`）。
+这种连续性很重要。它意味着即使 accumulator 在 compute phase 中位于 TMEM，Blackwell epilogue 仍然可以复用
+Ampere `mma` 或 Hopper `wgmma` 已经使用过的同一种 register-level cast 和 store 结构。
+
+![tcgen05.ld / st 以 m8n8 fragment 形式在 TMEM 累加器和寄存器之间移动数据（lane l → row l/4，两列）](../img/tcgen05_ldst.svg)
+
+第二条路径是 `tcgen05.st`，它把数据从寄存器 store 回 TMEM。这是 `tcgen05.ld` 的反方向。
+当 thread 已经持有 register fragment，并且需要把它放入 TMEM 时会使用它。
+例如，某些 operand 或 intermediate value 可能会先通过寄存器 stage，然后写入 TMEM，供后续 `tcgen05` 操作使用。
+
+第三条路径是 `tcgen05.cp`，它把数据从 shared memory copy 到 TMEM。这是一条 bulk copy 路径，
+常用于 block-scaled MMA 中的 scale factor。在这种情况下，TMA 或普通 thread code 先在 shared memory 中准备 scale data，
+`tcgen05.cp` 再把它移入 Tensor Core 期望的 TMEM layout。
+
+三条路径都是异步的。`tcgen05.ld`、`tcgen05.st` 或 `tcgen05.cp` 指令都可能在数据移动完成前返回。
+因此，kernel 必须在消费结果或复用 storage 前使用正确的 completion mechanism（{ref}`zh_chap_async_barriers`）。
+
+wait path 取决于指令。`tcgen05.ld` 通过 `tcgen05.wait::ld` 完成。`tcgen05.st` 通过 `tcgen05.wait::st` 完成。
+`tcgen05.cp` 像 `tcgen05.mma` 一样，通过 commit group 和 `mbarrier` 完成。
+如果数据从一组 thread 交给另一组 thread，kernel 可能还需要 fence，确保接收方 thread 按预期顺序看到已经完成的写入。
+
+TMEM 位于 Blackwell Tensor Core 数据路径的中间。TMA 把 operand stage 到 shared memory。
+`tcgen05.mma` 读取 operand，并累加到 TMEM 中。对于 block-scaled MMA，scale factor 也可以 stage 到 TMEM 中。
+compute phase 之后，`tcgen05.ld` 把 accumulator 带回寄存器，epilogue 转换并存储最终输出。
diff --git a/zh/img/flash_attention_main_handoff.png b/zh/img/flash_attention_main_handoff.png
new file mode 100644
index 00000000..471795ff
Binary files /dev/null and b/zh/img/flash_attention_main_handoff.png differ
diff --git a/zh/img/flash_attention_pipeline_v2.png b/zh/img/flash_attention_pipeline_v2.png
new file mode 100644
index 00000000..056f78ed
Binary files /dev/null and b/zh/img/flash_attention_pipeline_v2.png differ
diff --git a/zh/img/flash_attention_softmax_correction.png b/zh/img/flash_attention_softmax_correction.png
new file mode 100644
index 00000000..68a4c3e5
Binary files /dev/null and b/zh/img/flash_attention_softmax_correction.png differ
diff --git a/zh/img/gemm_perf.png b/zh/img/gemm_perf.png
new file mode 100644
index 00000000..e93fefa0
Binary files /dev/null and b/zh/img/gemm_perf.png differ
diff --git a/zh/img/ldstmatrix.svg b/zh/img/ldstmatrix.svg
new file mode 100644
index 00000000..9dae1d8f
--- /dev/null
+++ b/zh/img/ldstmatrix.svg
@@ -0,0 +1,3643 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:06.033932
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/memory_dataflow.png b/zh/img/memory_dataflow.png
new file mode 100644
index 00000000..2570f301
Binary files /dev/null and b/zh/img/memory_dataflow.png differ
diff --git a/zh/img/mma_block_scaled.svg b/zh/img/mma_block_scaled.svg
new file mode 100644
index 00000000..a27d9228
--- /dev/null
+++ b/zh/img/mma_block_scaled.svg
@@ -0,0 +1,2907 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:06.586642
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/mma_cg1_m128.svg b/zh/img/mma_cg1_m128.svg
new file mode 100644
index 00000000..ceef1d01
--- /dev/null
+++ b/zh/img/mma_cg1_m128.svg
@@ -0,0 +1,1065 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:06.196203
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/mma_cg1_m64.svg b/zh/img/mma_cg1_m64.svg
new file mode 100644
index 00000000..f0ad7484
--- /dev/null
+++ b/zh/img/mma_cg1_m64.svg
@@ -0,0 +1,1804 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:06.263085
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+ 
+  
+   
+   
+  
+ 
+
diff --git a/zh/img/mma_cg2_m128.svg b/zh/img/mma_cg2_m128.svg
new file mode 100644
index 00000000..343a2cf2
--- /dev/null
+++ b/zh/img/mma_cg2_m128.svg
@@ -0,0 +1,2013 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:06.463635
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/mma_cg2_m256.svg b/zh/img/mma_cg2_m256.svg
new file mode 100644
index 00000000..8ac6083a
--- /dev/null
+++ b/zh/img/mma_cg2_m256.svg
@@ -0,0 +1,2228 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:06.350390
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/pipe_depth2.png b/zh/img/pipe_depth2.png
new file mode 100644
index 00000000..af015f59
Binary files /dev/null and b/zh/img/pipe_depth2.png differ
diff --git a/zh/img/roofline.png b/zh/img/roofline.png
new file mode 100644
index 00000000..930a75f5
Binary files /dev/null and b/zh/img/roofline.png differ
diff --git a/zh/img/sf_scale_vec.svg b/zh/img/sf_scale_vec.svg
new file mode 100644
index 00000000..6a9ca563
--- /dev/null
+++ b/zh/img/sf_scale_vec.svg
@@ -0,0 +1,3319 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:27:09.026671
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/smem_descriptor.svg b/zh/img/smem_descriptor.svg
new file mode 100644
index 00000000..62aaf0a9
--- /dev/null
+++ b/zh/img/smem_descriptor.svg
@@ -0,0 +1,4405 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:58.635220
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+    
+    
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+    
+   
+   
+    
+    
+     
+     
+     
+    
+    
+    
+     
+      
+      
+      
+     
+     
+     
+     
+     
+    
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/swizzle_conflict.svg b/zh/img/swizzle_conflict.svg
new file mode 100644
index 00000000..e8ebd2c3
--- /dev/null
+++ b/zh/img/swizzle_conflict.svg
@@ -0,0 +1,4324 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:07.039615
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/tcgen05_ldst.svg b/zh/img/tcgen05_ldst.svg
new file mode 100644
index 00000000..cb849402
--- /dev/null
+++ b/zh/img/tcgen05_ldst.svg
@@ -0,0 +1,3847 @@
+
+
+
+ 
+  
+   
+    
+    2026-06-25T13:24:07.223393
+    image/svg+xml
+    
+     
+      Matplotlib v3.8.3, https://matplotlib.org/
+     
+    
+   
+  
+ 
+ 
+  
+ 
+ 
+  
+   
+  
+  
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+     
+    
+   
+   
+    
+    
+     
+      
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+   
+   
+    
+    
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+   
+    
+    
+     
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+      
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+     
+    
+   
+  
+ 
+ 
+  
+   
+  
+ 
+
diff --git a/zh/img/tma_sync_flow.png b/zh/img/tma_sync_flow.png
new file mode 100644
index 00000000..4e2f7ba0
Binary files /dev/null and b/zh/img/tma_sync_flow.png differ
diff --git a/zh/img/tmem_grid.png b/zh/img/tmem_grid.png
new file mode 100644
index 00000000..61a7c86f
Binary files /dev/null and b/zh/img/tmem_grid.png differ
diff --git a/zh/img/tmem_layout_v3.png b/zh/img/tmem_layout_v3.png
new file mode 100644
index 00000000..fad86331
Binary files /dev/null and b/zh/img/tmem_layout_v3.png differ
diff --git a/zh/img/warp_specialization_timeline.png b/zh/img/warp_specialization_timeline.png
new file mode 100644
index 00000000..99c11853
Binary files /dev/null and b/zh/img/warp_specialization_timeline.png differ
diff --git a/zh/index.md b/zh/index.md
new file mode 100644
index 00000000..1cd841be
--- /dev/null
+++ b/zh/index.md
@@ -0,0 +1,83 @@
+---
+orphan: true
+---
+
+# 面向 MLSys 的现代 GPU 编程
+
+机器学习系统位于现代 AI 工作负载的核心。在这些系统中，性能往往取决于少数几个
+GPU kernel 的质量。注意力 kernel、LLM prefill 和 decode kernel、低精度块缩放
+GEMM、融合 MoE 层，以及其他大型融合 kernel，都会直接影响训练和服务中的端到端速度。
+
+然而，要让这些 kernel 跑得快，仅有一串优化技巧还不够。现代 GPU 已经不再只是旧式设计的
+简单变体。近年的架构引入了更丰富的内存空间、新的访问模式，以及越来越专门化的执行单元。
+要写好它们，我们既需要对硬件形成清晰的心智模型，也需要实际理解高性能 kernel 是如何构建出来的。
+本书的目标正是同时培养这两种能力。
+
+本书遵循一条简单的路线：先理解 GPU 硬件，再学习我们将使用的编程模型，最后一步步构建
+state-of-the-art 的 kernel。我们的主要目标是 Blackwell 这一代 GPU，贯穿全书的主要例子是
+高速矩阵乘法（GEMM）和 FlashAttention。在这个过程中，我们也会研究 GPU 优化背后的核心要素：
+数据布局、异步数据移动和异步协同。
+
+这些材料源自卡内基梅隆大学的 [Machine Learning Systems](https://mlsyscourse.org/) 课程系列。
+为了让这些思想更容易学习、也更容易运行，本书使用 **TIRx** Python DSL，一步步构建真实的
+GPU kernel 示例。TIRx 贴近硬件，因此我们既能通过可运行代码学习，又能推理底层控制细节。
+
+## 本书结构
+
+- **第一部分，理解 GPU。** 本部分介绍 GPU 的整体组织方式、编写高速 kernel 的通用方法，以及
+  数据布局、异步内存操作和协同等关键概念。它建立了后续章节都依赖的硬件直觉。
+- **第二部分，TIRx 概览。** 本部分介绍 TIRx 的关键组成，它们是全书代码示例的基础。
+- **第三部分，GEMM：从分块到 SOTA。** 这是优化 tiled GEMM 的完整指南，逐步引入
+  TMA 流水线、持久化调度、warp specialization 和 2-CTA cluster。
+- **第四部分，Flash Attention 4。** 使用第三部分技术构建完整的注意力 kernel：两个 MMA，
+  中间插入 softmax，包含 online-softmax rescaling、causal masking 和 GQA。
+- **参考。** TIRx 语言参考和编译器内部机制。
+
+```{toctree}
+:caption: 第一部分，理解 GPU
+:maxdepth: 1
+
+chapter_background/index
+chapter_performance/index
+chapter_data_layout/index
+chapter_layout_generations/index
+chapter_tma/index
+chapter_tensor_cores/index
+chapter_tmem/index
+chapter_async_barriers/index
+chapter_clc/index
+```
+
+```{toctree}
+:caption: 第二部分，TIRx 概览
+:maxdepth: 1
+
+chapter_intro_tirx/index
+chapter_tirx_layout_api/index
+```
+
+```{toctree}
+:caption: "第三部分，GEMM：从分块到 SOTA"
+:maxdepth: 2
+
+chapter_gemm_basics/index
+chapter_gemm_async/index
+chapter_gemm_advanced/index
+```
+
+```{toctree}
+:caption: 第四部分，Flash Attention 4
+:maxdepth: 2
+
+chapter_flash_attention/index
+```
+
+```{toctree}
+:caption: 参考
+:maxdepth: 1
+
+appendix/index
+appendix/debugging_warp_specialized
+tirx_guide/arch/index
+tirx_guide/language_reference/index
+```
diff --git a/zh/tirx_guide/arch/index.rst b/zh/tirx_guide/arch/index.rst
new file mode 100644
index 00000000..3f6e573a
--- /dev/null
+++ b/zh/tirx_guide/arch/index.rst
@@ -0,0 +1,28 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _zh_chap_arch:
+
+编译器内部结构
+==============
+
+面向贡献者的 TIRx 编译器内部结构说明。
+
+.. toctree::
+   :maxdepth: 1
+
+   lowering_pipeline
diff --git a/zh/tirx_guide/arch/lowering_pipeline.rst b/zh/tirx_guide/arch/lowering_pipeline.rst
new file mode 100644
index 00000000..a4ce95dd
--- /dev/null
+++ b/zh/tirx_guide/arch/lowering_pipeline.rst
@@ -0,0 +1,200 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+TIRx lowering pipeline
+======================
+
+``tvm.compile(mod, target, tir_pipeline="tirx")`` 会把你写好的 TIRx module
+送过 **tirx pipeline**。这是一串有序的 TIR pass，会把你写的高层结构
+（tile primitive、带 ``TileLayout`` 类型的 buffer、execution-scope id）转换成
+拆分后的 **host** + **device** 函数，然后由 CUDA backend 渲染成源码。
+pipeline 定义在 ``python/tvm/tirx/compilation_pipeline.py``（``tirx_pipeline``）
+中；本页按顺序走过这些 pass。
+
+它所在的位置
+------------
+
+``tvm.compile`` 会先绑定 target，运行 **tirx pipeline**（下面这些 module-level
+pass），然后分别对 host 和 device 函数应用 **finalization** pass，最后把每个
+device 函数交给 CUDA code generator：
+
+.. code-block:: text
+
+    authored TIRx  ──BindTarget──▶  tirx_pipeline  ──▶  host func  ──host finalize──▶  C/LLVM
+                                          │
+                                          └──────────▶  device func ──device finalize──▶  CUDA
+
+Pass 列表
+---------
+
+``tirx_pipeline`` module pass 会应用下面这个精确顺序（其中少数 pass 受
+``PassContext`` config 控制）：
+
+.. list-table::
+   :header-rows: 1
+   :widths: 6 32 62
+
+   * - #
+     - Pass
+     - 作用
+   * - 1
+     - ``LowerTIRx``
+     - 核心 lowering，见下文 `Inside LowerTIRx`_
+   * - 2
+     - ``UnifyThreadBinding``
+     - 合并等价的 thread-axis binding，让每个 ``threadIdx`` / ``blockIdx``
+       轴只声明一次
+   * - 3
+     - ``StmtSimplify``
+     - 语句级算术化简（arith analyzer）
+   * - 4
+     - ``LowerTIRxOpaque``
+     - 将剩余 opaque TIRx construct lower 成普通 TIR
+   * - 5
+     - ``FlattenBuffer``
+     - 把多维 ``BufferLoad`` / ``BufferStore`` flatten 成 1-D
+   * - 6
+     - ``BF16ComputeLegalize``
+     - 把 ``bfloat16`` compute 重写成合法形式（上转为 f32）
+   * - 7
+     - ``NarrowDataType(32)``
+     - 在可证明安全时，把 index/loop ``PrimExpr`` dtype 缩窄到 32-bit
+   * - 8
+     - ``VectorizeLoop``
+     - 把 ``T.vectorized`` loop 转成 vector op（若设置 ``tir.disable_vectorize``
+       则跳过）
+   * - 9
+     - ``UnrollLoop``
+     - 展开标记为 ``T.unroll`` 的 loop（以及小的常量 loop）
+   * - 10
+     - ``StmtSimplify``
+     - 再次化简，因为 vectorize/unroll 暴露了常量
+   * - 11
+     - ``CommonSubexprElim``
+     - 把重复子表达式 hoist 成临时变量（若设置 ``tir.disable_cse_tir``
+       则跳过）
+   * - 12
+     - ``FP8ComputeLegalize``
+     - 把 ``float8`` compute 重写成合法形式
+   * - 13
+     - ``VerifyMemory``
+     - 检查 host-side 代码不会直接解引用 device memory（安全闸门）
+   * - 14
+     - ``AnnotateEntryFunc``
+     - 将单个 PrimFunc 标记为 module entry point
+   * - 15
+     - ``SplitHostDevice``
+     - 在 ``launch_thread`` 边界处，把每个 kernel 拆成 **host** 函数和
+       **device** 函数
+   * - 16
+     - ``MakePackedAPI``
+     - 将 host 函数重写成 packed-func ABI（TVM launcher 调用的形式）
+   * - 17
+     - ``FP8StorageLegalize``
+     - legalize ``float8`` storage（打包进受支持的 container type）
+   * - 18
+     - ``BF16StorageLegalize``
+     - legalize ``bfloat16`` storage
+
+随后 **Finalization** 会按函数类型运行：
+
+- **host**：``LowerTVMBuiltin``（lower ``tvm_*`` builtin）、``LowerIntrin``
+  （target-specific intrinsic）
+- **device**：``LowerWarpMemory``（warp-scoped buffer → shuffle）、``StmtSimplify``、
+  ``LowerIntrin``
+
+Inside LowerTIRx
+----------------
+
+``LowerTIRx`` 本身也是一个小序列（``src/tirx/transform/lower_tirx.cc``）：
+
+.. code-block:: text
+
+    LowerTIRx = Sequential([ TilePrimitiveDispatch, LowerTIRxCleanup ])
+
+- **``TilePrimitiveDispatch``** 会把每个 ``TilePrimitiveCall``（``copy``、
+  ``gemm``、``reduction`` 等）替换成所选 backend dispatch 发出的 body，
+  也就是它的 variant-selection 和 codegen。
+- **``LowerTIRxCleanup``** 会运行 ``LayoutApplier``：把每个带
+  ``TileLayout`` 类型的 buffer access 解析成具体物理地址算术
+  （``addr = data + elem_offset + layout.apply(coord)``），flatten buffer，并
+  lower execution-scope id（``T.cta_id`` / ``T.thread_id`` / … 通过
+  ``launch_thread`` 变成 ``blockIdx`` / ``threadIdx``）。
+
+因此经过 ``LowerTIRx`` 后，module 就是普通 TIR：不再有 tile primitive，
+不再有 ``TileLayout`` 间接层，scope id 也解析成了 thread axis。
+
+一个完整例子
+------------
+
+来看一个一行 scale kernel：
+
+.. code-block:: python
+
+    @T.prim_func
+    def scale(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (256,), "float32")
+        B = T.match_buffer(B_ptr, (256,), "float32")
+        T.device_entry(); bx = T.cta_id([1]); tx = T.thread_id([256])
+        B[tx] = A[tx] * T.float32(2.0)
+
+**经过 ``LowerTIRx`` 后**，scope id 已经是真实 thread axis，layout 也已经应用
+（``A_1`` / ``B_1`` 是 flattened 1-D view）：
+
+.. code-block:: python
+
+    with T.launch_thread("blockIdx.x", 1) as blockIdx_x:
+        threadIdx_x = T.launch_thread("threadIdx.x", 256)
+        bx: T.let = blockIdx_x
+        tx: T.let = threadIdx_x
+        B_1[threadIdx_x] = A_1[threadIdx_x] * T.float32(2.0)
+
+**经过 ``SplitHostDevice`` + ``MakePackedAPI`` 后**，一个函数变成两个：
+一个 host launcher 和一个 device kernel：
+
+.. code-block:: python
+
+    @I.ir_module
+    class Module:
+        def main(...):          # host: packed-API launcher (computes the grid/block, launches)
+            ...
+        def scale_kernel(...):  # device: the __global__ body, run on the GPU
+
+随后 CUDA backend 会把 ``scale_kernel`` 渲染成 ``__global__`` 函数
+（``B_ptr[threadIdx.x] = A_ptr[threadIdx.x] * 2.0f``）。
+
+自己复现
+--------
+
+你可以手动运行 pipeline 的任意前缀来检查某个阶段；这些文档中的 IR snippet
+就是这样生成的：
+
+.. code-block:: python
+
+    from tvm.tirx import transform as TT
+
+    target = tvm.target.Target("cuda")
+    mod = TT.BindTarget(target.with_host("llvm"))(tvm.IRModule({"main": scale}))
+    mod = TT.LowerTIRx()(mod)         # tile primitives dispatched, layouts applied
+    print(mod.script())               # inspect the lowered TIRx IR
+
+或者编译整个 module，然后读取生成的 CUDA：
+
+.. code-block:: python
+
+    exe = tvm.compile(tvm.IRModule({"main": scale}), target=target, tir_pipeline="tirx")
+    print(exe.mod.imports[0].inspect_source())
diff --git a/zh/tirx_guide/language_reference/cuda/buffers.rst b/zh/tirx_guide/language_reference/cuda/buffers.rst
new file mode 100644
index 00000000..32494f28
--- /dev/null
+++ b/zh/tirx_guide/language_reference/cuda/buffers.rst
@@ -0,0 +1,471 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+Buffer 与内存
+=============
+
+参数 buffer 通过 ``T.match_buffer`` 绑定；scratch buffer 则在函数体里用下面
+两类声明 API 之一创建。可以用 ``A[i, j]`` 索引 buffer，用
+``A[m0:m0+BM, 0:BK]`` 切片（得到 ``BufferRegion``），也可以用
+``A.ptr_to([i, j])`` 取得指针，或用 ``A.data`` 取得原始 data pointer。
+
+声明 buffer
+-----------
+
+创建 buffer 有两个基础 API：
+
+- ``T.alloc_buffer(shape, dtype, scope=..., ...)`` — **分配新的存储空间**
+  （发出一个 ``AllocBuffer`` 节点）并返回 ``Buffer``。``T.alloc_shared`` /
+  ``T.alloc_local`` 只是 ``alloc_buffer`` 加上 ``scope="shared"`` /
+  ``scope="local"`` 的简写。
+- ``T.decl_buffer(shape, dtype, data=..., ...)`` — 在已有指针 ``data`` 上
+  **声明一个 view** （不分配）；用于 alias 或 reinterpret 存储空间，例如 pool
+  的子区域或 tensor-memory address。若 ``data=None``，它会像 ``alloc_buffer``
+  一样分配。
+
+buffer 的 ``data`` pointer 是一个 immutable ``Var``（``alloc_buffer`` 会定义它；
+``decl_buffer`` 接收它）。如果要让 buffer 背后使用一个指针 *表达式*，请先
+绑定该表达式——见 :doc:`data_types`。
+
+二者共享同一种 descriptor；最重要的参数如下：
+
+.. list-table::
+   :header-rows: 1
+   :widths: 28 72
+
+   * - 参数
+     - 含义
+   * - ``dtype``
+     - 元素类型，例如 ``"float32"``、``"float16"``、``"float4_e2m1fn"`` 等
+   * - ``shape``
+     - 逻辑形状（一组 extent）
+   * - ``layout``
+     - 物理映射（:ref:`TileLayout `）；``"default"`` = dense
+       row-major
+   * - ``elem_offset`` / ``allocated_addr``
+     - ``elem_offset``（或 ``byte_offset``）把一个 *view* 放到 ``data`` 内部的
+       某个偏移；``allocated_addr`` 携带预先分配的地址（tensor memory）
+   * - ``align``
+     - data pointer 的对齐字节数
+
+``scope`` 参数选择内存空间：
+
+.. list-table::
+   :header-rows: 1
+   :widths: 26 22 52
+
+   * - Scope
+     - 简写
+     - 内存
+   * - ``"global"``
+     - (default)
+     - device global memory
+   * - ``"shared"``
+     - ``T.alloc_shared``
+     - static shared memory（``__shared__``）
+   * - ``"shared.dyn"``
+     - (pool)
+     - dynamic shared memory（pooled，见下文）
+   * - ``"local"``
+     - ``T.alloc_local``
+     - per-thread register
+   * - ``"tmem"``
+     - (TMEM pool)
+     - Blackwell tensor memory（见下文）
+
+.. code-block:: python
+
+    A = T.match_buffer(A_ptr, (M, K), "float16", align=16)   # parameter buffer
+    As = T.alloc_shared((BM, BK), "float16")                 # new shared tile
+    acc = T.alloc_local((4,), "float32")                     # register accumulator
+    view = T.decl_buffer((BM, BK), "float16", data=As.data)  # a view over As
+
+**基于 pointer 的 buffer 本质上只是 pointer 上的一层 metadata。** 对任何
+非 tmem buffer，声明都只是一个 pointer 加一个 layout；索引会解析成地址：
+
+    addr(buffer[coord]) = buffer.data + elem_offset + layout.apply(coord, shape=shape)["m"]
+
+（``layout.apply`` 返回每个轴的映射；其中 ``"m"`` 分量是元素偏移。）因此
+*同一个* 逻辑访问会完全根据 buffer metadata 编译成不同的地址算术。对一个
+4×8 区域写 ``B[i, j] = A[i, j] + 1``，并用四种方式声明 ``B``：
+
+.. code-block:: python
+
+    from tvm.tirx.layout import TileLayout, S
+
+    B = T.match_buffer(p, (4, 8), "float32")                                       # row-major
+    B = T.match_buffer(p, (4, 8), "float32", layout=TileLayout(S[(4, 8):(1, 4)]))  # column-major
+    B = T.match_buffer(p, (4, 8), "float32", elem_offset=64)                       # shifted view
+    B = T.match_buffer(p, (4, 8), "float32", layout=TileLayout(S[(4, 8):(16, 1)])) # row stride 16
+
+每种声明都会让 ``B[i, j]`` 在生成的 CUDA 中 lower 成不同索引（``A[i, j]``
+load 仍是 ``i*8 + j``，只有 ``B`` 的 metadata 改了）：
+
+.. code-block:: c++
+
+    B_ptr[((i * 8) + j)]        = ...;   // row-major:        i*8 + j
+    B_ptr[((j * 4) + i)]        = ...;   // column-major:     j*4 + i
+    B_ptr[(((i * 8) + j) + 64)] = ...;   // elem_offset=64:   i*8 + j + 64
+    B_ptr[((i * 16) + j)]       = ...;   // row stride 16:    i*16 + j
+
+Shared memory
+-------------
+
+shared memory 有两种形式：**static** （编译期固定大小）和 **dynamic** （launch
+时确定大小）；此外还有一个 pool helper 用来管理 dynamic 情况。
+
+Static
+~~~~~~
+
+最简单的 shared buffer 是 **static** 形式：``T.alloc_shared``（也就是
+``scope="shared"``），大小在编译期确定。把数据 stage 进去，调用
+``cta_sync`` 让整个 block 都看到这些写入，然后再读出：
+
+.. code-block:: python
+
+    @T.prim_func
+    def smem_demo(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (128,), "float32")
+        B = T.match_buffer(B_ptr, (128,), "float32")
+        T.device_entry()
+        bx = T.cta_id([1])
+        tx = T.thread_id([128])
+        sm = T.alloc_shared((128,), "float32")   # static shared memory
+        sm[tx] = A[tx]
+        T.cuda.cta_sync()
+        B[tx] = sm[tx] * T.float32(2.0)
+
+它会 lower 成普通 ``__shared__`` array（省略生成 CUDA 的样板部分）：
+
+.. code-block:: c++
+
+    extern "C" __global__ void __launch_bounds__(128)
+    smem_demo_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) {
+      int tx = ((int)threadIdx.x);
+      __shared__ alignas(64) float sm_ptr[128];      // T.alloc_shared
+      sm_ptr[tx] = A_ptr[tx];
+      __syncthreads();                               // T.cuda.cta_sync()
+      B_ptr[tx] = sm_ptr[tx] * 2.0f;
+    }
+
+Dynamic
+~~~~~~~
+
+**Dynamic** shared memory（``scope="shared.dyn"``）的大小按 launch 确定（即
+``sharedMemBytes`` launch 参数），不是编译期确定。一个 kernel **只能有一个**
+dynamic-shared allocation，也就是 *arena*。因此你只分配一次 arena，再用
+``T.decl_buffer`` 把每个 buffer 声明成它内部的一个 view：``data=`` 传 arena
+pointer，并设置 ``elem_offset``：
+
+.. code-block:: python
+
+    arena = T.alloc_buffer((128,), "float32", scope="shared.dyn")   # the one arena
+    As = T.decl_buffer((64,), "float32", data=arena.data, scope="shared.dyn")                 # offset 0
+    Bs = T.decl_buffer((64,), "float32", data=arena.data, elem_offset=64, scope="shared.dyn") # offset 64
+    As[tx] = A[tx]
+    Bs[tx] = B[tx]
+    T.cuda.cta_sync()
+    C[tx] = As[tx] + Bs[tx]
+
+两个 view 共享同一个 ``extern __shared__`` arena（省略生成 CUDA 的样板部分；
+为了清楚起见，arena 命名为 ``smem``）：
+
+.. code-block:: c++
+
+    extern __shared__ __align__(64) float smem[];   // the one dynamic-shared arena
+    smem[tx]      = A_ptr[tx];                       // As — view at offset 0
+    smem[tx + 64] = B_ptr[tx];                       // Bs — view at offset 64
+    __syncthreads();
+    C_ptr[tx] = smem[tx] + smem[tx + 64];
+
+（两次单独调用 ``alloc_buffer(scope="shared.dyn")`` 是错误的——*只允许一个
+dynamic shared memory allocation*。）因此 static shared memory 在编译期定大小
+（``__shared__ T x[N];``）；dynamic shared memory 则是这个 launch-sized arena，
+其内部不同偏移处声明多个 view。
+
+.. note::
+
+   **TVM 如何标注 dynamic-shared 大小。** arena 的大小在编译期已知（这里
+   ``128`` 个 float = ``512`` bytes）。lowering 期间，TVM 会向 device kernel 的
+   ``tirx.kernel_launch_params`` 追加一个 ``"tirx.use_dyn_shared_memory"`` tag；
+   host launcher 会计算总字节数，并把它作为最后一个 launch 参数传入：
+
+   .. code-block:: python
+
+       # device kernel attribute:
+       "tirx.kernel_launch_params": ["blockIdx.x", "threadIdx.x", "tirx.use_dyn_shared_memory"]
+
+       # host-side launch call  (..., gridDim.x, blockDim.x, dyn_shared_bytes):
+       T.call_packed("dyn_kernel", A.data, B.data, C.data, 1, 64, 512)
+
+   运行时，这个 ``512`` 会成为 ``cuLaunchKernelEx`` 调用中的
+   ``config.sharedMemBytes``。你不需要手动设置它；它由 ``shared.dyn``
+   allocation 的大小推导而来。
+
+Pool sugar
+~~~~~~~~~~
+
+``T.SMEMPool`` 会自动处理 arena bookkeeping：它用 bump allocator 分配偏移，
+因此你不用手写 ``decl`` view。除了 ``alloc`` / ``commit`` 之外，它还提供
+per-buffer ``align=``、一个为你构建 MMA-compatible swizzle layout 的
+``alloc_mma`` helper，以及用于回退 cursor、复用空间的 ``move_base_to``：
+
+.. code-block:: python
+
+    pool = T.SMEMPool()                          # bump allocator over shared.dyn
+    As = pool.alloc((BM, BK), "float16", align=128)   # carve a tile
+    Bs = pool.alloc((BK, BN), "float16", align=128)
+    Cs = pool.alloc_mma((BM, BN), "float16")     # MMA-compatible, swizzle inferred
+    pool.commit()                                 # finalize the pool's size
+    # pool.move_base_to(offset) rewinds the cursor to reuse space
+
+TMEM pool（见下文 `Tensor memory`_）构建在 ``SMEMPool`` 之上。
+
+Registers
+---------
+
+per-thread scratch 存在 register 中。用 ``T.alloc_local(shape, dtype)``
+（也就是 ``scope="local"``）分配：它对每个线程私有，并 lower 成保存在
+register 中的 local array。
+
+.. code-block:: python
+
+    r = T.alloc_local((4,), "float32")   # per-thread register array
+    for k in T.unroll(4):
+        r[k] = A[tx, k]
+    # ... compute on r[0..3] ...
+
+.. code-block:: c++
+
+    alignas(64) float r_ptr[4];          // per-thread, register-resident
+    r_ptr[0] = A_ptr[tx * 4 + 0];
+    r_ptr[1] = A_ptr[tx * 4 + 1];
+    // ...
+
+.. note::
+
+   ``alignas(64)`` 是 *默认* buffer alignment：buffer 的 ``data_alignment``
+   默认是 ``runtime::kAllocAlignment``（64 bytes），CUDA codegen 会把它印到每个
+   allocation 上，包括这种对齐没有意义的 per-thread ``local`` array。对于这些
+   register-resident array，它 **没有性能影响**：带静态可解析索引的 thread-local
+   array 会被 nvcc/ptxas 提升到 register（scalar replacement of aggregates，SROA），
+   因此它从不进入可寻址 local memory，对齐就是 no-op。（如果动态索引 array
+   spill 到 local memory，它确实会带上这种过度对齐，但那是不常见情况。）这种
+   register local 的过度对齐是一个已知粗糙边角，我们计划修掉（对 ``local``
+   scope 使用 dtype 的自然对齐）。
+
+Scalar
+~~~~~~
+
+scalar 只是一个 **单元素** register array；严格来说，不需要单独概念。你可以
+分配一个 size-1 的 ``local`` buffer，并用 ``[0]`` 索引：
+
+.. code-block:: python
+
+    phase = T.alloc_local((1,), "int32")   # 1-element register array
+    phase[0] = 0
+    while phase[0] < 4:
+        acc = acc + A[tx, phase[0]]
+        phase[0] += 1
+
+但到处写 ``phase[0]`` 很笨重，所以 **scalar** 正是这件事的语法糖：一个可以
+**按名字** 读写的单元素 register buffer：
+
+.. code-block:: python
+
+    phase: T.int32 = 0                 # mutable scalar (sugar for the above)
+    while phase < 4:
+        acc = acc + A[tx, phase]
+        phase += 1
+
+    s = T.local_scalar("int32")        # explicit form; assign by name (s = ..., not s[0])
+    acc: T.float32 = 0.0               # a type-annotated assignment also makes one
+
+二者不只是相似，而是会 parse 成 **结构完全相同的 TIRx**。这个语法糖完全在
+parser 中消解：``phase: T.int32`` *就是* 那个单元素 ``local`` buffer，
+``phase`` / ``phase += 1`` *就是* ``phase[0]`` / ``phase[0] += 1``。两个
+kernel 上的 ``tvm.ir.assert_structural_equal`` 会通过，printer 甚至会把显式
+``alloc_local`` + ``[0]`` 形式 **重新打印回** scalar 形式；因此 parsing 完成
+后没有任何差别。二者都会 lower 成同一个
+``alignas(64) int phase_ptr[1];``；scalar 只是让你省掉 ``[0]``。
+（``T.local_scalar`` / ``T.shared_scalar`` / ``T.alloc_scalar`` 可以显式选择 scope。）
+
+.. note::
+
+   **为什么不用** ``Var``\ **？** TIRx ``Var`` 是 *immutable* 的：它是一个单次
+   static binding（也就是下面的 ``T.let`` 产生的东西）。scalar 需要是
+   *mutable* 的，因为你会在循环和 accumulator 中反复给它赋值；所以它必须
+   由一个可重复 store 的单元素 buffer 支撑，而不是 ``Var``。
+
+``let``
+~~~~~~~
+
+``T.let`` binding 是 **immutable** 的：一个单独的 ``LetStmt`` （有名字的值，
+不是 buffer）。它适合派生常量：
+
+.. code-block:: python
+
+    n: T.let = M * K               # immutable binding (LetStmt)
+    half: T.let[T.int32] = N // 2  # ... with an explicit type
+
+它会 lower 成 **普通 scalar C 变量**，不是 buffer（没有 array，没有 ``[0]``）。
+对于 ``half: T.let = m * 2``（其中 ``m`` 是运行时值）：
+
+.. code-block:: c++
+
+    int half = m * 2;     // the `let` -> a const-like local
+
+因为值是 immutable 的，simplifier 可以自由传播并对它做 CSE，所以在使用点你
+经常会直接看到 ``m * 2`` 被替换进去（或通过 common-subexpression 临时变量
+共享），而不是引用 ``half``。
+
+.. note::
+
+   **为什么还需要 immutable binding？** 因为值不会变化，arithmetic analyzer 可以
+   把 var 绑定到这个值上（化简 ``LetStmt`` 时调用
+   ``analyzer.Bind(var, value)``），所以关于这个值证明出的事实——常量边界、
+   modular set（可整除性 / 对齐）、范围——都会 **传播到每次使用**。这会反过来
+   支持 index simplification、bounds-check elimination，以及 alignment/vectorization
+   决策。*mutable* scalar 是一次 memory load（``buf[0]``）：analyzer 不能假设它
+   保持不变，所以这些性质不会传递。``let`` 也是一个纯值：没有 allocation，
+   可以自由 inline / substitute / CSE；而 scalar 是一个带 load/store 语义的
+   单元素 buffer。
+
+Tensor memory
+-------------
+
+Blackwell *tensor memory* 不是普通 scratch scope：它必须用 warp-uniform 的
+``T.ptx.tcgen05.alloc`` / ``tcgen05.dealloc`` intrinsic 显式 reserve 和 free；
+每个 tensor 都是其中的一个 view，通过
+``T.decl_buffer(..., scope="tmem", allocated_addr=, layout=)``
+声明。``allocated_addr``（列偏移）是必需的，tensor-core dispatch 会 assert 它；
+因此 ``T.alloc_buffer(scope="tmem")``（不会设置它）不能工作。不同于 shared
+memory，tensor memory 不能直接寻址：只能通过 ``tcgen05`` 的 ``mma`` / ``ld`` /
+``st`` / ``cp`` 读写。
+
+手写时，一个 warp 会把 allocation 发到一个 shared slot 中；你再把每个 tensor
+``decl`` 成某个列偏移处的 view；最后由一个 warp 释放它：
+
+.. code-block:: python
+
+    addr = T.alloc_shared((1,), "uint32")             # slot for the allocated base
+    if warp_id == alloc_warp:                         # tcgen05.alloc is warp-uniform
+        T.ptx.tcgen05.alloc(T.address_of(addr), n_cols=512, cta_group=cta_group)
+    acc = T.decl_buffer((CTA_M, 512), "float32", scope="tmem",
+                        allocated_addr=0, layout=tmem_layout)   # view at column 0
+    # ... use acc as a gemm_async / copy_async operand ...
+    if warp_id == alloc_warp:
+        T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+        T.ptx.tcgen05.dealloc(addr, n_cols=512, cta_group=cta_group)
+
+列偏移和 ``tmem_layout`` （datapath D/F layout）由你自己管理。下面的 pool 发出
+的正是这套序列。
+
+Pool
+~~~~
+
+``T.TMEMPool`` 把这些全部封装起来：warp-uniform alloc/dealloc、列方向
+bump-allocation，以及 datapath layout：
+
+.. code-block:: python
+
+    tmem_addr = pool.alloc((1,), "uint32")          # pool = the kernel's smem pool
+    tmem_pool = T.TMEMPool(pool, total_cols=512, cta_group=cta_group,
+                           tmem_addr=tmem_addr)
+    acc = tmem_pool.alloc((CTA_M, 512), "float32")  # allocated_addr set for you
+    tmem_pool.commit()                               # emits tcgen05.alloc (one warp)
+    # ... use acc ...
+    tmem_pool.dealloc()                              # emits tcgen05.dealloc (one warp)
+
+完整示例见 Part III 的 GEMM kernel。
+
+Buffer API
+----------
+
+``Buffer`` 是 pointer 上的 metadata（见上文 *声明 buffer*），因此它的大部分
+方法都是 *compile-time* reshape/reinterpret：要么改变索引算术，要么把指针交给
+你；它们本身不会发出 runtime op。常见方法如下：
+
+.. list-table::
+   :header-rows: 1
+   :widths: 34 66
+
+   * - 方法
+     - 含义
+   * - ``B.data``
+     - 原始 data pointer（一个 ``Var``）；打印为 ``B_ptr``
+   * - ``B.ptr_to([i, j])``
+     - 指向某个元素的 typed pointer（``address_of``）；打印为 ``&B_ptr[…]``
+   * - ``B.vload([i], dtype="float32x4")`` / ``B.vstore([i], v)``
+     - vectorized load / store；打印为 ``*(float4*)(B_ptr + …)``
+   * - ``B.view(*shape, layout=…)``
+     - 以新的 shape/layout reinterpret 同一份存储（不 copy）
+   * - ``B.local(*shape, layout=…)``
+     - 调用线程在 ``local`` buffer 中的私有 register slice
+   * - ``B.permute(*dims)``
+     - 轴被 permute 后的 view（转置 layout）
+   * - ``B.access_ptr(mask, …)``
+     - masked access pointer（``tvm_access_ptr`` builtin），用于把 region 传给
+       intrinsic
+
+**Pointer — ``ptr_to`` / ``data``。** ``ptr_to`` 用来把元素地址交给 intrinsic
+或 inline function；``data`` 是 base pointer：
+
+.. code-block:: python
+
+    B[tx] = T.cuda.func_call("ld", A.ptr_to([tx]), source_code=SRC, return_type="float32")
+
+.. code-block:: c++
+
+    B_ptr[tx] = ld(&A_ptr[tx]);          // ptr_to([tx]) -> &A_ptr[tx];  A.data -> A_ptr
+
+**Vectorized access — ``vload`` / ``vstore``。** 把多个元素作为一次宽传输来
+移动（另见 :doc:`data_types`）：
+
+.. code-block:: python
+
+    B.vstore([tx * 4], A.vload([tx * 4], dtype="float32x4"))
+
+.. code-block:: c++
+
+    *(float4*)(B_ptr + tx * 4) = *(float4*)(A_ptr + tx * 4);
+
+**Reshape / reinterpret — ``view`` / ``permute``。** 二者都是纯 metadata；
+data pointer 不变，只有索引算术不同。``A.view(64, 4)`` 会把 256 元素 buffer
+看成 ``64×4``；``A.permute(1, 0)`` 会转置轴：
+
+.. code-block:: python
+
+    A2 = A.view(64, 4);     y = A2[tx, 0] + A2[tx, 3]   # A2[tx, j] -> A_ptr[tx*4 + j]
+    At = A.permute(1, 0);   z = At[i, j]                # At[i, j]  -> A_ptr[j*4 + i]
+
+.. code-block:: c++
+
+    A2_ptr[tx * 4]  /* +3 */                 // view: row-major 64x4 index
+    At_ptr[(j * 4) + i]                       // permute: swapped strides
+
+**Register — ``local``。** 把 thread-axis ``local`` layout 分解成调用线程的
+flat register bundle（tile primitive 中大量使用）：
+
+.. code-block:: python
+
+    R  = T.alloc_buffer((32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)]))
+    Rl = R.local(8)          # this lane's 8 registers
+
+.. code-block:: c++
+
+    alignas(64) float Rl_ptr[8];             // the lane's private registers
diff --git a/zh/tirx_guide/language_reference/cuda/control_flow.rst b/zh/tirx_guide/language_reference/cuda/control_flow.rst
new file mode 100644
index 00000000..e5e360df
--- /dev/null
+++ b/zh/tirx_guide/language_reference/cuda/control_flow.rst
@@ -0,0 +1,122 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+控制流
+======
+
+控制流包括 ``if``、循环家族和 ``while``；它们都会映射到直观对应的
+CUDA 结构。
+
+if
+--
+
+Python 的 ``if`` / ``else`` 会变成 CUDA 的 ``if`` / ``else``。可以用
+thread/lane 比较来保护某段工作，也可以用 ``T.ptx.elect_sync()`` 选出
+一个发出指令的线程：
+
+.. code-block:: python
+
+    if tx < 128:
+        A[tx] = A[tx] * T.float32(2.0)
+    else:
+        A[tx] = A[tx] + T.float32(1.0)
+
+    if T.ptx.elect_sync():
+        ...                              # one elected lane (e.g. to issue TMA/MMA)
+
+.. code-block:: c++
+
+    if (((int)threadIdx.x) < 128) {
+      A_ptr[tx] = A_ptr[tx] * 2.0f;
+    } else {
+      A_ptr[tx] = A_ptr[tx] + 1.0f;
+    }
+
+如果只是表达式级选择（不产生分支），使用 ``T.if_then_else(cond, a, b)``。
+它会 lower 成三元表达式，因此不会引入 control-flow divergence：
+
+.. code-block:: c++
+
+    O_ptr[tx] = (A_ptr[tx] > 0.0f) ? A_ptr[tx] : 0.0f;
+
+Uniform 与 divergent 控制流
+---------------------------
+
+像 ``if tx < 128`` 这样的 per-thread guard 对普通工作没有问题，但
+**collective** 操作必须被它要同步的所有线程 *一致地* 到达。
+
+例如，``T.cuda.cta_sync()`` 会映射到 ``__syncthreads()``，它要求 thread block
+中的所有线程都到达。它绝不能放在 thread-divergent 或 warpgroup-divergent
+分支里：如果放进 ``if wg_id == 0:``，其他 warpgroup 永远不会到达，kernel
+就会 deadlock。若只需要同步一个 warpgroup，请使用 warpgroup-scoped
+``T.cuda.warpgroup_sync(id)`` （见 :ref:`zh_chap_gemm_advanced` 和
+:doc:`threads_sync`）。
+
+barrier 初始化也要同样小心。``mbarrier`` 的 ``.init()`` 会 lower 成一个
+single-thread guard（``if (threadIdx.x < 1)``）。如果再把它嵌进另一个
+divergent 分支，barrier 可能保持未初始化，导致未定义的 launch failure。
+
+loop
+----
+
+循环有四种形式；普通 Python ``range`` 会变成 ``T.serial``：
+
+- ``T.serial(n)`` — 顺序循环（ptxas 仍可能 unroll 它）。
+- ``T.unroll(n)`` — 完全 unroll（展开成直线语句）。
+- ``T.vectorized(n)`` — vectorized loop。
+- ``T.grid(*extents)`` — 嵌套循环。
+
+循环内部可以使用 ``break`` / ``continue``。
+
+.. code-block:: python
+
+    for i, j in T.grid(8, 8):
+        B[i, j] = T.max(A[i, j], T.float32(0.0))
+
+.. code-block:: c++
+
+    for (int i = 0; i < 8; ++i)
+      for (int j = 0; j < 8; ++j)
+        B_ptr[i * 8 + j] = max(A_ptr[i * 8 + j], 0.0f);
+
+``T.unroll(4)`` 则会展开成四条直线语句，不再保留循环。
+
+while
+-----
+
+``while`` 循环会一直运行到条件为 false。请使用 mutable scalar counter
+（见 :doc:`buffers`）：
+
+.. code-block:: python
+
+    i: T.int32 = 0
+    while i < 64:
+        A[i] = A[i] + T.float32(1.0)
+        i += 1
+
+它会 lower 成一个带 early-exit ``break`` 的 ``while (1)`` （counter 是一个
+单元素 register buffer）：
+
+.. code-block:: c++
+
+    int i_ptr[1];
+    i_ptr[0] = 0;
+    while (1) {
+      if (!(i_ptr[0] < 64)) { break; }
+      A_ptr[i_ptr[0]] = A_ptr[i_ptr[0]] + 1.0f;
+      i_ptr[0] = i_ptr[0] + 1;
+    }
diff --git a/zh/tirx_guide/language_reference/cuda/data_types.rst b/zh/tirx_guide/language_reference/cuda/data_types.rst
new file mode 100644
index 00000000..18588928
--- /dev/null
+++ b/zh/tirx_guide/language_reference/cuda/data_types.rst
@@ -0,0 +1,117 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+数据类型与表达式
+================
+
+每个 TIRx 表达式都同时携带一个低层 **dtype** 和一个高层 **type**。
+
+表达式 dtype
+------------
+
+``PrimExpr`` 的 ``.dtype`` 是它的标量（或向量）元素类型，例如 ``float32``、
+``float16``、``bfloat16``、``int32``、``uint8``、``bool``、低精度的
+``float8_e4m3fn`` / ``float4_e2m1fn``、``handle``（指针），以及
+``float32x4`` 这类向量形式。每种 dtype 都会打印成对应的 CUDA 类型。下面
+示例展示跨多种 dtype 分配 local/shared buffer，以及一次 vectorized
+``float32x4`` load/store：
+
+.. code-block:: python
+
+    @T.prim_func
+    def dtypes(A_ptr: T.handle, O_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (256,), "float32")
+        O = T.match_buffer(O_ptr, (256,), "float32")
+        T.device_entry(); bx = T.cta_id([1]); tx = T.thread_id([64])
+        f16  = T.alloc_local((1,), "float16")        # register scalars ...
+        bf16 = T.alloc_local((1,), "bfloat16")
+        i32  = T.alloc_local((1,), "int32")
+        u8   = T.alloc_local((1,), "uint8")
+        b1   = T.alloc_local((1,), "bool")
+        sm   = T.alloc_shared((64,), "float16")      # ... and a shared tile
+        v    = T.alloc_local((1,), "float32x4")      # a vector-dtype register (float4)
+        v[0] = A.vload([tx * 4], dtype="float32x4")  # vectorized load
+        O.vstore([tx * 4], v[0])                     # vectorized store
+        # ... (use f16/bf16/i32/u8/b1/sm) ...
+
+会 lower 成下面的代码（省略生成 CUDA 的样板部分）：
+
+.. code-block:: c++
+
+    half          f16_ptr[1];               // float16
+    nv_bfloat16   bf16_ptr[1];              // bfloat16
+    int           i32_ptr[1];               // int32
+    uchar         u8_ptr[1];                // uint8
+    signed char   b1_ptr[1];                // bool
+    __shared__ alignas(64) half sm_ptr[64]; // shared float16
+    float4        v_ptr[1];                 // float32x4  (vector)
+    v_ptr[0]                  = *(float4*)(A_ptr + tx * 4);   // vectorized load
+    *(float4*)(O_ptr + tx * 4) = v_ptr[0];                   // vectorized store
+
+buffer 的 dtype 本身也可以是 **vector type**：``T.alloc_local((1,), "float32x4")``
+会直接声明一个 ``float4`` register（用 ``v[0]`` 访问），而 ``float32x4``
+的 ``vload`` / ``vstore`` 会把它作为一次 16-byte 访问来搬运。vector dtype
+并不绑定在 ``vload`` 上；任意 buffer 或 scalar 都可以携带它。
+
+因此 dtype → CUDA 的映射如下：
+
+.. list-table::
+   :header-rows: 1
+   :widths: 34 33 33
+
+   * - dtype → CUDA
+     - dtype → CUDA
+     - dtype → CUDA
+   * - ``float32`` → ``float``
+     - ``float16`` → ``half``
+     - ``bfloat16`` → ``nv_bfloat16``
+   * - ``int32`` → ``int``
+     - ``uint8`` → ``uchar``
+     - ``bool`` → ``signed char``
+   * - ``float32x4`` → ``float4``
+     - ``handle`` → ``T*`` (pointer)
+     - (vector dtypes → CUDA vector types)
+
+dtype 与 type
+----------------
+
+``dtype`` 是 *低层* 信息，说明“这些 bit 如何解释”。此外，值还拥有高层
+**type**：标量是 ``PrimType(dtype)``，指针是
+``PointerType(PrimType(dtype), scope)``。大多数表达式都是标量
+（``PrimType``）；类型系统主要在 **指针** 上变得重要。
+
+指针（``handle``）
+------------------
+
+buffer 的 ``data``，也就是它的指针，是一个 pointer type 的 ``Var``，并且
+它是 **immutable** 的（指针不会被重新赋值）。这决定了你如何获得它：
+
+- ``T.alloc_buffer(...)`` 会分配存储空间，**并** 定义它的 ``data`` 指针。
+- ``T.decl_buffer(..., data=ptr)`` 会在已有指针 ``Var`` ``ptr`` 上声明一个
+  buffer。
+- 如果要让 buffer 背后使用一个指针 **表达式**，例如 ``T.ptx.map_shared_rank``
+  （PTX ``mapa``）返回另一个 cluster CTA 的 shared address，你必须先用
+  ``PointerType`` 的 ``T.let`` 把该表达式绑定成一个指针 ``Var``（``data``
+  必须是 ``Var``，不能是表达式）：
+
+  .. code-block:: python
+
+      from tvm.ir.type import PointerType, PrimType
+
+      ptr: T.let[T.Var(name="ptr", dtype=PointerType(PrimType("uint64")))] = \
+          T.reinterpret("handle", T.ptx.map_shared_rank(mbar.ptr_to([0]), 0))
+      remote_mbar = T.decl_buffer([1], "uint64", data=ptr, scope="shared")
diff --git a/zh/tirx_guide/language_reference/cuda/parser_utils.rst b/zh/tirx_guide/language_reference/cuda/parser_utils.rst
new file mode 100644
index 00000000..cf2ba190
--- /dev/null
+++ b/zh/tirx_guide/language_reference/cuda/parser_utils.rst
@@ -0,0 +1,79 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+Parser 工具
+===========
+
+有几个 helper 会在 **parse time** 生效（也就是 TVMScript 转成 TIRx 时），
+让你内联 Python 计算出的值、抽出可复用片段，并打包 parser 侧状态。
+
+``T.meta_var`` — 内联 Python 值
+----------------------------------------
+
+``T.meta_var(x)`` 会告诉 parser：把 ``x`` 这个由 **Python** 计算出的值当作
+compile-time *meta* 值，直接内联进 IR，而不是把它解析成 script 变量。这样
+可以避免一个临时 local，也能驱动 metaprogramming：围绕 meta value 的普通
+Python ``for`` 会在 parser 中展开。
+
+.. code-block:: python
+
+    n = T.meta_var(4)              # n is a Python int, inlined
+    for j in range(n):            # unrolled at parse time
+        acc[0] = acc[0] + A[tx, j]
+
+``@T.inline`` — 内联函数
+--------------------------------
+
+``@T.inline`` 定义的函数会在 parsing 期间把函数体 **内联到每个调用点**，
+生成代码中不会出现调用。它遵循 Python 的 lexical（LEGB）scope 和 late
+binding，因此参数会遮蔽外层变量：
+
+.. code-block:: python
+
+    @T.inline
+    def add_into(acc, x):
+        acc[0] = acc[0] + x
+
+    add_into(acc, A[tx, j])       # inlined -> acc[0] = acc[0] + A[tx, j]
+
+``@T.meta_class`` — parser 侧状态对象
+---------------------------------------------
+
+``@T.meta_class`` 标记一个普通 Python class，使其 **实例成为 parser meta
+value**：字段可以持有 buffer 和 scalar，因此你可以把相关 allocation 与状态
+打包进一个对象，并在 kernel body 中使用它。
+
+.. code-block:: python
+
+    @T.meta_class
+    class State:
+        def __init__(self, smem):
+            self.acc = T.alloc_local([1], "float32")
+            self.buf = T.decl_buffer([64], "float16", smem, scope="shared.dyn")
+
+    s = State(smem.data)
+    s.acc[0] = T.float32(0.0)     # use its fields like ordinary buffers
+    # ... s.buf[i] ...
+
+这很适合把 kernel 的 pipeline state（barrier、accumulator、scratch view）
+组织到一起，而不是让许多分散的 local 穿过整个 body。
+
+``T.constexpr``
+---------------
+
+``T.constexpr`` 标记 compile-time kernel 参数，它会通过 ``@T.jit`` 的
+``.specialize(...)`` 固化进去。细节见 :ref:`zh_chap_tirx_primer`。
diff --git a/zh/tirx_guide/language_reference/cuda/threads_sync.rst b/zh/tirx_guide/language_reference/cuda/threads_sync.rst
new file mode 100644
index 00000000..69822c1c
--- /dev/null
+++ b/zh/tirx_guide/language_reference/cuda/threads_sync.rst
@@ -0,0 +1,139 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+CUDA C++/PTX intrinsic
+======================
+
+当没有现成 tile primitive 覆盖你的需求时，有两条 escape hatch 可以直接触达
+硬件：**调用 backend intrinsic**（来自 ``tvm.backend.cuda`` 的 ``T.cuda.*`` /
+``T.ptx.*`` 命名空间），或者 **内联原始 CUDA** 源码。
+
+调用 backend intrinsic
+----------------------
+
+``T.cuda.*`` 和 ``T.ptx.*`` 会直接暴露 CUDA backend 的 device intrinsic：
+同步、mbarrier、reduction，以及 PTX data-movement / MMA 家族：
+
+.. code-block:: python
+
+    T.cuda.cta_sync()                    # block barrier (__syncthreads)
+    T.cuda.warp_sync()                   # __syncwarp
+    T.cuda.warpgroup_sync(8)             # warpgroup barrier
+    T.cuda.cta_sum(val, num_warps, scratch.ptr_to([0]))   # block-level reduction
+
+    bar = T.alloc_shared((1,), "uint64")
+    T.ptx.mbarrier.init(bar.data, 1)     # mbarrier for async completion
+    T.ptx.mbarrier.try_wait(bar.data, phase)
+
+下面是一个完整可运行的例子：通过 ``T.tvm_warp_shuffle_xor`` 做 warp all-reduce：
+
+.. code-block:: python
+
+    @T.prim_func
+    def warp_reduce(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (32,), "float32", align=16)
+        T.device_entry()
+        cta_id = T.cta_id([1]); warp_id = T.warp_id([1]); lane_id = T.lane_id([32])
+        v = T.alloc_local((1,), "float32"); i = T.alloc_local((1,), "int32")
+        v[0] = T.float32(31 - lane_id)
+        i[0] = 16
+        while i[0] >= 1:
+            v[0] += T.tvm_warp_shuffle_xor(0xFFFFFFFF, v[0], i[0], 32, 32)
+            i[0] = i[0] // 2
+        A[lane_id] = v[0]
+
+shuffle 会直接 lower 成 ``__shfl_xor_sync``：
+
+.. code-block:: c++
+
+    v_ptr[0] = v_ptr[0] + __shfl_xor_sync(0xFFFFFFFF, v_ptr[0], i_ptr[0], 32);
+
+``T.ptx.*`` / ``T.cuda.*`` 下还有其他家族：``cp_async``（LDGSTS）、
+``cp_async.bulk.tensor``（TMA）、``ldmatrix`` / ``stmatrix``、``tcgen05.*``
+（Blackwell MMA）、``atomic_add``、``fence`` 等。完整 ``tvm.backend.cuda``
+参考请见 backend API reference。
+
+同步语义
+--------
+
+GEMM 和 Flash Attention kernel 中反复出现四种同步机制。它们控制异步引擎和
+并行线程组，所以任何一种用错，通常都会导致静默数据损坏或 deadlock。
+
+**Mbarrier phase。** Mbarrier 使用一个内部 phase bit 跟踪 arrival。
+``T.ptx.mbarrier.try_wait(bar, phase)`` intrinsic 会阻塞，直到 barrier 的内部
+phase 与调用者提供的 ``phase`` 参数 *不同*。因此，当跨 loop iteration 复用
+barrier 时，调用者必须在每次 wait 之后翻转自己的本地 phase tracker
+（``phase ^= 1``）。如果忘了翻转，后续 wait 会立即返回，导致引擎读取半写入
+的 memory。:ref:`zh_chap_gemm_basics` 中完整走了一遍 phase-tracking 表。
+
+**Election。** ``T.ptx.elect_sync()`` 会在一个 warp 内选出 *单个 active lane*，
+不是 lane 0，也不是每个 CTA 一个线程。若要把 issuer 缩小到精确一个线程，
+必须配合 warp-level guard。:ref:`zh_chap_gemm_basics` 中使用
+``if warp_id == 0:`` 后接 ``if T.ptx.elect_sync():`` 的模式来发出
+``Tx.gemm_async`` 和 ``tcgen05.commit``。
+
+**Named Warpgroup Barrier。** ``T.cuda.cta_sync()`` 会映射到 ``__syncthreads()``，
+并要求 *每个* CTA 线程都到达。一旦 warpgroup 被 specialize 到不同代码路径，
+把 ``cta_sync()`` 放进某个 warpgroup 分支就会让 kernel deadlock，因为其他
+warpgroup 永远到不了它。硬件提供 16 个 named barrier（ID 0 到 15）；
+``T.cuda.warpgroup_sync(10)`` 只同步一个 warpgroup 的线程。不同 warpgroup
+使用不同 ID（例如 ``warpgroup_sync(wg_id + 10)``），避免撞到同一个硬件
+barrier。见 :ref:`zh_chap_gemm_advanced`。
+
+**Fence。** Fence 保证 producer 的写入排在 consumer（通常是异步引擎）读取
+之前：
+
+.. list-table::
+   :header-rows: 1
+   :widths: 50 50
+
+   * - Fence
+     - 保证的顺序
+   * - ``T.ptx.fence.proxy_async("shared::cta")``
+     - 线程写入 shared memory，先于 async proxy（TMA store / MMA）读取它
+   * - ``T.ptx.fence.mbarrier_init()``
+     - mbarrier 初始化，先于后续 arrival 或 wait 使用该 barrier
+   * - ``T.ptx.tcgen05.fence.after_thread_sync()``
+     - ``tcgen05`` writeback 边上的保守 ordering fence（Steps 8 和 9 会加入；TMA-to-MMA 路径不需要）
+
+内联原始 CUDA
+-------------
+
+如果某个功能完全没有 intrinsic，可以用
+``T.cuda.func_call(name, *args, source_code=..., return_type=...)`` 从源码字符串
+注入一个 ``__device__`` 函数：
+
+.. code-block:: python
+
+    SRC = r"""
+    __device__ __forceinline__ float my_relu(float x) { return x > 0.f ? x : 0.f; }
+    """
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (256,), "float32")
+        B = T.match_buffer(B_ptr, (256,), "float32")
+        T.device_entry(); bx = T.cta_id([1]); tx = T.thread_id([256])
+        B[tx] = T.cuda.func_call("my_relu", A[tx], source_code=SRC, return_type="float32")
+
+源码会原样发出，并把调用接进去：
+
+.. code-block:: c++
+
+    __device__ __forceinline__ float my_relu(float x) { return x > 0.f ? x : 0.f; }
+    // ...
+    B_ptr[tx] = my_relu(A_ptr[tx]);
diff --git a/zh/tirx_guide/language_reference/index.rst b/zh/tirx_guide/language_reference/index.rst
new file mode 100644
index 00000000..103e357a
--- /dev/null
+++ b/zh/tirx_guide/language_reference/index.rst
@@ -0,0 +1,35 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+.. _zh_chap_language_reference:
+
+TIRx 语言参考
+==============
+
+这里收录编写 TIRx device kernel 所需的完整语言特性，并从
+:ref:`zh_chap_tirx_primer` walkthrough 中拆分出来：parser 工具、数据类型与表达式、
+buffer 与内存、控制流，以及线程同步。当你需要确认某个特性的精确写法或语义时，
+可以查阅这些页面。
+
+.. toctree::
+   :maxdepth: 1
+
+   cuda/parser_utils
+   cuda/data_types
+   cuda/buffers
+   cuda/control_flow
+   cuda/threads_sync