npm - mlx-bun - Versions diffs - 0.0.5 → 0.0.6 - Mend

mlx-bun 0.0.5 → 0.0.6

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (57) hide show

package/README.md +11 -2
package/docs/dag/training-inference-map.html +575 -0
package/docs/design/memory-system.md +193 -295
package/docs/design/orpo-dynamic-lambda.md +179 -0
package/docs/design/orpo-future-enhancements.md +42 -0
package/docs/design/orpo-training.md +13 -0
package/docs/design/parity-tier-dag.md +143 -0
package/docs/design/w4a16-compute-precision-spike.md +232 -0
package/docs/investigations/orpo-base-uf-experiment-and-directions.md +120 -0
package/docs/investigations/orpo-flash-cce-pin-leak.md +167 -0
package/docs/investigations/orpo-uf-testing-handoff.md +126 -0
package/docs/investigations/steel-flash-cce-handoff.md +58 -0
package/docs/investigations/trainer-validation-experiment.md +173 -0
package/docs/reference/memory.md +77 -0
package/docs/reference/orpo-quickstart.md +44 -5
package/docs/reference/training.md +9 -5
package/package.json +1 -1
package/src/cli.ts +651 -18
package/src/eval/runner.ts +76 -6
package/src/eval/tasks/bfcl.ts +49 -15
package/src/eval/tasks/gsm8k.ts +21 -9
package/src/eval/tasks/hashhop.ts +33 -4
package/src/eval/tasks/humaneval.ts +30 -8
package/src/eval/tasks/ifeval.ts +52 -18
package/src/eval/tasks/mmlu.ts +85 -74
package/src/jobs/runner.ts +6 -0
package/src/memory/chunk.ts +25 -0
package/src/memory/cluster.ts +24 -0
package/src/memory/db.ts +114 -0
package/src/memory/ingest.ts +22 -0
package/src/memory/pipeline.ts +63 -0
package/src/memory/prompts.ts +23 -0
package/src/memory/schedule.ts +163 -0
package/src/memory/synthesize.ts +29 -0
package/src/memory/tools.ts +399 -0
package/src/memory/vault.ts +694 -0
package/src/memory/wikify.ts +23 -0
package/src/mlx/metal-kernel.ts +13 -0
package/src/model/flash-attention.ts +5 -3
package/src/pi-provider.ts +29 -13
package/src/pi-session.ts +81 -0
package/src/pi-terminal.ts +13 -24
package/src/pi-web.ts +246 -166
package/src/quantize/job.ts +17 -2
package/src/quantize/quantizer.ts +26 -2
package/src/server.ts +213 -44
package/src/tokenizer.ts +6 -2
package/src/train/dataset.ts +11 -3
package/src/train/flash-cce.ts +18 -5
package/src/train/job.ts +41 -4
package/src/train/loss.ts +10 -4
package/src/train/trainer.ts +62 -3
package/src/train/watch.ts +542 -0
package/src/web/app.html +264 -14
package/src/web/skills/memory/SKILL.md +97 -0
package/src/web/skills/web-research/SKILL.md +3 -1
package/src/web/skills.ts +46 -8

package/README.md CHANGED Viewed

@@ -134,6 +134,14 @@ back in responses).
 - **The background agent runtime** — leave a local OpenAI-compatible endpoint
   available for scripts, recurring jobs, app integrations, and per-request
   adapter routing.
+- **The local memory layer** — `mlx-bun memory init` creates a git-tracked
+  Markdown wiki at `~/.mlx-bun/wiki`; `Reference/` is seeded with read-only
+  symlinks to mlx-bun's own docs, and the built-in pi agents read both those
+  docs and your articles quietly (`memory_search`, `memory_section`, etc.) and can open it in Obsidian for you
+  (`mlx-bun memory open [article]`). Chat-time memory is read-only; synthesis is
+  the separate writer path. The planned synthesis pipeline uses the same local
+  models with stage-specific LoRA adapters (chunk, bucket, synthesize, edit)
+  trained by the project creator and distributed from Hugging Face.
 The default experience is intentionally opinionated. Power users can override
 models, budgets, KV modes, adapters, and flags; the average user should not
@@ -407,8 +415,9 @@ prompt cache); registry / fit / KV persistence; quantized +
 mixed-precision KV serving (rotating-cache KV-quant, Phase 9) with fused
 quantized prefill (Phase 10); LoRA hot-swap with per-request selection;
 segmented-backward LoRA training (SFT / DPO / ORPO, with an `[M,vocab]`-free
-flash-CCE head + prefix-sharing for long-context preference data — see
-[training guide](docs/reference/training.md)); resumable
+flash-CCE head + prefix-sharing for long-context preference data — `mlx-bun
+train <model> --data <dir>`, see [training guide](docs/reference/training.md));
+resumable
 verified downloads (`mlx-bun get`); memory admission control; the
 head-to-head benchmark harness; the decode-gap root-cause fix
 (2026-06-11); Anthropic Messages + Responses API (`/v1/messages` and

package/docs/dag/training-inference-map.html ADDED Viewed

@@ -0,0 +1,575 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8" />
+<meta name="viewport" content="width=device-width, initial-scale=1" />
+<title>mlx-bun · control-flow DAG (MiniCPM5-1B)</title>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/cytoscape/3.28.1/cytoscape.min.js"></script>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/dagre/0.8.5/dagre.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/cytoscape-dagre@2.5.0/cytoscape-dagre.min.js"></script>
+<style>
+  :root{--bg:#0b0f14;--panel:#11161d;--line:#243040;--txt:#cdd9e5;--dim:#7d8aa0;--cond:#39c5cf;}
+  *{box-sizing:border-box;}
+  html,body{margin:0;height:100%;background:var(--bg);color:var(--txt);font:13px/1.5 ui-monospace,SFMono-Regular,Menlo,monospace;}
+  #top{display:flex;align-items:center;gap:12px;padding:8px 14px;flex-wrap:wrap;border-bottom:1px solid var(--line);background:var(--panel);}
+  #top b{color:#fff;font-size:14px;}
+  .seg{display:flex;gap:2px;background:#0b0f14;border:1px solid var(--line);border-radius:6px;}
+  .seg button,.btn{background:transparent;color:var(--dim);border:0;padding:5px 11px;cursor:pointer;font:inherit;border-radius:5px;}
+  .seg button.on{background:#388bfd;color:#fff;}
+  .btn{border:1px solid var(--line);}
+  .legend{margin-left:auto;display:flex;gap:11px;color:var(--dim);align-items:center;flex-wrap:wrap;font-size:11px;}
+  .legend span{display:flex;align-items:center;gap:5px;}
+  .dot{width:9px;height:9px;border-radius:50%;display:inline-block;}
+  .dash{width:15px;border-top:2px dashed var(--cond);display:inline-block;}
+  #wrap{display:flex;height:calc(100% - 46px);}
+  #cy{flex:1;height:100%;}
+  #side{width:340px;border-left:1px solid var(--line);background:var(--panel);padding:14px;overflow:auto;}
+  #side h2{margin:0 0 4px;font-size:14px;color:#fff;}
+  #side .file{color:#388bfd;font-size:12px;word-break:break-all;margin-bottom:8px;}
+  #side .k{color:var(--dim);text-transform:uppercase;font-size:10.5px;letter-spacing:.06em;margin-top:10px;}
+  #side .v{white-space:pre-wrap;}
+  .hint{color:var(--dim);font-size:12px;margin-top:6px;}
+  code{color:#e3b341;}
+</style>
+</head>
+<body>
+<div id="top">
+  <b>mlx-bun control-flow DAG</b><span style="color:var(--dim);font-size:11px">MiniCPM5-1B</span>
+  <div class="seg" id="graph">
+    <button data-g="infer" class="on">inference</button>
+    <button data-g="train">training</button>
+  </div>
+  <button class="btn" id="fit">⤢ fit</button>
+  <div class="legend" id="legend"></div>
+</div>
+<div id="wrap">
+  <div id="cy"></div>
+  <div id="side">
+    <h2 id="s-title">Click a step</h2>
+    <div class="file" id="s-file"></div>
+    <div class="hint" id="hint">Every node = one actual step the code runs. Solid arrows = the next step in execution order. <b style="color:var(--cond)">Dashed labelled arrows</b> = a branch: the code can go several ways, and the label is the exact condition (flag / variable / check) that takes that edge. Loops are back-edges. Scroll to zoom, drag to pan. Colour = subsystem (legend). Traced from the actual source — file:line on each node.</div>
+    <div id="s-body" style="display:none">
+      <div class="k">file</div><div class="v" id="s-fileb"></div>
+      <div class="k">subsystem</div><div class="v" id="s-grp"></div>
+      <div class="k">outgoing</div><div class="v" id="s-out"></div>
+    </div>
+  </div>
+</div>
+<script>
+// node = [id, label, group, file].  edge = [from, to, condition('' = unconditional)].
+// A node with several edges of different conditions IS a branch point.
+const GRP={ // subsystem → colour
+  drive:'#388bfd', fwd:'#2ea043', layer:'#3fb950', attn:'#d29922', sdpa:'#a371f7',
+  unfused:'#6e7b91', tiled:'#39c5cf', mlp:'#56d364', kv:'#f0883e', samp:'#db61a2',
+  setup:'#6e7b91', loss:'#d29922', seg:'#a371f7', accum:'#388bfd', bwd:'#f85149',
+  norm:'#39c5cf', opt:'#f0883e', metrics:'#db61a2' };
+const INFER={ n:[
+ ['n1','generate() entry','drive','generate.ts:168'],
+ ['n2','generateInner() entry','drive','generate.ts:208'],
+ ['n3','? wired-scope: FORCE_WIRE ∨ weights>thresh','drive','generate.ts:175'],
+ ['n4','enter wired scope','drive','generate.ts:200'],
+ ['n5','? adapters?.length','drive','generate.ts:174'],
+ ['n6','set loraState.active','drive','generate.ts:187'],
+ ['n7','makeCache/Sampler/LogitsProcessors','drive','generate.ts:218'],
+ ['n8','maybeQuantizeKv (prefill)','drive','generate.ts:281'],
+ ['n9','? prompt-pos > prefillChunkSize','drive','generate.ts:298'],
+ ['n10','chunked prefill: loop chunk','drive','generate.ts:299'],
+ ['n11','forwardHidden(chunk, cache)','drive','generate.ts:301'],
+ ['n12','evalAll(cache) after chunk','drive','generate.ts:304'],
+ ['n13','clearCache() after chunk','drive','generate.ts:312'],
+ ['n14','? options.promptEmbeddings','drive','generate.ts:284'],
+ ['n15','forwardEmbeddings(...)','drive','generate.ts:292'],
+ ['n16','forwardHidden(lastChunk, cache)','drive','generate.ts:320'],
+ ['n17','? needsTokenHistory init','drive','generate.ts:315'],
+ ['n18','slice last position h0[:,L-1]','drive','generate.ts:324'],
+ ['n19','logitsFromHidden(hLast)','drive','generate.ts:326'],
+ ['n20','sampleStep(logits0, 0)','drive','generate.ts:328'],
+ ['n21','asyncEvalAll([pending])','drive','generate.ts:338'],
+ ['n22','? CompiledDecode.for (env·adapters·gemma4)','drive','generate.ts:352'],
+ ['n23','decode loop: while !stop','drive','generate.ts:360'],
+ ['n24','? generated+1 < maxTokens','drive','generate.ts:364'],
+ ['n25','maybeQuantizeKv (decode)','drive','generate.ts:365'],
+ ['n26','pushHistory(cur)','drive','generate.ts:366'],
+ ['n27','? compiled && supports(cache)','drive','generate.ts:369'],
+ ['n28','compiled.step(cur, cache)','drive','generate.ts:371'],
+ ['n29','catch → compiled=null (fallback)','drive','generate.ts:374'],
+ ['n30','forwardHidden(ids, cache)','drive','generate.ts:383'],
+ ['n31','logitsFromHidden(h)','drive','generate.ts:385'],
+ ['n32','sampleStep(logits, n+1)','drive','generate.ts:388'],
+ ['n33','asyncEvalAll([nextPending,…])','drive','generate.ts:390'],
+ ['n34','itemUint32(cur) sync-read token','drive','generate.ts:394'],
+ ['n35','? first token: prefillMs/tDecode','drive','generate.ts:399'],
+ ['n36','cur.dispose; generated++','drive','generate.ts:402'],
+ ['n37','? nextPending !== null','drive','generate.ts:406'],
+ ['n38','? eosTokenIds.includes(token)','drive','generate.ts:408'],
+ ['n39','yield {token, index}','drive','generate.ts:413'],
+ ['n40','? clearCache every 256','drive','generate.ts:416'],
+ ['n42','stop = true (no next)','drive','generate.ts:418'],
+ ['n43','pending = nextPending','drive','generate.ts:420'],
+ ['n44','stop = true (EOS / done)','drive','generate.ts:411'],
+ ['n45','decodeMs','drive','generate.ts:425'],
+ ['n46','return makeStats()','drive','generate.ts:427'],
+ ['n47','finally: dispose pending/cache/history','drive','generate.ts:432'],
+ ['n49','loraState.active=[] (scope exit)','drive','generate.ts:191'],
+ ['n50','exitWiredScope()','drive','generate.ts:204'],
+ // model forward
+ ['n51','embed.encode(ids)','fwd','minicpm5.ts:250'],
+ ['n52','forwardLayers(h, cache)','fwd','minicpm5.ts:258'],
+ ['n53','runLayerRange(h,0,N,cache)','fwd','minicpm5.ts:259'],
+ ['n54','? makeMask(L) (N==1? offset==0?)','fwd','minicpm5.ts:275'],
+ ['n55','for i in [0,N): layers','fwd','minicpm5.ts:277'],
+ ['n56','layers[i].forward(cur,mask,cache[i])','fwd','minicpm5.ts:278'],
+ ['n57','? dispose(cur) if i>0','fwd','minicpm5.ts:279'],
+ ['n58','mask.arr?.dispose()','fwd','minicpm5.ts:282'],
+ ['n59','finalNorm.forward(h)','fwd','minicpm5.ts:261'],
+ ['n60','lmHead.forward(h) → logits','fwd','minicpm5.ts:287'],
+ // layer
+ ['n61','inputNorm RMSNorm','layer','minicpm5.ts:189'],
+ ['n62','attn.forward(xn,mask,cache)','layer','minicpm5.ts:190'],
+ ['n63','add(x, attnOut) residual','layer','minicpm5.ts:192'],
+ ['n64','postAttnNorm RMSNorm','layer','minicpm5.ts:194'],
+ ['n65','mlp.forward(hn)','layer','minicpm5.ts:195'],
+ ['n66','add(h, mlpOut) residual','layer','minicpm5.ts:197'],
+ // attention
+ ['n67','q_proj quantized_matmul','attn','minicpm5.ts:89'],
+ ['n68','k_proj quantized_matmul','attn','minicpm5.ts:90'],
+ ['n69','v_proj quantized_matmul','attn','minicpm5.ts:91'],
+ ['n70','reshape q [B,L,H,D]','attn','minicpm5.ts:93'],
+ ['n71','reshape k [B,L,Kv,D]','attn','minicpm5.ts:94'],
+ ['n72','reshape v [B,L,Kv,D]','attn','minicpm5.ts:95'],
+ ['n73','transpose q →[B,H,L,D]','attn','minicpm5.ts:97'],
+ ['n74','transpose k','attn','minicpm5.ts:98'],
+ ['n75','transpose v','attn','minicpm5.ts:99'],
+ ['n76','? rope: prefixPlan / offsetArr / static','attn','minicpm5.ts:106'],
+ ['n77','ropeBlocks(q) (prefix)','attn','minicpm5.ts:108'],
+ ['n78','ropeDynamic(q) (offsetArr)','attn','minicpm5.ts:110'],
+ ['n79','rope(q, cache.offset)','attn','minicpm5.ts:111'],
+ ['n80','rope k (same offset)','attn','minicpm5.ts:113'],
+ ['n81','? cache instanceof QuantizedKVCache','attn','minicpm5.ts:115'],
+ ['n82','updateAndFetchQuantized(k,v)','attn','minicpm5.ts:116'],
+ ['n83','quantizedSdpa(q,kq,vq,…)','attn','minicpm5.ts:119'],
+ ['n84','updateAndFetch(k,v) bf16','attn','minicpm5.ts:123'],
+ ['n85','? getTrainingAttn()=="flash"','attn','minicpm5.ts:128'],
+ ['n86','flashAttention(...)','attn','minicpm5.ts:130'],
+ ['n87','ops.sdpa(q,keys,values,…)','attn','minicpm5.ts:132'],
+ ['n88','transpose attn →[B,L,H,D]','attn','minicpm5.ts:138'],
+ ['n89','reshape merged [B,L,H*D]','attn','minicpm5.ts:140'],
+ ['n90','o_proj quantized_matmul','attn','minicpm5.ts:142'],
+ // sdpa dispatch
+ ['n91','? mask=="" && scale==1.0','sdpa','gemma4-base.ts:1514'],
+ ['n92','? perfKernelEnabled()','sdpa','fused-decode-kernel.ts:36'],
+ ['n93','? !isCompiledTrace()','sdpa','gemma4-base.ts:1515'],
+ ['n94','? fusedDecodeKernelSupported','sdpa','gemma4-base.ts:1515'],
+ ['n95','⚙ fusedDecodeSdpa (L=1)','sdpa','gemma4-base.ts:1517'],
+ ['n96','? tile: L>1 ∨ FUSED_DECODE','sdpa','gemma4-base.ts:1518'],
+ ['n97','? fusedSdpaSupported','sdpa','gemma4-base.ts:1519'],
+ ['n98','⚙ quantizedSdpaTiled','sdpa','gemma4-base.ts:1520'],
+ ['n99','quantizedSdpaUnfused','sdpa','gemma4-base.ts:1521'],
+ // unfused internals
+ ['n100','dims [B,H,L,D,KV,nRep,N]','unfused','gemma4-base.ts:1221'],
+ ['n101','? scale !== 1.0','unfused','gemma4-base.ts:1229'],
+ ['n102','mulScalar(q, scale)','unfused','gemma4-base.ts:1229'],
+ ['n103','? nRep > 1 (GQA)','unfused','gemma4-base.ts:1233'],
+ ['n104','reshape q [B,KV,nRep,L,D]','unfused','gemma4-base.ts:1234'],
+ ['n105','expandDims(k.packed,-3)','unfused','gemma4-base.ts:1245'],
+ ['n106','quantizedMatmulQT(q,kᵀ) → scores','unfused','gemma4-base.ts:1251'],
+ ['n107','? mask.mode=="causal"','unfused','gemma4-base.ts:1256'],
+ ['n108','greaterEqual → causal mask','unfused','gemma4-base.ts:1261'],
+ ['n109','? mask.mode=="array"','unfused','gemma4-base.ts:1264'],
+ ['n110','where(mask, scores, -inf)','unfused','gemma4-base.ts:1269'],
+ ['n111','softmaxAxis(scores,-1)','unfused','gemma4-base.ts:1276'],
+ ['n112','quantizedMatmulQT(probs,vᵀ) → out','unfused','gemma4-base.ts:1278'],
+ ['n113','? nRep > 1','unfused','gemma4-base.ts:1279'],
+ ['n114','reshape out [B,H,L,D]','unfused','gemma4-base.ts:1280'],
+ // tiled internals
+ ['n115','dims [B,H,L,D,KV,nRep,N]','tiled','gemma4-base.ts:1304'],
+ ['n116','? scale !== 1.0','tiled','gemma4-base.ts:1312'],
+ ['n117','mulScalar(q, scale)','tiled','gemma4-base.ts:1312'],
+ ['n118','? nRep > 1 (GQA)','tiled','gemma4-base.ts:1317'],
+ ['n119','reshape q [B,KV,nRep,L,D]','tiled','gemma4-base.ts:1318'],
+ ['n120','expandDims(k.packed,-3)','tiled','gemma4-base.ts:1324'],
+ ['n121','? mask.mode=="causal"','tiled','gemma4-base.ts:1340'],
+ ['n122','greaterEqual → causal mask','tiled','gemma4-base.ts:1345'],
+ ['n123','? mask.mode=="array"','tiled','gemma4-base.ts:1348'],
+ ['n124','for n0=0..N step 512 (KV chunks)','tiled','gemma4-base.ts:1378'],
+ ['n125','slice kᵀ.packed chunk','tiled','gemma4-base.ts:1381'],
+ ['n126','slice kᵀ.scales chunk','tiled','gemma4-base.ts:1382'],
+ ['n127','slice kᵀ.biases chunk','tiled','gemma4-base.ts:1383'],
+ ['n128','slice vᵀ.packed chunk','tiled','gemma4-base.ts:1386'],
+ ['n129','quantizedMatmulQT(q,kChunk) → scores','tiled','gemma4-base.ts:1391'],
+ ['n130','? maskChunk dtype bool','tiled','gemma4-base.ts:1397'],
+ ['n131','where(mask, scores, -inf)','tiled','gemma4-base.ts:1399'],
+ ['n132','add(scores, maskChunk) float','tiled','gemma4-base.ts:1402'],
+ ['n133','maxAxis(scores) → chunkMax','tiled','gemma4-base.ts:1409'],
+ ['n134','? oAcc===null (first tile)','tiled','gemma4-base.ts:1410'],
+ ['n135','init rowMax/rowSum/oAcc','tiled','gemma4-base.ts:1411'],
+ ['n136','maximum(rowMax,chunkMax)=newMax','tiled','gemma4-base.ts:1419'],
+ ['n137','exp(rowMax-newMax)=factor','tiled','gemma4-base.ts:1422'],
+ ['n138','exp(scores-newMax)=exps','tiled','gemma4-base.ts:1425'],
+ ['n139','newSum = factor·rowSum + Σexps','tiled','gemma4-base.ts:1428'],
+ ['n140','nextAcc = oAcc·factor + exps·vChunk','tiled','gemma4-base.ts:1436'],
+ ['n141','rowMax/rowSum = new','tiled','gemma4-base.ts:1443'],
+ ['n142','div(oAcc, rowSum) → out','tiled','gemma4-base.ts:1452'],
+ ['n143','? nRep > 1','tiled','gemma4-base.ts:1459'],
+ ['n144','reshape out [B,H,L,D]','tiled','gemma4-base.ts:1460'],
+ // mlp
+ ['n145','gate_proj quantized_matmul','mlp','minicpm5.ts:160'],
+ ['n146','up_proj quantized_matmul','mlp','minicpm5.ts:161'],
+ ['n147','sigmoid(gate)','mlp','minicpm5.ts:162'],
+ ['n148','mul → SiLU (gate·σ(gate))','mlp','minicpm5.ts:163'],
+ ['n149','mul(silu, up)','mlp','minicpm5.ts:166'],
+ ['n150','down_proj quantized_matmul','mlp','minicpm5.ts:169'],
+ // KV bf16
+ ['n151','offset=prev, L=k.shape[2]','kv','gemma4-base.ts:252'],
+ ['n152','? capacity: !keys ∨ prev+L>cap','kv','gemma4-base.ts:254'],
+ ['n153','alloc newK/newV zeros','kv','gemma4-base.ts:257'],
+ ['n154','? cache exists','kv','gemma4-base.ts:260'],
+ ['n155','? prev%STEP!=0 (trim)','kv','gemma4-base.ts:263'],
+ ['n156','slice old, dispose','kv','gemma4-base.ts:264'],
+ ['n157','concat([old,new]) grow','kv','gemma4-base.ts:271'],
+ ['n158','sliceUpdate keys←k','kv','gemma4-base.ts:283'],
+ ['n159','sliceUpdate values←v','kv','gemma4-base.ts:284'],
+ ['n160','slice fetch keys','kv','gemma4-base.ts:290'],
+ ['n161','slice fetch values','kv','gemma4-base.ts:291'],
+ // KV quantized
+ ['n162','prev/L/elPerInt','kv','gemma4-base.ts:443'],
+ ['n163','? capacity (packed)','kv','gemma4-base.ts:446'],
+ ['n164','grow cache (6 comps)','kv','gemma4-base.ts:447'],
+ ['n165','? prev%STEP!=0 (trim)','kv','gemma4-base.ts:448'],
+ ['n166','trimTo (slice 6)','kv','gemma4-base.ts:449'],
+ ['n167','quantize(k) → kq','kv','gemma4-base.ts:467'],
+ ['n168','quantize(v) → vq','kv','gemma4-base.ts:468'],
+ ['n169','write k packed/scales/biases','kv','gemma4-base.ts:479'],
+ ['n175','fetch keys (3 comps)','kv','gemma4-base.ts:488'],
+ ['n176','fetch values (3 comps)','kv','gemma4-base.ts:488'],
+ // sampler
+ ['n177','? curve sampler active','samp','sampler.ts:418'],
+ ['n178','applyCurve(lp, curve)','samp','sampler.ts:420'],
+ ['n179','randomCategorical (curve)','samp','sampler.ts:421'],
+ ['n180','? temperature === 0','samp','sampler.ts:429'],
+ ['n181','argmaxAxis(lp) greedy','samp','sampler.ts:430'],
+ ['n182','? hlg?.enabled','samp','sampler.ts:437'],
+ ['n183','applyHlgShaper','samp','sampler.ts:453'],
+ ['n184','applyHlgPipeline','samp','sampler.ts:454'],
+ ['n185','applyHlgEotf','samp','sampler.ts:454'],
+ ['n186','applyHlgOetf','samp','sampler.ts:455'],
+ ['n187','applyHlg (parametric)','samp','sampler.ts:455'],
+ ['n188','randomCategorical (HLG)','samp','sampler.ts:457'],
+ ['n189','? topP in (0,1)','samp','sampler.ts:467'],
+ ['n190','applyTopP(lp, topP)','samp','sampler.ts:467'],
+ ['n191','? topK > 0','samp','sampler.ts:468'],
+ ['n192','applyTopK(lp, topK)','samp','sampler.ts:468'],
+ ['n193','mulScalar(lp, 1/temp)','samp','sampler.ts:469'],
+ ['n194','randomCategorical (temp)','samp','sampler.ts:471'],
+],e:[
+ ['n1','n2',''],['n2','n3',''],
+ ['n3','n4','FORCE_WIRE ∨ weights>thresh'],['n3','n5','else'],
+ ['n4','n5',''],
+ ['n5','n6','adapters?.length'],['n5','n7','else'],
+ ['n6','n7',''],['n7','n8',''],['n8','n9',''],
+ ['n9','n10','pos > prefillChunkSize'],['n9','n14','else (one-shot)'],
+ ['n10','n11',''],['n11','n51','call forward'],['n11','n12',''],['n12','n13',''],
+ ['n13','n10','more chunks'],['n13','n14','done'],
+ ['n14','n15','promptEmbeddings'],['n14','n16','else'],
+ ['n15','n17',''],['n16','n51','call forward'],['n16','n17',''],
+ ['n17','n18',''],['n18','n19',''],['n19','n60','call lmHead'],['n19','n20',''],['n20','n21',''],['n21','n22',''],['n22','n23',''],
+ ['n23','n24',''],
+ ['n24','n25','generated+1<maxTokens'],['n24','n34','else'],
+ ['n25','n26',''],['n26','n27',''],
+ ['n27','n28','compiled && supports'],['n27','n30','else (uncompiled)'],
+ ['n28','n29','throws'],['n28','n31','ok'],['n29','n30',''],
+ ['n30','n51','call forward'],['n30','n31',''],['n31','n60','call lmHead'],['n31','n32',''],['n32','n33',''],['n33','n34',''],
+ ['n34','n35','first token'],['n34','n36','else'],['n35','n36',''],['n36','n37',''],
+ ['n37','n38','nextPending!=null'],['n37','n42','else'],
+ ['n38','n39','!EOS'],['n38','n44','EOS'],
+ ['n39','n40',''],['n40','n43','nextPending'],['n40','n42','none'],
+ ['n42','n44',''],['n43','n23','loop'],
+ ['n44','n45',''],['n45','n46',''],['n46','n47',''],['n47','n49',''],['n49','n50',''],
+ // forward
+ ['n51','n52',''],['n52','n53',''],['n53','n54',''],['n54','n55',''],
+ ['n55','n56','i<N'],['n56','n57',''],['n57','n55','i++'],['n57','n58','done'],
+ ['n58','n59',''],['n59','n12','return (prefill)'],['n59','n17','return (prefill)'],['n59','n31','return (decode)'],
+ // layer (called by n56)
+ ['n56','n61','call layer'],
+ ['n61','n62',''],['n62','n63',''],['n63','n64',''],['n64','n65',''],['n65','n66',''],
+ // attention (called by n62)
+ ['n62','n67','call attn'],
+ ['n67','n68',''],['n68','n69',''],['n69','n70',''],['n70','n71',''],['n71','n72',''],['n72','n73',''],['n73','n74',''],['n74','n75',''],['n75','n76',''],
+ ['n76','n77','_prefixPlan'],['n76','n78','offsetArr'],['n76','n79','else (static)'],
+ ['n77','n80',''],['n78','n80',''],['n79','n80',''],['n80','n81',''],
+ ['n81','n82','QuantizedKVCache'],['n81','n84','else (bf16)'],
+ ['n82','n83',''],['n83','n88',''],
+ ['n84','n85',''],
+ ['n85','n86','flash & supported'],['n85','n87','else'],
+ ['n86','n88',''],['n87','n88',''],['n88','n89',''],['n89','n90',''],
+ // sdpa dispatch (called by n83)
+ ['n83','n91','call quantizedSdpa'],
+ ['n91','n92','mask=="" & scale==1'],['n91','n96','else'],
+ ['n92','n93','perfKernel'],['n92','n96','else'],
+ ['n93','n94','!compiledTrace'],['n93','n96','else'],
+ ['n94','n95','supported'],['n94','n96','else'],
+ ['n96','n97','tile (L>1 ∨ FUSED_DECODE)'],['n96','n99','else'],
+ ['n97','n98','fusedSdpaSupported'],['n97','n99','else'],
+ // unfused (n99)
+ ['n99','n100','call'],['n100','n101',''],
+ ['n101','n102','scale!=1'],['n101','n103','else'],['n102','n103',''],
+ ['n103','n104','nRep>1'],['n103','n106','else'],['n104','n105',''],['n105','n106',''],
+ ['n106','n107',''],['n107','n108','causal'],['n107','n109','else'],
+ ['n108','n110',''],['n109','n110','array'],['n109','n111','none'],['n110','n111',''],
+ ['n111','n112',''],['n112','n113',''],['n113','n114','nRep>1'],
+ // tiled (n98)
+ ['n98','n115','call'],['n115','n116',''],
+ ['n116','n117','scale!=1'],['n116','n118','else'],['n117','n118',''],
+ ['n118','n119','nRep>1'],['n118','n121','else'],['n119','n120',''],['n120','n121',''],
+ ['n121','n122','causal'],['n121','n123','else'],['n122','n124',''],['n123','n124',''],
+ ['n124','n125','n0<N'],['n125','n126',''],['n126','n127',''],['n127','n128',''],['n128','n129',''],
+ ['n129','n130',''],['n130','n131','bool'],['n130','n132','float'],['n131','n133',''],['n132','n133',''],
+ ['n133','n134',''],['n134','n135','first'],['n134','n136','else'],
+ ['n135','n124','next chunk'],
+ ['n136','n137',''],['n137','n138',''],['n138','n139',''],['n139','n140',''],['n140','n141',''],['n141','n124','next chunk'],
+ ['n124','n142','n0>=N'],['n142','n143',''],['n143','n144','nRep>1'],
+ // mlp (called by n65)
+ ['n65','n145','call mlp'],
+ ['n145','n146',''],['n146','n147',''],['n147','n148',''],['n148','n149',''],['n149','n150',''],
+ // KV bf16 (called by n84)
+ ['n84','n151','call updateAndFetch'],
+ ['n151','n152',''],['n152','n153','grow'],['n152','n158','fits'],
+ ['n153','n154',''],['n154','n155','exists'],['n154','n157','first'],
+ ['n155','n156','prev%STEP!=0'],['n155','n157','aligned'],['n156','n157',''],['n157','n158',''],
+ ['n158','n159',''],['n159','n160',''],['n160','n161',''],
+ // KV quant (called by n82)
+ ['n82','n162','call updateAndFetchQuantized'],
+ ['n162','n163',''],['n163','n164','grow'],['n163','n167','fits'],
+ ['n164','n165',''],['n165','n166','prev%STEP!=0'],['n165','n167','aligned'],['n166','n167',''],
+ ['n167','n168',''],['n168','n169',''],['n169','n175',''],['n175','n176',''],
+ // sampler (called by n20, n32)
+ ['n20','n177','call sampleStep'],['n32','n177','call sampleStep'],
+ ['n177','n178','curve'],['n177','n180','else'],['n178','n179',''],
+ ['n180','n181','temp==0'],['n180','n182','else'],
+ ['n182','n183','shaper'],['n182','n184','pipeline'],['n182','n185','eotf'],['n182','n186','oetf'],['n182','n187','parametric'],['n182','n189','hlg off'],
+ ['n183','n188',''],['n184','n188',''],['n185','n188',''],['n186','n188',''],['n187','n188',''],
+ ['n189','n190','topP'],['n189','n191','else'],['n190','n191',''],
+ ['n191','n192','topK'],['n191','n193','else'],['n192','n193',''],['n193','n194',''],
+]};
+const TRAIN={ n:[
+ ['t1','step loop begins','setup','trainer.ts:479'],
+ ['t2','per-step dropout seed','setup','trainer.ts:876'],
+ ['t3','enter accumulation loop','accum','trainer.ts:487'],
+ ['t4','fetch micro-batch','accum','trainer.ts:490'],
+ ['t5','? segmentSize > 0','setup','trainer.ts:409'],
+ ['t6','segmented forward start','seg','segmented.ts:196'],
+ ['t7','embed.encode(ids)','seg','segmented.ts:196'],
+ ['t8','detach embedding boundary','seg','segmented.ts:196'],
+ ['t9','for segment k=0..nSeg','seg','segmented.ts:197'],
+ ['t10','runLayerRange[lo:hi] (swap LoRA)','seg','segmented.ts:199'],
+ ['t11','detach output boundary k','seg','segmented.ts:199'],
+ ['t12','? ORPO prefix-shared & segmented','seg','trainer.ts:880'],
+ ['t13','splitPrefixBatch','seg','trainer.ts:881'],
+ ['t14','single forward [P;C;R] block-mask','seg','prefix-shared.ts:159'],
+ ['t15','prefix fallback (two-forward)','seg','trainer.ts:884'],
+ ['t16','non-segmented forward start','accum','trainer.ts:492'],
+ ['t17','value_and_grad: swap primals','accum','trainer.ts:432'],
+ ['t18','? B==1 response-only vs full','loss','loss.ts:123'],
+ ['t19','forwardHidden → [B,T,hid]','fwd','forward.ts:72'],
+ ['t20','embedding forward','fwd','forward.ts:92'],
+ ['t21','layer-stack forward','fwd','forward.ts:92'],
+ ['t22','? MLX_BUN_TRAIN_ATTN=="flash"','fwd','trainer.ts:469'],
+ ['t23','sdpa attention (O(L²) bwd)','fwd','loss.ts:131'],
+ ['t24','flash attention (O(L) bwd)','fwd','trainer.ts:469'],
+ ['t25','final norm','fwd','forward.ts:113'],
+ ['t26','lm-head → logits','fwd','loss.ts:131'],
+ ['t27','? B==1 responseCe vs B>1 maskedCe','loss','loss.ts:123'],
+ ['t28','slice hidden to response span','loss','loss.ts:160'],
+ ['t29','head at response positions','loss','loss.ts:161'],
+ ['t30','gather target logits','loss','loss.ts:170'],
+ ['t31','logsumexp vocab','loss','loss.ts:169'],
+ ['t32','CE = lse - target','loss','loss.ts:172'],
+ ['t33','sum/normalize response CE','loss','loss.ts:174'],
+ ['t34','build response mask [B,T]','loss','loss.ts:233'],
+ ['t35','reshape logits [B*T,V]','loss','loss.ts:251'],
+ ['t36','gather targets','loss','loss.ts:252'],
+ ['t37','logsumexp','loss','loss.ts:256'],
+ ['t38','CE * mask','loss','loss.ts:262'],
+ ['t39','sum masked CE / norm','loss','loss.ts:264'],
+ ['t40','DPO ref logps (scale=0, detach)','loss','loss.ts:361'],
+ ['t41','DPO policy logps','loss','loss.ts:389'],
+ ['t42','logr_c = piC - refC','loss','loss.ts:392'],
+ ['t43','logr_r = piR - refR','loss','loss.ts:393'],
+ ['t44','adv = β(logr_c - logr_r)','loss','loss.ts:394'],
+ ['t45','-logσ(adv) = softplus; mean','loss','loss.ts:399'],
+ ['t46','ORPO chosen logp','loss','loss.ts:1242'],
+ ['t47','? orpoChunkSize > 0','loss','loss.ts:481'],
+ ['t48','? orpoFusedCe ∨ orpoFlashCe','loss','loss.ts:482'],
+ ['t49','chunked head (Checkpoint)','loss','loss.ts:567'],
+ ['t50','fused linear-CE head (CustomVjp)','loss','loss.ts:726'],
+ ['t51','vocab-blocked head','loss','loss.ts:822'],
+ ['t52','flash-CCE Metal head','loss','loss.ts:1009'],
+ ['t53','whole-vocab head (oracle)','loss','loss.ts:1087'],
+ ['t54','ORPO rejected logp','loss','loss.ts:1243'],
+ ['t55','NLL = mean(-ℓw)','loss','loss.ts:1225'],
+ ['t56','log_odds calc','loss','loss.ts:1207'],
+ ['t57','OR = mean softplus(-log_odds)','loss','loss.ts:1228'],
+ ['t58','L = NLL + λ·OR','loss','loss.ts:1231'],
+ ['t59','seg loss-head VJP (cot=1)','seg','segmented.ts:141'],
+ ['t60','seg loss scalar (detached)','seg','segmented.ts:204'],
+ ['t61','seg dh_out = dLoss/dhidden','seg','segmented.ts:205'],
+ ['t62','evalAll([loss, …grads])','accum','trainer.ts:1027'],
+ ['t63','loss scalar toFloat32','accum','trainer.ts:1028'],
+ ['t64','scale grad by 1/accum','accum','trainer.ts:1042'],
+ ['t65','accumulate grad','accum','trainer.ts:1048'],
+ ['t66','dispose checkpoint closures','accum','trainer.ts:496'],
+ ['t67','sum micro losses','accum','trainer.ts:1061'],
+ ['t69','value_and_grad backward','bwd','trainer.ts:492'],
+ ['t70','MLX reverse-mode AD','bwd','trainer.ts:492'],
+ ['t71','restore LoRA leaves','bwd','trainer.ts:436'],
+ ['t72','seg backward: k=nSeg-1..0','bwd','segmented.ts:210'],
+ ['t73','vjp segment k (cot=dh_out)','bwd','segmented.ts:215'],
+ ['t74','scatter per-seg LoRA grads','bwd','segmented.ts:220'],
+ ['t75','dh_in → cotangent for k-1','bwd','segmented.ts:225'],
+ ['t76','dispose boundary k','bwd','segmented.ts:226'],
+ ['t80','globalNorm(grads)','norm','trainer.ts:1076'],
+ ['t81','square each grad','norm','trainer.ts:1078'],
+ ['t82','sum squares, sqrt','norm','trainer.ts:1086'],
+ ['t83','? ‖g‖ > maxNorm','norm','trainer.ts:1099'],
+ ['t84','scale g *= max/‖g‖','norm','trainer.ts:1103'],
+ ['t85','AdamW.step(grads)','opt','trainer.ts:504'],
+ ['t86','t++ step counter','opt','optimizer.ts:84'],
+ ['t87','bias-correct β1^t','opt','optimizer.ts:85'],
+ ['t88','bias-correct β2^t','opt','optimizer.ts:86'],
+ ['t89','for each param i (A,B)','opt','optimizer.ts:88'],
+ ['t90','grad → float32','opt','optimizer.ts:92'],
+ ['t91','m ← β1·m + (1-β1)g','opt','optimizer.ts:95'],
+ ['t92','v ← β2·v + (1-β2)g²','opt','optimizer.ts:105'],
+ ['t93','? weight decay','opt','optimizer.ts:119'],
+ ['t94','m̂ = m/(1-β1^t)','opt','optimizer.ts:125'],
+ ['t95','v̂ = v/(1-β2^t)','opt','optimizer.ts:126'],
+ ['t96','√v̂ + ε','opt','optimizer.ts:127'],
+ ['t97','m̂/(√v̂+ε)','opt','optimizer.ts:130'],
+ ['t98','scale by lr (LoRA+ per-param)','opt','optimizer.ts:131'],
+ ['t99','p = p - lr·update','opt','optimizer.ts:132'],
+ ['t100','writeParam → swap into LoRA','opt','trainer.ts:402'],
+ ['t101','opt.evalState()','opt','trainer.ts:505'],
+ ['t102','clearCache()','metrics','trainer.ts:511'],
+ ['t103','emit step metrics','metrics','trainer.ts:516'],
+ ['t104','? step%report==0 ∨ step==1','metrics','trainer.ts:514'],
+ ['t105','recompute dpo/orpo metrics','metrics','trainer.ts:900'],
+ ['t106','? step%eval==0 & valid>0','metrics','trainer.ts:529'],
+ ['t107','disable grad-ckpt for eval','metrics','trainer.ts:532'],
+ ['t108','loop validation examples','metrics','trainer.ts:554'],
+ ['t109','eval loss per example','metrics','trainer.ts:564'],
+ ['t110','reduce eval (mean/acc/margin)','metrics','trainer.ts:684'],
+ ['t111','re-enable grad-ckpt','metrics','trainer.ts:534'],
+ ['t112','step iteration completes','setup','trainer.ts:538'],
+],e:[
+ ['t1','t2',''],['t2','t3',''],['t3','t4',''],['t4','t5',''],
+ ['t5','t6','segmentSize>0'],['t5','t16','else'],
+ ['t6','t7',''],['t7','t8',''],['t8','t9',''],
+ ['t9','t10','k<nSeg'],['t10','t11',''],['t11','t9','k++'],
+ ['t11','t12','done (ORPO)'],['t11','t59','done (SFT seg)'],
+ ['t12','t13','prefixShared & segmented'],['t12','t59','else'],
+ ['t13','t14','split non-null'],['t13','t15','null (mismatch)'],
+ ['t14','t59',''],['t15','t72',''],
+ ['t16','t17',''],['t17','t18',''],
+ ['t18','t19','B==1'],['t18','t20','full-logits'],
+ ['t19','t20',''],['t20','t21',''],['t21','t22',''],
+ ['t22','t24','flash'],['t22','t23','else'],
+ ['t23','t25',''],['t24','t25',''],['t25','t26',''],
+ ['t19','t27','SFT B==1'],['t26','t27','SFT B>1'],
+ ['t27','t28','B==1'],['t27','t34','B>1'],
+ ['t28','t29',''],['t29','t30',''],['t30','t31',''],['t31','t32',''],['t32','t33',''],
+ ['t34','t35',''],['t35','t36',''],['t36','t37',''],['t37','t38',''],['t38','t39',''],
+ // DPO
+ ['t4','t40','method==dpo'],['t40','t41',''],['t41','t42',''],['t42','t43',''],['t43','t44',''],['t44','t45',''],
+ // ORPO
+ ['t4','t12','method==orpo & prefixShared'],['t12','t46','else'],['t13','t46','null → two-forward'],
+ ['t46','t47',''],
+ ['t47','t48','chunkSize>0'],['t47','t54','batched'],
+ ['t48','t49','neither'],['t48','t50','fused'],['t48','t52','flash'],
+ ['t50','t51','vocabBlock>0'],['t50','t53','else'],
+ ['t49','t54',''],['t51','t54',''],['t52','t54',''],['t53','t54',''],
+ ['t14','t55','prefix-shared done'],['t54','t55',''],
+ ['t55','t56',''],['t56','t57',''],['t57','t58',''],
+ // seg head
+ ['t59','t60',''],['t60','t61',''],
+ // accumulation
+ ['t33','t62','SFT B==1'],['t39','t62','SFT B>1'],['t45','t62','DPO'],['t58','t62','ORPO'],['t60','t62','seg'],
+ ['t62','t63',''],
+ ['t63','t64','accum>1'],['t63','t66','else'],
+ ['t64','t65',''],['t65','t66',''],
+ ['t66','t64','next micro'],['t66','t67','last micro'],
+ // backward
+ ['t62','t69','non-seg'],['t69','t70',''],['t70','t71',''],
+ ['t61','t72','seg'],['t72','t73','k>=0'],['t73','t74',''],['t74','t75',''],['t75','t76',''],['t76','t72','k--'],
+ // norm/clip
+ ['t71','t80',''],['t72','t80','seg done'],['t67','t80','accum done'],
+ ['t80','t81',''],['t81','t82',''],['t82','t83',''],
+ ['t83','t84','‖g‖>max'],['t83','t85','else'],['t84','t85',''],
+ // optimizer
+ ['t85','t86',''],['t86','t87',''],['t87','t88',''],['t88','t89',''],
+ ['t89','t90','i<len'],['t90','t91',''],['t91','t92',''],['t92','t93',''],
+ ['t93','t94','wd>0 / else'],['t94','t95',''],['t95','t96',''],['t96','t97',''],['t97','t98',''],['t98','t99',''],['t99','t100',''],['t100','t89','i++'],
+ ['t89','t101','done'],
+ // metrics
+ ['t101','t102',''],['t102','t103',''],['t103','t104',''],
+ ['t104','t105','report'],['t104','t106','skip'],['t105','t106',''],
+ ['t106','t107','eval'],['t106','t112','skip'],
+ ['t107','t108',''],['t108','t109','ex'],['t109','t108','next'],['t108','t110','done'],['t110','t111',''],['t111','t112',''],
+ ['t112','t1','step++'],
+]};
+let MODE_G='infer';
+const cy=cytoscape({container:document.getElementById('cy'),minZoom:0.06,maxZoom:3,wheelSensitivity:0.25,
+  autoungrabify:true, boxSelectionEnabled:false, selectionType:'single',
+  style:[
+    {selector:'node',style:{'background-color':d=>GRP[d.data('grp')]||'#6e7b91','label':'data(label)',
+      'color':'#0b0f14','font-family':'ui-monospace, Menlo, monospace','font-size':'11px','font-weight':'600',
+      'text-valign':'center','text-halign':'center','shape':'round-rectangle','width':'label','height':22,
+      'padding':'7px','text-wrap':'wrap','text-max-width':'210px','border-width':1,'border-color':'#0b0f14'}},
+    {selector:'node[branch]',style:{'border-width':2.5,'border-color':'#39c5cf'}},
+    {selector:'node:selected',style:{'border-width':3,'border-color':'#fff'}},
+    {selector:'edge',style:{'width':1.6,'line-color':'#3a4760','target-arrow-color':'#3a4760',
+      'target-arrow-shape':'triangle','curve-style':'bezier','arrow-scale':0.8}},
+    {selector:'edge[cond]',style:{'line-color':'#39c5cf','target-arrow-color':'#39c5cf','line-style':'dashed',
+      'label':'data(cond)','font-size':'8px','font-family':'ui-monospace, Menlo, monospace','color':'#39c5cf',
+      'text-background-color':'#0b0f14','text-background-opacity':0.85,'text-background-padding':'1px','text-rotation':'autorotate'}},
+  ]});
+function build(g){
+  // count outgoing per node to mark branch points (>1 distinct outgoing edge)
+  const out={}; g.e.forEach(([a])=>out[a]=(out[a]||0)+1);
+  const nodes=g.n.map(([id,label,grp,file])=>({data:{id,label,grp,file,branch:out[id]>1?1:undefined}}));
+  const edges=g.e.map(([a,b,c])=>({data:{id:a+'>'+b+(c||''),source:a,target:b,...(c?{cond:c}:{})}}));
+  return {nodes,edges};
+}
+function load(g){const el=build(g);cy.elements().remove();cy.add(el.nodes);cy.add(el.edges);
+  cy.layout({name:'dagre',rankDir:'TB',nodeSep:18,rankSep:30,padding:30}).run();}
+function show(d){
+  document.getElementById('s-title').textContent=d.label;
+  document.getElementById('s-file').textContent=d.file||'';
+  document.getElementById('hint').style.display='none';
+  document.getElementById('s-body').style.display='block';
+  document.getElementById('s-fileb').textContent=d.file||'';
+  document.getElementById('s-grp').textContent=d.grp||'';
+  const outs=cy.getElementById(d.id).outgoers('edge').map(e=>{
+    const t=e.target().data('label'); const c=e.data('cond'); return c?`• [${c}] → ${t}`:`• → ${t}`;}).join('\n');
+  document.getElementById('s-out').textContent=outs||'(terminal)';
+}
+cy.on('tap','node',e=>show(e.target.data()));
+document.getElementById('fit').onclick=()=>cy.fit(undefined,40);
+const LABELS={drive:'generate',fwd:'forward',layer:'layer',attn:'attention',sdpa:'sdpa dispatch',unfused:'unfused sdpa',tiled:'tiled sdpa',mlp:'mlp',kv:'kv cache',samp:'sampler',
+  setup:'setup/loop',loss:'loss head',seg:'segmented',accum:'accum',bwd:'backward',norm:'grad norm/clip',opt:'optimizer',metrics:'metrics/eval'};
+function legend(){
+  const groups = MODE_G==='infer' ? ['drive','fwd','attn','sdpa','tiled','unfused','mlp','kv','samp'] : ['setup','fwd','loss','seg','accum','bwd','norm','opt','metrics'];
+  document.getElementById('legend').innerHTML='<span><i class="dash"></i>branch (condition)</span>'+
+    groups.map(g=>`<span><i class="dot" style="background:${GRP[g]}"></i>${LABELS[g]}</span>`).join('');
+}
+document.querySelectorAll('#graph button').forEach(b=>b.onclick=()=>{
+  document.querySelectorAll('#graph button').forEach(x=>x.classList.remove('on'));b.classList.add('on');
+  MODE_G=b.dataset.g;load(MODE_G==='infer'?INFER:TRAIN);legend();});
+legend();load(INFER);
+</script>
+</body>
+</html>