| export default "// Decode-step MEGAKERNEL: one workgroup computes one batch row's ENTIRE\n// decoder layer — 8 dispatches collapse into 1 (small-B latency: the b1 step\n// is dispatch-overhead-bound, ~15µs of fixed cost per dispatch against\n// ~5-10µs kernels). With IF_EMBED (layer 0) the decode embedding folds in\n// too, so a b1 step becomes: mega L0 → mega L1 → lm_head → argmax.\n//\n// Stages (workgroupBarrier between each; all math f32; every stage boundary\n// value is ROUNDED THROUGH THE STORAGE TYPE first — replicating the unfused\n// path's activation-buffer round trips, the gemm_row_ln precedent):\n// S0 x ← embed (L0: ring/table/pos, matching embed.wgsl DECODE) or the\n// global hidden buffer X (L1)\n// S1 qkv = x·Wqkvᵀ + b; q stays in shared, k|v quads store to the caches\n// at position t (the same rounded value — the kv_append contract)\n// S2 self-attention over positions 0..t (phase structure and score/exp/\n// fold order mirror attention.wgsl; heads loop serially)\n// S3 x = LN1(x + self_out(attn))\n// S4 q = cross_q(x)\n// S5 cross-attention over lens[b] encoder positions (fused crossKV k|v)\n// S6 x = LN2(x + cross_out(attn))\n// S7 ffn = SiLU(fc1(x)) (SiLU in f32 BEFORE the f16 round, as gemv)\n// S8 X ← LN3(x + fc2(ffn)) (written back to the global hidden buffer)\n//\n// All projections read the ORIGINAL [K, N] row-major `.weight` tensors with\n// the gemm_row_ln access pattern: thread q owns output quad q, and at each k\n// the threads read CONSECUTIVE quads of W's k-row — fully coalesced. (The\n// first version walked the transposed [N,K] copies, one row per thread —\n// every load touched 32 distinct lines and the whole kernel ran at ~4.5GB/s,\n// 4.3× SLOWER than the chain it replaced. One workgroup has only ~8 warps of\n// latency-hiding; coalescing is everything here.) Four independent\n// accumulators per k-quad keep 4 loads in flight per thread. Every tensor is\n// addressed inside the ONE weights buffer via compile-time vec4-element\n// offsets (…4 defines = byteOffset/8; manifest offsets are 256-aligned so /8\n// is exact). One pipeline per layer.\n//\n// NOT bit-exact vs the unfused chain (accumulation/reduction ORDER differs\n// at every site) — routed like every kernel change: m3/golden gates + the\n// step-0 + divergence-rate equiv (mega_equiv), e2e A/B decides the batch\n// threshold.\n//\n// Shared budget (16KB): (2·HD4 + TMP4)·16 + SCORES_CAP·4 + WG·4 — 13,184B\n// for MoxhiMT-30 (448/1792), 16,256B for HachimiMT-60 (576/2304); checked at\n// dispatch. tmp4 is max(FFN4, HD4 + WG) quads: the attention phase-3 partial\n// scratch tmp4[HD4 .. HD4+WG) must fit even when the model's FFN is small\n// (q lives in [0..HD4) for self, out4 for cross; the fold result lands in\n// [0..HD4) only after all partial reads).\n//\n// Template placeholders (buildShader in pipelines.js):\n// ENABLE_F16, T (must be f16 — the .wt copies only exist as f16), WG (256)\n// ENABLE_SG + IF_SG/IF_NOSG subgroup wgMax/wgSum (flags.sg) — see below\n// IF_EMBED / IF_NOEMBED layer-0 embedding fold (TABLE4/POS4/EMBED_SCALE/\n// DECODER_START live inside IF_EMBED)\n// H, D, FFN4, LMAX, SCORES_CAP, ATTN_SCALE, EPS\n// QKVW4 QKVB4 OUTW4 OUTB4 LN1G4 LN1B4 CQW4 CQB4 COW4 COB4 LN2G4 LN2B4\n// FC1W4 FC1B4 FC2W4 FC2B4 LN3G4 LN3B4 per-tensor vec4 offsets into W\n//\n// SG mode (flags.sg): wgMax/wgSum are where this kernel's barriers live —\n// each tree call is 10 workgroupBarriers, and one layer makes 38 of them\n// (2 per attention head × 2 sides × H, 2 per LN × 3), ~400 barriers per\n// step per layer. That is exactly what the megakernel pays on Apple GPUs\n// (Metal mega_sweep: mega LOSES b1 there while winning −12% on NVIDIA).\n// With sg each call is subgroupMax/Add → one partial per subgroup → ONE\n// barrier → serial fold over ≤ WG/4 partials, ~5× fewer barriers overall.\n{{ENABLE_IMMEDIATE}}\n{{ENABLE_SG}}\n{{ENABLE_F16}}\n\nstruct Params {\n B: u32, // batch rows (grid.x)\n t: u32, // decode step (cache position; embed pos; self len = t+1)\n S: u32, // encoder crossKV position capacity (padded S)\n _pad: u32,\n}\n\n{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params;\n@group(0) @binding(1) var<storage, read> W: array<vec4<{{T}}>>; // whole weights buffer\n@group(0) @binding(2) var<storage, read> ring: array<u32>; // token ring\n@group(0) @binding(3) var<storage, read_write> Kc: array<vec4<{{T}}>>; // [B, LMAX, H·D]\n@group(0) @binding(4) var<storage, read_write> Vc: array<vec4<{{T}}>>;\n@group(0) @binding(5) var<storage, read> CKV: array<vec4<{{T}}>>; // [B·S, 2·H·D] fused k|v\n@group(0) @binding(6) var<storage, read> lens: array<u32>;\n@group(0) @binding(7) var<storage, read_write> X: array<vec4<{{T}}>>; // hidden [B, H·D]\n\nconst H: u32 = {{H}}u;\nconst D: u32 = {{D}}u;\nconst D4: u32 = D / 4u; // quads per head\nconst HD4: u32 = H * D4; // quads per d_model row\nconst QKV4: u32 = 3u * HD4; // fused q|k|v quads\nconst FFN4: u32 = {{FFN4}}u; // ffn quads (FFN/4)\nconst KQ_FFN: u32 = FFN4; // fc2 K quads (K = FFN)\nconst LMAX: u32 = {{LMAX}}u;\nconst SCORES_MAX: u32 = {{SCORES_CAP}}u;\nconst ATTN_SCALE: f32 = {{ATTN_SCALE}};\nconst WG: u32 = {{WG}}u;\nconst JT: u32 = WG / D4; // attention phase-3 j-lanes per d-quad\nconst TMP4: u32 = max(FFN4, HD4 + WG); // ffn row AND attn partial scratch fit\n\nvar<workgroup> xs4: array<vec4<f32>, HD4>; // hidden state (residual base)\nvar<workgroup> tmp4: array<vec4<f32>, TMP4>; // q / vbuf / ffn / attn partials\nvar<workgroup> out4: array<vec4<f32>, HD4>; // stage outputs\nvar<workgroup> scores: array<f32, SCORES_MAX>;\n{{IF_NOSG}}\nvar<workgroup> red: array<f32, WG>;\n{{/IF_NOSG}}\n{{IF_SG}}\n// One partial per subgroup, in TWO alternating slots of NSG_CAP (WG/4 covers\n// the spec-minimum subgroup size 4). The alternation is what buys the single\n// barrier per call: call N's fold reads slot A strictly before every thread\n// passes call N+1's barrier (slot B), and call N+2's elect-writes to slot A\n// happen strictly after it — so no trailing barrier is needed to protect\n// reuse. sgId = tid/sgSize assumes the linear tid→subgroup layout (same bet\n// as add_layernorm.wgsl; the equiv gates catch a violating backend).\nconst NSG_CAP: u32 = WG / 4u;\nvar<workgroup> red: array<f32, 2u * NSG_CAP>;\nvar<private> sgId: u32;\nvar<private> nSg: u32;\nvar<private> redSlot: u32 = 0u;\n{{/IF_SG}}\n\nfn wgMax(tid: u32, v: f32) -> f32 {\n{{IF_SG}}\n let s1 = subgroupMax(v);\n let base = redSlot * NSG_CAP;\n if (subgroupElect()) { red[base + sgId] = s1; }\n workgroupBarrier();\n var r = red[base];\n for (var i = 1u; i < nSg; i = i + 1u) { r = max(r, red[base + i]); }\n redSlot = 1u - redSlot;\n return r;\n{{/IF_SG}}\n{{IF_NOSG}}\n red[tid] = v;\n workgroupBarrier();\n for (var s = WG / 2u; s > 0u; s = s >> 1u) {\n if (tid < s) { red[tid] = max(red[tid], red[tid + s]); }\n workgroupBarrier();\n }\n let r = red[0];\n workgroupBarrier(); // red[0] reads done before the next reduction reuses red\n return r;\n{{/IF_NOSG}}\n}\n\nfn wgSum(tid: u32, v: f32) -> f32 {\n{{IF_SG}}\n let s1 = subgroupAdd(v);\n let base = redSlot * NSG_CAP;\n if (subgroupElect()) { red[base + sgId] = s1; }\n workgroupBarrier();\n var r = red[base];\n for (var i = 1u; i < nSg; i = i + 1u) { r = r + red[base + i]; }\n redSlot = 1u - redSlot;\n return r;\n{{/IF_SG}}\n{{IF_NOSG}}\n red[tid] = v;\n workgroupBarrier();\n for (var s = WG / 2u; s > 0u; s = s >> 1u) {\n if (tid < s) { red[tid] = red[tid] + red[tid + s]; }\n workgroupBarrier();\n }\n let r = red[0];\n workgroupBarrier();\n return r;\n{{/IF_NOSG}}\n}\n\n// One GEMV output quad, gemm_row_ln-style: thread computes outputs\n// 4·n4 .. 4·n4+3 from the [K, N] row-major W — at each k, threads read\n// consecutive quads of the k-row (coalesced across the workgroup). srcSel\n// picks the shared source (0 = xs4, 1 = out4, 2 = tmp4); kq = K/4 source\n// quads, nq = N/4 output quads (the W row stride). Four independent\n// accumulators keep 4 loads in flight; the fold order is fixed\n// (a0+a1)+(a2+a3). Returns f32 WITHOUT rounding — the caller rounds/routes.\nfn gemvQuad(wOff: u32, bOff: u32, n4: u32, kq: u32, nq: u32, srcSel: u32) -> vec4<f32> {\n var a0 = vec4<f32>(0.0);\n var a1 = vec4<f32>(0.0);\n var a2 = vec4<f32>(0.0);\n var a3 = vec4<f32>(0.0);\n for (var k4 = 0u; k4 < kq; k4 = k4 + 1u) {\n var xq: vec4<f32>;\n if (srcSel == 0u) { xq = xs4[k4]; }\n else if (srcSel == 1u) { xq = out4[k4]; }\n else { xq = tmp4[k4]; }\n let kBase = wOff + (k4 << 2u) * nq + n4;\n a0 = fma(vec4<f32>(xq.x), vec4<f32>(W[kBase]), a0);\n a1 = fma(vec4<f32>(xq.y), vec4<f32>(W[kBase + nq]), a1);\n a2 = fma(vec4<f32>(xq.z), vec4<f32>(W[kBase + 2u * nq]), a2);\n a3 = fma(vec4<f32>(xq.w), vec4<f32>(W[kBase + 3u * nq]), a3);\n }\n return (a0 + a1) + (a2 + a3) + vec4<f32>(W[bOff + n4]);\n}\n\n@compute @workgroup_size({{WG}})\nfn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{{IF_SG}}, @builtin(subgroup_size) sgSize: u32{{/IF_SG}}) {\n // Uniform per workgroup — safe early return before the first barrier.\n if (wid.x >= params.B) { return; }\n let b = wid.x;\n let tid = lid.x;\n let t = params.t;\n{{IF_SG}}\n sgId = tid / sgSize;\n nSg = (WG + sgSize - 1u) / sgSize;\n{{/IF_SG}}\n\n // ---- S0: hidden state into xs4 ----\n{{IF_EMBED}}\n // embed.wgsl DECODE semantics: id = DECODER_START at t=0, else the ring\n // token; pos = t; y = f16round(table·EMBED_SCALE + pos_embed).\n var id: u32 = {{DECODER_START}}u;\n if (t != 0u) { id = ring[(t - 1u) * params.B + b]; }\n for (var i = tid; i < HD4; i = i + WG) {\n let e = vec4<f32>(W[{{TABLE4}}u + id * HD4 + i]) * {{EMBED_SCALE}}\n + vec4<f32>(W[{{POS4}}u + t * HD4 + i]);\n xs4[i] = vec4<f32>(vec4<{{T}}>(e));\n }\n{{/IF_EMBED}}\n{{IF_NOEMBED}}\n // Phony use: only the embed fold reads the ring, but the binding must stay\n // statically used or layout 'auto' drops @binding(2) and the bind group\n // (which always supplies it) fails validation — killing the whole submit.\n _ = ring[0];\n for (var i = tid; i < HD4; i = i + WG) {\n xs4[i] = vec4<f32>(X[b * HD4 + i]);\n }\n{{/IF_NOEMBED}}\n workgroupBarrier();\n\n // ---- S1: fused qkv projection; q → tmp4[0..HD4), k|v quads → caches ----\n let kvBase = (b * LMAX + t) * HD4;\n for (var n4 = tid; n4 < QKV4; n4 = n4 + WG) {\n let g = vec4<{{T}}>(gemvQuad({{QKVW4}}u, {{QKVB4}}u, n4, HD4, QKV4, 0u));\n if (n4 < HD4) { tmp4[n4] = vec4<f32>(g); }\n else if (n4 < 2u * HD4) { Kc[kvBase + n4 - HD4] = g; }\n else { Vc[kvBase + n4 - 2u * HD4] = g; }\n }\n workgroupBarrier();\n\n // ---- S2: self-attention over positions 0..t (attention.wgsl phases) ----\n {\n let len = min(t + 1u, LMAX);\n for (var h = 0u; h < H; h = h + 1u) {\n let hq = h * D4;\n var lm: f32 = -1e30;\n for (var j = tid; j < len; j = j + WG) {\n let koff = (b * LMAX + j) * HD4 + hq;\n var dot4 = vec4<f32>(0.0);\n for (var i = 0u; i < D4; i = i + 1u) {\n dot4 = dot4 + tmp4[hq + i] * vec4<f32>(Kc[koff + i]);\n }\n let sc = (dot4.x + dot4.y + dot4.z + dot4.w) * ATTN_SCALE;\n scores[j] = sc;\n lm = max(lm, sc);\n }\n let rowMax = wgMax(tid, lm);\n var ls: f32 = 0.0;\n for (var j = tid; j < len; j = j + WG) {\n let e = exp(scores[j] - rowMax);\n scores[j] = e;\n ls = ls + e;\n }\n let denom = wgSum(tid, ls);\n let dq = tid % D4;\n let jg = tid / D4;\n var acc = vec4<f32>(0.0);\n if (jg < JT) {\n for (var j = jg; j < len; j = j + JT) {\n acc = acc + scores[j] * vec4<f32>(Vc[(b * LMAX + j) * HD4 + hq + dq]);\n }\n }\n tmp4[HD4 + tid] = acc; // partial scratch; q region [0..HD4) untouched\n workgroupBarrier();\n if (tid < D4) {\n var o = vec4<f32>(0.0);\n for (var g = 0u; g < JT; g = g + 1u) { o = o + tmp4[HD4 + g * D4 + tid]; }\n out4[hq + tid] = vec4<f32>(vec4<{{T}}>(o / denom));\n }\n workgroupBarrier(); // out4 + partial reads done before the next head\n }\n }\n\n // ---- S3: x = LN1(x + self_out(attn)); vbuf = tmp4 ----\n for (var n4 = tid; n4 < HD4; n4 = n4 + WG) {\n let g = vec4<{{T}}>(gemvQuad({{OUTW4}}u, {{OUTB4}}u, n4, HD4, HD4, 1u));\n tmp4[n4] = vec4<f32>(g) + xs4[n4];\n }\n workgroupBarrier();\n {\n var s: f32 = 0.0;\n for (var i = tid; i < HD4; i = i + WG) {\n let v = tmp4[i];\n s = s + v.x + v.y + v.z + v.w;\n }\n let mu = wgSum(tid, s) / f32(H * D);\n var sq: f32 = 0.0;\n for (var i = tid; i < HD4; i = i + WG) {\n let dv = tmp4[i] - vec4<f32>(mu);\n sq = sq + dot(dv, dv);\n }\n let inv = inverseSqrt(wgSum(tid, sq) / f32(H * D) + {{EPS}});\n for (var i = tid; i < HD4; i = i + WG) {\n let o = vec4<f32>(W[{{LN1G4}}u + i]) * (tmp4[i] - vec4<f32>(mu)) * inv\n + vec4<f32>(W[{{LN1B4}}u + i]);\n xs4[i] = vec4<f32>(vec4<{{T}}>(o));\n }\n }\n workgroupBarrier();\n\n // ---- S4: q = cross_q(x) → out4 ----\n for (var n4 = tid; n4 < HD4; n4 = n4 + WG) {\n out4[n4] = vec4<f32>(vec4<{{T}}>(gemvQuad({{CQW4}}u, {{CQB4}}u, n4, HD4, HD4, 0u)));\n }\n workgroupBarrier();\n\n // ---- S5: cross-attention over lens[b] encoder positions → tmp4[0..HD4) ----\n {\n let len = min(lens[b], SCORES_MAX);\n for (var h = 0u; h < H; h = h + 1u) {\n let hq = h * D4;\n var lm: f32 = -1e30;\n for (var j = tid; j < len; j = j + WG) {\n let koff = (b * params.S + j) * 2u * HD4 + hq; // k slice at offset 0\n var dot4 = vec4<f32>(0.0);\n for (var i = 0u; i < D4; i = i + 1u) {\n dot4 = dot4 + out4[hq + i] * vec4<f32>(CKV[koff + i]);\n }\n let sc = (dot4.x + dot4.y + dot4.z + dot4.w) * ATTN_SCALE;\n scores[j] = sc;\n lm = max(lm, sc);\n }\n let rowMax = wgMax(tid, lm);\n var ls: f32 = 0.0;\n for (var j = tid; j < len; j = j + WG) {\n let e = exp(scores[j] - rowMax);\n scores[j] = e;\n ls = ls + e;\n }\n let denom = wgSum(tid, ls);\n let dq = tid % D4;\n let jg = tid / D4;\n var acc = vec4<f32>(0.0);\n if (jg < JT) {\n for (var j = jg; j < len; j = j + JT) {\n // v slice at element offset H·D within the fused k|v position\n acc = acc + scores[j] * vec4<f32>(CKV[(b * params.S + j) * 2u * HD4 + HD4 + hq + dq]);\n }\n }\n tmp4[HD4 + tid] = acc; // fold results land in [0..HD4) only afterwards\n workgroupBarrier();\n if (tid < D4) {\n var o = vec4<f32>(0.0);\n for (var g = 0u; g < JT; g = g + 1u) { o = o + tmp4[HD4 + g * D4 + tid]; }\n tmp4[hq + tid] = vec4<f32>(vec4<{{T}}>(o / denom));\n }\n workgroupBarrier();\n }\n }\n\n // ---- S6: x = LN2(x + cross_out(attn)); vbuf = out4 ----\n for (var n4 = tid; n4 < HD4; n4 = n4 + WG) {\n let g = vec4<{{T}}>(gemvQuad({{COW4}}u, {{COB4}}u, n4, HD4, HD4, 2u));\n out4[n4] = vec4<f32>(g) + xs4[n4];\n }\n workgroupBarrier();\n {\n var s: f32 = 0.0;\n for (var i = tid; i < HD4; i = i + WG) {\n let v = out4[i];\n s = s + v.x + v.y + v.z + v.w;\n }\n let mu = wgSum(tid, s) / f32(H * D);\n var sq: f32 = 0.0;\n for (var i = tid; i < HD4; i = i + WG) {\n let dv = out4[i] - vec4<f32>(mu);\n sq = sq + dot(dv, dv);\n }\n let inv = inverseSqrt(wgSum(tid, sq) / f32(H * D) + {{EPS}});\n for (var i = tid; i < HD4; i = i + WG) {\n let o = vec4<f32>(W[{{LN2G4}}u + i]) * (out4[i] - vec4<f32>(mu)) * inv\n + vec4<f32>(W[{{LN2B4}}u + i]);\n xs4[i] = vec4<f32>(vec4<{{T}}>(o));\n }\n }\n workgroupBarrier();\n\n // ---- S7: ffn = SiLU(fc1(x)) → tmp4[0..FFN4) (SiLU in f32, then round) ----\n for (var n4 = tid; n4 < FFN4; n4 = n4 + WG) {\n var v = gemvQuad({{FC1W4}}u, {{FC1B4}}u, n4, HD4, FFN4, 0u);\n v = v / (vec4<f32>(1.0) + exp(-v));\n tmp4[n4] = vec4<f32>(vec4<{{T}}>(v));\n }\n workgroupBarrier();\n\n // ---- S8: X ← LN3(x + fc2(ffn)); vbuf = out4 ----\n for (var n4 = tid; n4 < HD4; n4 = n4 + WG) {\n let g = vec4<{{T}}>(gemvQuad({{FC2W4}}u, {{FC2B4}}u, n4, KQ_FFN, HD4, 2u));\n out4[n4] = vec4<f32>(g) + xs4[n4];\n }\n workgroupBarrier();\n {\n var s: f32 = 0.0;\n for (var i = tid; i < HD4; i = i + WG) {\n let v = out4[i];\n s = s + v.x + v.y + v.z + v.w;\n }\n let mu = wgSum(tid, s) / f32(H * D);\n var sq: f32 = 0.0;\n for (var i = tid; i < HD4; i = i + WG) {\n let dv = out4[i] - vec4<f32>(mu);\n sq = sq + dot(dv, dv);\n }\n let inv = inverseSqrt(wgSum(tid, sq) / f32(H * D) + {{EPS}});\n for (var i = tid; i < HD4; i = i + WG) {\n let o = vec4<f32>(W[{{LN3G4}}u + i]) * (out4[i] - vec4<f32>(mu)) * inv\n + vec4<f32>(W[{{LN3B4}}u + i]);\n X[b * HD4 + i] = vec4<{{T}}>(o);\n }\n }\n}\n"; | |