diff --git "a/assets/preload-helper-DpiWnaYd.js" "b/assets/preload-helper-DpiWnaYd.js" new file mode 100644--- /dev/null +++ "b/assets/preload-helper-DpiWnaYd.js" @@ -0,0 +1,2530 @@ +function gr(e=globalThis.navigator?.gpu){return!!e?.wgslLanguageFeatures?.has?.("immediate_address_space")}const Qn=new WeakSet;function wr(e){return!!e&&Qn.has(e)}async function jt({requireF16:e=!1}={}){if(!navigator.gpu)throw new Error("WebGPU unavailable");const n=await navigator.gpu.requestAdapter({powerPreference:"high-performance"});if(!n)throw new Error("no adapter");const t=["shader-f16","timestamp-query","subgroups"].filter(s=>n.features.has(s));if(e&&!t.includes("shader-f16"))throw new Error("shader-f16 unsupported");const o=await n.requestDevice({requiredFeatures:t,requiredLimits:{maxStorageBufferBindingSize:Math.min(536870912,n.limits.maxStorageBufferBindingSize),maxBufferSize:Math.min(1073741824,n.limits.maxBufferSize)}}),a=gr(navigator.gpu);return a&&Qn.add(o),{device:o,adapterInfo:{vendor:n.info?.vendor??"",architecture:n.info?.architecture??"",device:n.info?.device??"",description:n.info?.description??""},hasF16:t.includes("shader-f16"),hasTimestamps:t.includes("timestamp-query"),hasSubgroups:t.includes("subgroups"),hasImmediates:a,subgroupMinSize:n.info?.subgroupMinSize??0,subgroupMaxSize:n.info?.subgroupMaxSize??0,limits:{maxStorageBufferBindingSize:o.limits.maxStorageBufferBindingSize,maxBufferSize:o.limits.maxBufferSize,maxComputeWorkgroupStorageSize:o.limits.maxComputeWorkgroupStorageSize,maxComputeInvocationsPerWorkgroup:o.limits.maxComputeInvocationsPerWorkgroup},adapterLimits:(()=>{const s={};try{for(const u of Object.getOwnPropertyNames(Object.getPrototypeOf(n.limits))){const l=n.limits[u];typeof l=="number"&&Number.isFinite(l)&&(s[u]=l)}}catch{}return s})(),features:(()=>{try{return[...n.features]}catch{return[]}})()}}let R,ae,se,fe,qe,Fn,_e,Xn,Ge,ne,fn,Oe,en,xe,pn,Mn,cn,ke;const Yn=1.2,xn=["dModel","heads","headDim","ffn","encLayers","decLayers","vocab","maxPos","srcCap","decodeCap","eos","pad","decoderStart"];function Jn(e){if(!e||typeof e!="object")throw new Error("model config: not an object");const n={};for(const r of xn){const t=e[r];if(!Number.isInteger(t)||t<0)throw new Error(`model config: ${r} must be a non-negative integer, got ${t}`);n[r]=t}for(const r of xn.slice(0,10))if(n[r]===0)throw new Error(`model config: ${r} must be positive`);if(!(typeof e.lnEps=="number"&&e.lnEps>0&&e.lnEps<.1))throw new Error(`model config: lnEps must be in (0, 0.1), got ${e.lnEps}`);if(n.lnEps=e.lnEps,n.heads*n.headDim!==n.dModel)throw new Error(`model config: heads·headDim = ${n.heads*n.headDim} != dModel ${n.dModel}`);for(const r of["dModel","headDim","ffn"])if(n[r]%4!==0)throw new Error(`model config: ${r} = ${n[r]} must be a multiple of 4 (vec4 kernels)`);if(n.srcCap%32!==0)throw new Error(`model config: srcCap ${n.srcCap} must be a multiple of 32 (bucketFor invariant)`);if(n.srcCap>n.maxPos||n.decodeCap>n.maxPos)throw new Error(`model config: srcCap ${n.srcCap} / decodeCap ${n.decodeCap} must not exceed maxPos ${n.maxPos}`);for(const r of["eos","pad","decoderStart"])if(n[r]>=n.vocab)throw new Error(`model config: ${r} = ${n[r]} out of vocab ${n.vocab}`);return e.scaleEmbedding===!1&&(n.scaleEmbedding=!1),n.embedScale=n.scaleEmbedding===!1?1:Math.sqrt(n.dModel),n}let Zn=null;function $t(){return{...Zn}}function er(e){const n=Jn(e);return Zn=n,R=n.dModel,ae=n.heads,se=n.headDim,fe=n.ffn,qe=n.vocab,Fn=n.encLayers,_e=n.decLayers,Xn=n.maxPos,Ge=n.srcCap,ne=n.decodeCap,fn=Math.ceil(Math.max(n.srcCap,n.decodeCap)/32)*32+32,Oe=n.eos,en=n.pad,xe=n.decoderStart,pn=n.lnEps,Mn=n.embedScale,cn=1/Math.sqrt(n.headDim),ke=Math.ceil(n.vocab/32),n}function Dn(e,n="weights"){const r=e??{};if(r.dModel!==R||r.heads!==ae||r.headDim!==se||r.ffn!==fe||r.vocab!==qe||r.encLayers!==Fn||r.decLayers!==_e||r.maxPos!==Xn||r.srcCap!==Ge||r.decodeCap!==ne||r.eos!==Oe||r.pad!==en||r.decoderStart!==xe)throw new Error(`${n} belong to a different model config than the active one — call applyModelConfig(weights.model) (loadWeights/uploadParsed do it) before dispatching`)}er({dModel:448,heads:8,headDim:56,ffn:1792,encLayers:8,decLayers:2,vocab:24e3,maxPos:512,srcCap:320,decodeCap:224,eos:2,pad:0,decoderStart:0,lnEps:1e-5});const br=134217728;function vr(e){if(e>Ge)throw new Error(`input too long: ${e} tokens (max ${Ge})`);return Math.max(32,Te(e,32))}function yr(e){return Math.max(24,Math.min(ne,Math.ceil(e*3.2)+12))}function nr(e,n,r){if(!Number.isFinite(n)||n<=0)throw new Error(`maxBatchForLimits: bad S=${n}`);if(!Number.isFinite(r)||r<=0)throw new Error(`maxBatchForLimits: bad dtypeBytes=${r}`);const t=e?.limits?.maxStorageBufferBindingSize??br,o=Math.max(n*fe*r,ne*R*r,qe*4);return Math.max(1,Math.floor(t/o))}function Te(e,n){return Math.ceil(e/n)*n}function kr(e){const n=e&32768?-1:1,r=e>>10&31,t=e&1023;return r===0?n*t*2**-24:r===31?t?NaN:n*(1/0):n*(1024+t)*2**(r-25)}const _r=new Float32Array(1);new Uint32Array(_r.buffer);function rr(e){const n=e.length,r=new Float32Array(n);for(let t=0;t>10&31,s=o&1023,u=o&32768?-1:1;a===0?r[t]=u*s*2**-24:a===31?r[t]=s?NaN:u*(1/0):r[t]=u*(1024+s)*2**(a-25)}return r}const tr={f16:2,f32:4};function An(e){if(!e||typeof e!="object")throw new Error("manifest: not an object");if(e.version!==1)throw new Error(`manifest: expected version 1, got ${e.version}`);if(!e.model||typeof e.model!="object")throw new Error("manifest: missing model");if(!Array.isArray(e.tensors)||e.tensors.length===0)throw new Error("manifest: tensors must be a non-empty array");const n=new Map;let r=0,t=null;for(const o of e.tensors){const{name:a,dtype:s,shape:u,byteOffset:l,byteLength:c}=o??{};if(typeof a!="string"||!a)throw new Error("manifest: tensor without a name");if(s!=="f16"&&s!=="f32")throw new Error(`manifest: tensor "${a}" has unsupported dtype "${s}"`);if(!Array.isArray(u)||u.length===0||!u.every(p=>Number.isInteger(p)&&p>0))throw new Error(`manifest: tensor "${a}" has invalid shape ${JSON.stringify(u)}`);if(!Number.isInteger(l)||l<0||!Number.isInteger(c)||c<=0)throw new Error(`manifest: tensor "${a}" has invalid byteOffset/byteLength`);const f=u.reduce((p,g)=>p*g,1),i=f*tr[s];if(c!==i)throw new Error(`manifest: tensor "${a}" byteLength ${c} != shape·dtype ${i}`);if(l%256!==0)throw new Error(`manifest: tensor "${a}" byteOffset ${l} is not 256-aligned`);if(n.has(a))throw new Error(`manifest: duplicate tensor name "${a}"`);if(l!e.has(s)),a=[...e.keys()].filter(s=>!t.has(s));if(o.length||a.length){const s=[];throw o.length&&s.push(`missing: ${o.join(", ")}`),a.length&&s.push(`unexpected: ${a.join(", ")}`),new Error(`weights inventory mismatch — ${s.join("; ")}`)}}function sn(e){let n=0;for(const r of e.values())n=Math.max(n,r.byteOffset+r.byteLength);return n}function Sr(e,n){let r=0;const t=new Map;for(const[a,s]of e){const u=Te(r,256),l=s.elems*4;t.set(a,{dtype:"f32",shape:s.shape,elems:s.elems,byteOffset:u,byteLength:l}),r=u+l}const o=new Uint8Array(Te(r,4));for(const[a,s]of e){const u=t.get(a);if(s.dtype==="f16"){const l=new Uint16Array(n.buffer,n.byteOffset+s.byteOffset,s.elems);new Float32Array(o.buffer,u.byteOffset,s.elems).set(rr(l))}else o.set(n.subarray(s.byteOffset,s.byteOffset+s.byteLength),u.byteOffset)}return{tensors:t,bytes:o}}function or(e,n,r,t,o,a,s){let u=0;for(let c=0;c0?u/127:1;a[s]=l;for(let c=0;c>>0}}function Tr(e,n,r){if(r%4!==0)throw new Error(`quantizeQ8Rows: cols=${r} not a multiple of 4`);const t=new Uint32Array(n*r/4),o=new Float32Array(n),a=r/4;for(let s=0;s=n)throw new Error(`quantizeQ8RowsFromF16: row id ${i} out of range ${n}`);const p=i*r;for(let g=0;g=p.vocab)throw new Error(`lmHeadIds: id ${M} out of vocab ${p.vocab}`);if(v.has(M))throw new Error(`lmHeadIds: duplicate id ${M}`);v.add(M)}for(const[M,E]of[["eos",p.eos],["pad",p.pad],["decoderStart",p.decoderStart]])if(!v.has(E))throw new Error(`lmHeadIds must contain ${M} (${E})`);const k=l.length,N=new Uint32Array(Math.ceil(k/64)*64).fill(p.eos);N.set(l),h={ids:N,real:k}}const b=[],y={directEnabled:c,directTensors:0,sourceBytes:0,packedBytes:0,scaleBytes:0,peakTemporaryBytes:0,legacyPeakTemporaryBytes:0,temporaryBytesSaved:0,quantizeMs:0,tensors:[]};for(const v of m){const k=w.get(v.src);if(!k)throw new Error(`q8: no ${v.src} tensor`);const N=globalThis.performance?.now?.()??Date.now(),M=k.byteLength;let E=0,D=0,q=!1,O,[I,W]=k.shape;if(c&&k.dtype==="f16"&&!v.transpose){const K=new Uint16Array(d.buffer,d.byteOffset+k.byteOffset,k.elems),U=v.base==="lm_head"&&h?h.ids:null;O=Er(K,I,W,U),U&&(I=U.length),q=!0,E=O.temporaryBytes,D=k.elems*Float32Array.BYTES_PER_ELEMENT+(U?I*W*Float32Array.BYTES_PER_ELEMENT:0)}else{const K=k.dtype==="f16"?rr(new Uint16Array(d.buffer,d.byteOffset+k.byteOffset,k.elems)):new Float32Array(d.buffer,d.byteOffset+k.byteOffset,k.elems);E+=k.dtype==="f16"?K.byteLength:0;let U=K;if(v.transpose){const[re,H]=k.shape;U=new Float32Array(k.elems),E+=U.byteLength;for(let j=0;j"u"?null:await caches.open(Dr)}catch{return null}}async function Nr(e,n){const r=globalThis.crypto?.subtle;if(!e.sha256||!r)return;const t=new Uint8Array(await r.digest("SHA-256",n));let o="";for(const a of t)o+=a.toString(16).padStart(2,"0");if(o!==e.sha256)throw new Error(`${e.file}: SHA-256 mismatch`)}async function Ir(e,n,{onProgress:r}={}){const{tensors:t}=An(n),o=sn(t),a=n.bins??null;let s=0;const u=async(d,m,h,b,y)=>{let _=0;const F=d.body.getReader();for(;;){const{done:A,value:B}=await F.read();if(A)break;if(_+B.byteLength>b)throw new Error(`weights part larger than expected ${b} bytes`);m.set(B,h+_),_+=B.byteLength,s+=B.byteLength,r&&r(s,y)}return _};if(!a){const d=await fetch(`${e}/weights.bin`);if(!d.ok)throw new Error(`fetch weights.bin: HTTP ${d.status}`);const m=Number(d.headers.get("content-length"))||o,h=new Uint8Array(Te(Math.max(m,o),4)),b=await u(d,h,0,h.byteLength,m);if(bd+m.byteLength,0);if(l{const m=a[d],h=`${e}/${m.file}`;for(let b=0;;b++){let y=null;if(i&&b===0)try{y=await i.match(h)??null}catch{y=null}const _=y!==null;if(!y&&(y=await fetch(h),!y.ok))throw new Error(`fetch ${m.file}: HTTP ${y.status}`);let F=0;try{if(F=await u(y,c,f[d],m.byteLength,l),F!==m.byteLength)throw new Error(`${m.file}: got ${F} of declared ${m.byteLength} bytes`);await Nr(m,c.subarray(f[d],f[d]+F))}catch(A){if(s-=F,_){try{await i.delete(h)}catch{}continue}throw A}if(i&&!_)try{await i.put(h,new Response(c.subarray(f[d],f[d]+F)))}catch{}return}};let g=0;const w=async()=>{for(;;){const d=g++;if(d>=a.length)return;await p(d)}};return await Promise.all(Array.from({length:Math.min(Mr,a.length)},w)),i&&await Gr(i,e,a),c}async function Gr(e,n,r){try{const t=new Set(r.map(a=>a.file)),o=`${n}/`;for(const a of await e.keys()){const s=new URL(a.url).pathname;if(!s.startsWith(o))continue;const u=s.slice(o.length);u.includes("/")||t.has(u)||await e.delete(a)}}catch{}}async function zt(e,n="/weights",{targetDtype:r="f16",lmHeadQ8:t=!1,ffnQ8:o=!1,ffnWT:a=!1,projWT:s=!1,lmHeadIds:u=null,directQ8:l=!0,onProgress:c}={}){const f=await fetch(`${n}/manifest.json`);if(!f.ok)throw new Error(`fetch manifest.json: HTTP ${f.status}`);const i=await f.json(),{model:p,tensors:g}=An(i);Br(g,Jn(p));const w=await Ir(n,i,{onProgress:c}),d=Fr(e,i,w,{targetDtype:r,lmHeadQ8:t,ffnQ8:o,ffnWT:a,projWT:s,lmHeadIds:u,directQ8:l});return d.weightsTag=i.bins?.[0]?.file??"dev",d}const Ie=(e,n)=>{if(!e)throw new Error(`spm_tokenizer: unsupported tokenizer.json: ${n}`)},xr=4096,Or=512;class Wr{constructor(n,r={}){const{model:t,normalizer:o,pre_tokenizer:a,post_processor:s,decoder:u,added_tokens:l}=n;Ie(t?.type==="BPE"&&!Array.isArray(t.vocab),"model must be BPE with object vocab"),Ie(Array.isArray(t.merges?.[0]),"merges must use the pair-array format"),Ie(!t.ignore_merges&&!t.end_of_word_suffix&&!t.continuing_subword_suffix,"BPE suffix/ignore_merges options"),Ie(o?.type==="Sequence"&&o.normalizers.map(f=>f.type).join(",")==="Precompiled,Strip,Replace"&&o.normalizers[1].strip_left===!1&&o.normalizers[1].strip_right===!0&&o.normalizers[2].pattern?.Regex===" {2,}","normalizer must be [Precompiled, Strip(right), Replace(/ {2,}/)]"),Ie(a?.type==="Metaspace"&&a.prepend_scheme==="always","pre_tokenizer must be Metaspace(always)"),Ie(u?.type==="Metaspace","decoder must be Metaspace");const c=s?.single;Ie(s?.type==="TemplateProcessing"&&c?.length===2&&c[0].Sequence?.id==="A"&&typeof c[1].SpecialToken?.id=="string","post_processor must be TemplateProcessing [A, eos]"),Ie((l??[]).every(f=>!f.normalized&&!f.lstrip&&!f.rstrip&&!f.single_word),"added tokens must be unnormalized, no lstrip/rstrip/single_word"),this.replacement=a.replacement??"▁",this.spaceRun=o.normalizers[2].content??"",this.eosToken=c[1].SpecialToken.id,this.cleanUp=r.clean_up_tokenization_spaces??!0,this.vocab=new Map(Object.entries(t.vocab)),this.byteFallback=t.byte_fallback??!1,this.unkToken=t.unk_token??null,this.merges=new Map,t.merges.forEach(([f,i],p)=>{let g=this.merges.get(f);g||this.merges.set(f,g=new Map),g.has(i)||g.set(i,p)}),this.addedContents=[],this.specialTokens=new Set;for(const f of l??[])this.addedContents.push(f.content),this.vocab.set(f.content,f.id),f.special&&this.specialTokens.add(f.content);this.idToToken=[];for(const[f,i]of this.vocab)this.idToToken[i]=f;this.unkId=this.vocab.get(this.unkToken),this.textEncoder=new TextEncoder,this.cache=new Map}normalize(n){return n=n.replace(/[\u0001-\u0008\u000B\u000E-\u001F\u007F\u008F\u009F]/g,""),n=n.replace(/[\u0009\u000A\u000C\u000D\u00A0\u1680\u2000-\u200F\u2028\u2029\u202F\u205F\u2581\u3000\uFEFF\uFFFD]/g," "),n=n.includes("～")?n.split("～").map(r=>r.normalize("NFKC")).join("～"):n.normalize("NFKC"),n.trimEnd().replace(/ {2,}/gu,this.spaceRun)}splitOnAdded(n){const r=[];let t=0;for(let o=0;oa.length)&&(a=s);if(!a){o++;continue}o>t&&r.push(n.slice(t,o)),r.push(a),o+=a.length,t=o}return t`<0x${a.toString(16).toUpperCase().padStart(2,"0")}>`);o.every(a=>this.vocab.has(a))?r.push(...o):this.unkToken!=null&&r.push(this.unkToken)}else this.unkToken!=null&&r.push(this.unkToken);return r}encode(n,{addSpecialTokens:r=!0}={}){const t=(r?"S":"R")+n,o=n.length<=Or;if(o){const s=this.cache.get(t);if(s!==void 0)return this.cache.delete(t),this.cache.set(t,s),s.slice()}const a=this.encodeUncached(n,r);return o&&(this.cache.set(t,a.slice()),this.cache.size>xr&&this.cache.delete(this.cache.keys().next().value)),a}encodeUncached(n,r){const t=[];for(const o of this.splitOnAdded(n)){if(this.addedContents.includes(o)){t.push(o);continue}const a=this.normalize(o);if(a.length===0)continue;let s=a.replaceAll(" ",this.replacement);s.startsWith(this.replacement)||(s=this.replacement+s),t.push(...this.piecesToTokens(this.bpe(Array.from(s))))}return r&&t.push(this.eosToken),t.map(o=>this.vocab.get(o)??this.unkId)}decode(n,{skip_special_tokens:r=!1}={}){let t=Array.from(n,a=>this.idToToken[Number(a)]??this.unkToken);r&&(t=t.filter(a=>!this.specialTokens.has(a)));let o="";for(let a=0;a`/model/${e}/${n}${r?`?v=${encodeURIComponent(r)}`:""}`;async function Vt(e="moxhi",n=null){const r=async a=>{const s=await fetch(Pr(e,a,n));if(!s.ok)throw new Error(`fetch ${e}/${a}: HTTP ${s.status}`);return s.json()},[t,o]=await Promise.all([r("tokenizer.json"),r("tokenizer_config.json")]);return new Wr(t,o)}async function Rr(e,n){const r=n.length,t=[],o=[];for(const l of n){let c=e.encode(l);c.length>Ge&&(console.warn(`tokenizeBatch: truncating row of ${c.length} tokens to SRC_CAP=${Ge}`),o.push(t.length),c=c.slice(0,Ge),c[Ge-1]=Oe),t.push(c)}const a=Uint32Array.from(t,l=>l.length),s=vr(Math.max(...a)),u=new Uint32Array(r*s).fill(en);return t.forEach((l,c)=>u.set(l,c*s)),{ids:u,lens:a,B:r,S:s,truncated:o}}const P={allocBytes:0,freedBytes:0,liveBytes:0,highWaterBytes:0,liveBuffers:0,retiredBytes:0,retiredBuffers:0,submittedHighWaterBytes:0};function Qt(){return{...P,submittedBytes:P.liveBytes+P.retiredBytes,submittedBuffers:P.liveBuffers+P.retiredBuffers}}function Xt(){P.highWaterBytes=P.liveBytes,P.submittedHighWaterBytes=P.liveBytes+P.retiredBytes}function On(e){P.allocBytes+=e,P.liveBytes+=e,P.liveBuffers+=1,P.liveBytes>P.highWaterBytes&&(P.highWaterBytes=P.liveBytes);const n=P.liveBytes+P.retiredBytes;n>P.submittedHighWaterBytes&&(P.submittedHighWaterBytes=n)}function qr(e){P.freedBytes+=e,P.liveBytes-=e,P.liveBuffers-=1}function Kr(e){P.liveBytes-=e,P.liveBuffers-=1,P.retiredBytes+=e,P.retiredBuffers+=1}function Cr(e){P.retiredBytes-=e,P.retiredBuffers-=1,P.freedBytes+=e}function Ze(e){const n=[];return{buf(r,t,o){const a=Math.ceil(r/4)*4,s=e.createBuffer({label:o,size:a,usage:t});return n.push([s,a]),On(a),s},uniform(r,t){const o=Math.max(16,Math.ceil(r.byteLength/16)*16),a=e.createBuffer({label:t,size:o,usage:GPUBufferUsage.UNIFORM,mappedAtCreation:!0});return new Uint8Array(a.getMappedRange()).set(new Uint8Array(r.buffer,r.byteOffset,r.byteLength)),a.unmap(),n.push([a,o]),On(o),a},destroy(){for(const[r,t]of n)r.destroy(),qr(t);n.length=0},destroyDeferred(r){if(!r||typeof r.then!="function")throw new Error("arena.destroyDeferred needs a queue-drain promise");if(n.length===0)return Promise.resolve();const t=n.splice(0);for(const[a,s]of t)a.destroy(),Kr(s);const o=()=>{for(const[,a]of t)Cr(a)};return Promise.resolve(r).then(o,o)}}}function Lr(e){const n=Ze(e),r=Ze(e);return{retained:n,scratch:r,destroyScratch(){r.destroy()},retireScratch(t){return r.destroyDeferred(t)},destroy(){r.destroy(),n.destroy()}}}const Ur=`// GEMM: Y[m,n] = Σ_k X[m,k] · W[k,n] (+ B[n]), optional SiLU. +// X: [M,K] row-major, W: [K,N] row-major, Y: [M,N] row-major. +// With the WT flag, W is stored TRANSPOSED — [N,K] row-major — and the W +// index becomes n·K + k (the tied-embedding LM head reads shared.weight +// [24000,448] directly, no transposed copy). +// Simple correctness-first kernel — one invocation per output element, +// f32 accumulation regardless of storage type. Tuning happens in Task 17. +// +// Template placeholders (substituted by buildShader in pipelines.js — note: +// do NOT write literal placeholder syntax in comments, the substitution is a +// dumb string replace over the whole file): +// ENABLE_F16 the f16 enable directive when T or OUT_T is f16, else empty +// T storage type of X/W/B (f16|f32) +// OUT_T storage type of Y (f16|f32) +// WG workgroup size in x (default 64) +// IF_BIAS/IF_SILU/IF_WT conditional blocks +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Dims { M: u32, K: u32, N: u32, _pad: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> dims: Dims; +@group(0) @binding(1) var X: array<{{T}}>; +@group(0) @binding(2) var W: array<{{T}}>; +@group(0) @binding(3) var B: array<{{T}}>; +@group(0) @binding(4) var Y: array<{{OUT_T}}>; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(global_invocation_id) gid: vec3) { + // Keep B in the shader interface (and thus in the auto bind group layout) + // even when the bias block is compiled out — a 4-byte dummy buffer is + // bound in that case. + _ = &B; + let n = gid.x; + let m = gid.y; + if (n >= dims.N || m >= dims.M) { return; } + var acc: f32 = 0.0; + let xoff = m * dims.K; + for (var k: u32 = 0u; k < dims.K; k = k + 1u) { + var wIdx = k * dims.N + n; // W [K,N] row-major + {{IF_WT}}wIdx = n * dims.K + k; // W transposed: [N,K] row-major{{/IF_WT}} + acc = acc + f32(X[xoff + k]) * f32(W[wIdx]); + } + {{IF_BIAS}}acc = acc + f32(B[n]);{{/IF_BIAS}} + {{IF_SILU}}acc = acc / (1.0 + exp(-acc));{{/IF_SILU}} + Y[m * dims.N + n] = {{OUT_T}}(acc); +} +`,Hr=`// GEMV-style GEMM for SMALL M (decode-step projections, M = batch ≤ ~16). +// Same math/semantics as gemm.wgsl: Y[m,n] = Σ_k X[m,k]·W[k,n] (+B[n]), +// optional SiLU, f32 accumulation. Chosen by dispatchGemm({flags.gemv}). +// +// Why a separate kernel (Task 17 profile data): at M=1..8 the plain kernel +// launches only N threads, each walking the whole K serially → latency-bound +// (~150ns/iter chains, e.g. fc2 268µs at b1); and in WT mode adjacent threads +// read 2-byte elements 896B apart → ~1/64 cache-line utilization (lm_head +// 1918µs at b8). Here a workgroup of TK·TN threads computes a tile of outputs +// with K split across TK lanes + a shared-memory tree reduction: TK× shorter +// dependency chains, TK·(TN)× more threads in flight, and coalesced loads. +// +// Two layouts (exactly one of WT / NWT): +// WT — W stored [N,K] row-major (lm_head / shared.weight). K%4 == 0 +// required: X and W are bound as vec4 arrays and the lane loop walks +// K/4 vec4s. tid = o·TK + lane (lane fastest) → the TK lanes of one +// output read CONSECUTIVE vec4s of the W row (fully coalesced). +// TN outputs per workgroup; grid (ceil(N/TN), ceil(M/MT)): a +// workgroup serves up to MT rows of X from one L1-resident W tile — +// without this every batch row re-streams the whole W from DRAM +// (lm_head at b8: 8×21.5MB). +// NWT — W stored [K,N] row-major (all other projections). N%4 == 0 +// required: each thread owns a QUAD of 4 consecutive n via vec4 W +// row loads; X[k] is a scalar broadcast. tid = lane·TN + oq (quad +// fastest) → adjacent threads read adjacent vec4s within a k-row +// (coalesced); lanes stride k by TK. TN quads (4·TN outputs) per +// workgroup; grid (ceil(N/(4·TN)), M). +// +// STORE_KV (decode self_qkv, either layout): the projection output is a +// fused q|k|v row [M, 3·H·D]; the epilogue additionally scatters the k and v +// outputs into the [B, Lmax, H·D] K/V caches at decode position t — +// replacing the separate kv_append dispatch (Task 17: per-dispatch fixed +// overhead dominates the b1 step). Values are rounded through the same +// storage-type conversion as Y, so parity with the old copy-from-Y kv_append +// is bit-exact within a layout (WT vs NWT differ in accumulation order). +// Mutually exclusive with WQ8 (both claim binding 5; dispatch enforces). +// +// Out-of-range tail outputs are computed on clamped indices (uniform control +// flow for the barriers) and simply not stored. +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 f16 enable directive when T or OUT_T is f16 +// T / OUT_T storage types (f16|f32 — vec4 is core, no fallback +// split needed; the plain gemm.wgsl path remains the +// unvectorized fallback) +// WG workgroup size — MUST equal TK·TN (dispatchGemm enforces) +// TK k-lanes per output (power of two) +// TN outputs (WT) / output quads (NWT) per workgroup +// MT WT only: X rows served per workgroup (W-tile reuse) +// IF_BIAS/IF_SILU/IF_WT/IF_NWT conditional blocks +// IF_WQ8/IF_WQF WT only: int8 W8A16 weights / float weights. WQ8 binds W +// as u32 words (4×i8 along K, quantizeQ8Rows layout) with +// per-N f32 scales at binding 5, applied in the epilogue +// BEFORE bias — same contract as gemm_tiled2's WQ8 path. +// Mutually exclusive with STORE_KV (dispatch enforces). +// IF_SG/IF_NOSG WT only: subgroup reduction (flags.sg) — the TK lanes of +// one output are consecutive tids, so a subgroupShuffleDown +// chain folds them with ZERO barriers (the shared-memory +// tree costs log2(TK) barriers per MT row). Requires a +// TK-slice to never straddle a subgroup: dispatch gates on +// TK ≤ ctx.subgroupMinSize. Reduction order differs → +// tolerance-equal, not bit-equal, vs the tree. +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} +{{ENABLE_SG}} + +struct Dims { M: u32, K: u32, N: u32, _pad: u32{{IF_STORE_KV}}, t: u32, Lmax: u32, _p1: u32, _p2: u32{{/IF_STORE_KV}} } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> dims: Dims; +{{IF_WT}} +@group(0) @binding(1) var X: array>; +{{IF_WQF}} +@group(0) @binding(2) var W: array>; +{{/IF_WQF}} +{{IF_WQ8}} +@group(0) @binding(2) var W: array; // 4×i8 along K +{{/IF_WQ8}} +{{/IF_WT}} +{{IF_NWT}} +@group(0) @binding(1) var X: array<{{T}}>; +@group(0) @binding(2) var W: array>; +{{/IF_NWT}} +@group(0) @binding(3) var B: array<{{T}}>; +@group(0) @binding(4) var Y: array<{{OUT_T}}>; +{{IF_STORE_KV}} +@group(0) @binding(5) var Kc: array<{{T}}>; +@group(0) @binding(6) var Vc: array<{{T}}>; +{{/IF_STORE_KV}} +{{IF_WQ8}} +@group(0) @binding(5) var S: array; // per-N scales +{{/IF_WQ8}} + +const TK: u32 = {{TK}}u; +const TN: u32 = {{TN}}u; +const WG: u32 = {{WG}}u; + +{{IF_WT}} +const MT: u32 = {{MT}}u; +{{IF_NOSG}} +var red: array; +{{/IF_NOSG}} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + _ = &B; // keep B in the auto layout even when bias is compiled out + let tid = lid.x; + let lane = tid % TK; // lane fastest → coalesced W reads + let o = tid / TK; + let n = wid.x * TN + o; + let nc = min(n, dims.N - 1u); // clamp tail (store is guarded below) + let K4 = dims.K / 4u; + let wbase = nc * K4; + // m loop: the same W tile (L1-hot after the first row) is swept for up to + // MT rows. The break is uniform across the workgroup (wid-derived), so the + // barriers below stay in uniform control flow. red is reused per row: safe, + // because between barriers each thread only touches red slots it owns. + for (var mm = 0u; mm < MT; mm = mm + 1u) { + let m = wid.y * MT + mm; + if (m >= dims.M) { break; } + let xbase = m * K4; + var acc = vec4(0.0); + for (var i = lane; i < K4; i = i + TK) { + {{IF_WQF}} + acc = acc + vec4(X[xbase + i]) * vec4(W[wbase + i]); + {{/IF_WQF}} + {{IF_WQ8}} + // One u32 = 4 raw int8 q values along K, sign-extended by shifts; + // the per-N scale is applied once in the epilogue. + let word = bitcast(W[wbase + i]); + acc = acc + vec4(X[xbase + i]) * vec4( + f32((word << 24u) >> 24u), + f32((word << 16u) >> 24u), + f32((word << 8u) >> 24u), + f32(word >> 24u)); + {{/IF_WQ8}} + } + {{IF_NOSG}} + red[tid] = acc.x + acc.y + acc.z + acc.w; + workgroupBarrier(); + for (var s = TK / 2u; s > 0u; s = s >> 1u) { + if (lane < s) { red[tid] = red[tid] + red[tid + s]; } + workgroupBarrier(); + } + let vr = red[tid]; + {{/IF_NOSG}} + {{IF_SG}} + // Butterfly over the TK-slice: after step s, lanes < s hold sums of + // [lane, lane+2s). Lanes ≥ TK/2 read across the slice edge — harmless, + // they never feed lane 0. No barriers, no shared memory. + var vr = acc.x + acc.y + acc.z + acc.w; + for (var s = TK / 2u; s > 0u; s = s >> 1u) { + vr = vr + subgroupShuffleDown(vr, s); + } + {{/IF_SG}} + if (lane == 0u && n < dims.N) { + var v = vr; + {{IF_WQ8}}v = v * S[n]; // dequantize before bias{{/IF_WQ8}} + {{IF_BIAS}}v = v + f32(B[n]);{{/IF_BIAS}} + {{IF_SILU}}v = v / (1.0 + exp(-v));{{/IF_SILU}} + Y[m * dims.N + n] = {{OUT_T}}(v); + {{IF_STORE_KV}} + // Fused q|k|v row: outputs in [HD, 2HD) are k, [2HD, 3HD) are v. + // Same cache layout as the NWT quad scatter (kv_append contract). + let HD = dims.N / 3u; + if (n >= HD) { + let off = n - HD; + let dstBase = (m * dims.Lmax + dims.t) * HD; + if (off < HD) { + Kc[dstBase + off] = {{T}}(v); + } else { + Vc[dstBase + (off - HD)] = {{T}}(v); + } + } + {{/IF_STORE_KV}} + } + } +} +{{/IF_WT}} + +{{IF_NWT}} +var red: array, WG>; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + _ = &B; + let tid = lid.x; + let m = wid.y; + let oq = tid % TN; // quad fastest → coalesced W reads within a k-row + let lane = tid / TN; + let N4 = dims.N / 4u; + let q = wid.x * TN + oq; + let qc = min(q, N4 - 1u); // clamp tail (store is guarded below) + let xbase = m * dims.K; + var acc = vec4(0.0); + for (var k = lane; k < dims.K; k = k + TK) { + acc = acc + f32(X[xbase + k]) * vec4(W[k * N4 + qc]); + } + red[tid] = acc; + workgroupBarrier(); + for (var s = TK / 2u; s > 0u; s = s >> 1u) { + if (lane < s) { red[tid] = red[tid] + red[tid + s * TN]; } + workgroupBarrier(); + } + if (lane == 0u && q < N4) { + var v = red[tid]; + let n0 = q * 4u; + {{IF_BIAS}} + v = v + vec4(f32(B[n0]), f32(B[n0 + 1u]), f32(B[n0 + 2u]), f32(B[n0 + 3u])); + {{/IF_BIAS}} + {{IF_SILU}}v = v / (vec4(1.0) + exp(-v));{{/IF_SILU}} + let ybase = m * dims.N + n0; + Y[ybase] = {{OUT_T}}(v.x); + Y[ybase + 1u] = {{OUT_T}}(v.y); + Y[ybase + 2u] = {{OUT_T}}(v.z); + Y[ybase + 3u] = {{OUT_T}}(v.w); + {{IF_STORE_KV}} + // Fused q|k|v row: quads in [HD, 2HD) are k, [2HD, 3HD) are v. A quad + // never straddles slices (HD % 4 == 0). Cache layout matches kv_append: + // dst = (b·Lmax + t)·HD + (h·D + d), with m = batch row. + let HD = dims.N / 3u; + if (n0 >= HD) { + var off = n0 - HD; + let dstBase = (m * dims.Lmax + dims.t) * HD; + if (off < HD) { + Kc[dstBase + off] = {{T}}(v.x); + Kc[dstBase + off + 1u] = {{T}}(v.y); + Kc[dstBase + off + 2u] = {{T}}(v.z); + Kc[dstBase + off + 3u] = {{T}}(v.w); + } else { + off = off - HD; + Vc[dstBase + off] = {{T}}(v.x); + Vc[dstBase + off + 1u] = {{T}}(v.y); + Vc[dstBase + off + 2u] = {{T}}(v.z); + Vc[dstBase + off + 3u] = {{T}}(v.w); + } + } + {{/IF_STORE_KV}} + } +} +{{/IF_NWT}} +`,jr=`// Tiled GEMM: Y[m,n] = Σ_k X[m,k] · W[k,n] (+ B[n]), optional SiLU — the +// large-M sibling of gemm.wgsl (encoder GEMMs, decode lm_head at big batch). +// X: [M,K] row-major, W: [K,N] row-major (WT flag: [N,K] row-major), Y: [M,N]. +// +// Each workgroup computes a BM×BN output tile: the K dimension is walked in +// BK-wide slices staged through workgroup memory (Xs/Ws, f32), and each of +// the (BM/4)·(BN/4) threads accumulates a FIXED 4×4 register subtile — 16 +// FMAs per 8 shared-memory reads instead of gemm.wgsl's 1 FMA per 2 global +// reads. f32 accumulation regardless of storage type. +// +// The 4×4 subtile is deliberately NOT templated: it lives in four vec4 +// accumulators with fully static indexing. A first version used +// array with loop indices — Tint/DXC kept it in scratch memory +// and the kernel came out SLOWER than the naive one (862ms vs 598ms encoder +// b64). Static vec4 registers are the whole point of the tile. +// +// Constraints: BM % 4 == 0, BN % 4 == 0, workgroup size = (BM/4)·(BN/4) +// (dispatchGemmTiled computes it), Xs+Ws = (BM+BN)·BK·4 bytes ≤ workgroup +// storage limit (16 KiB default). M/N/K tails are guarded: OOB loads stage +// 0, OOB stores are skipped. +// +// Template placeholders (see buildShader in pipelines.js): +// ENABLE_F16, T, OUT_T, WG as in gemm.wgsl (WG = thread count) +// BM, BN, BK tile geometry (u32 literals) +// IF_BIAS / IF_SILU / IF_WT conditional blocks +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Dims { M: u32, K: u32, N: u32, _pad: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> dims: Dims; +@group(0) @binding(1) var X: array<{{T}}>; +@group(0) @binding(2) var W: array<{{T}}>; +@group(0) @binding(3) var B: array<{{T}}>; +@group(0) @binding(4) var Y: array<{{OUT_T}}>; + +const BM = {{BM}}u; +const BN = {{BN}}u; +const BK = {{BK}}u; +const THREADS = (BM / 4u) * (BN / 4u); + +// Xs is stored TRANSPOSED — Xs[kk][row] — so the inner loop's 4-row read is +// one contiguous run; Ws[kk][col] likewise makes the 4-col read contiguous. +var Xs: array; +var Ws: array; + +fn storeRow(m: u32, n0: u32, v: vec4) { + if (m >= dims.M) { return; } + let base = m * dims.N; + if (n0 + 3u < dims.N) { + Y[base + n0] = {{OUT_T}}(v.x); + Y[base + n0 + 1u] = {{OUT_T}}(v.y); + Y[base + n0 + 2u] = {{OUT_T}}(v.z); + Y[base + n0 + 3u] = {{OUT_T}}(v.w); + return; + } + // N tail: per-lane guards. + if (n0 < dims.N) { Y[base + n0] = {{OUT_T}}(v.x); } + if (n0 + 1u < dims.N) { Y[base + n0 + 1u] = {{OUT_T}}(v.y); } + if (n0 + 2u < dims.N) { Y[base + n0 + 2u] = {{OUT_T}}(v.z); } +} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_index) lid: u32) { + // Keep B in the auto bind-group layout even when the bias block is + // compiled out — a 4-byte dummy buffer is bound in that case. + _ = &B; + let rowBase = wid.y * BM; + let colBase = wid.x * BN; + // This thread's 4×4 subtile origin inside the workgroup tile. + let tRow = (lid / (BN / 4u)) * 4u; + let tCol = (lid % (BN / 4u)) * 4u; + + var acc0 = vec4(0.0); // row tRow+0, cols tCol..tCol+3 + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + for (var k0: u32 = 0u; k0 < dims.K; k0 = k0 + BK) { + // Cooperative stage: Xs[BK][BM] (transposed) and Ws[BK][BN], linear + // thread sweep, zero-fill out-of-bounds lanes (M/K/N tails). + for (var i: u32 = lid; i < BK * BM; i = i + THREADS) { + let kk = i / BM; + let row = i % BM; + var v: f32 = 0.0; + if (rowBase + row < dims.M && k0 + kk < dims.K) { + v = f32(X[(rowBase + row) * dims.K + k0 + kk]); + } + Xs[i] = v; + } + for (var i: u32 = lid; i < BK * BN; i = i + THREADS) { + let kk = i / BN; + let col = i % BN; + var v: f32 = 0.0; + if (colBase + col < dims.N && k0 + kk < dims.K) { + var wIdx = (k0 + kk) * dims.N + colBase + col; // W [K,N] row-major + {{IF_WT}}wIdx = (colBase + col) * dims.K + k0 + kk; // W transposed: [N,K]{{/IF_WT}} + v = f32(W[wIdx]); + } + Ws[i] = v; + } + workgroupBarrier(); + + for (var kk: u32 = 0u; kk < BK; kk = kk + 1u) { + let xb = kk * BM + tRow; + let wb = kk * BN + tCol; + let xr = vec4(Xs[xb], Xs[xb + 1u], Xs[xb + 2u], Xs[xb + 3u]); + let wr = vec4(Ws[wb], Ws[wb + 1u], Ws[wb + 2u], Ws[wb + 3u]); + acc0 = fma(vec4(xr.x), wr, acc0); + acc1 = fma(vec4(xr.y), wr, acc1); + acc2 = fma(vec4(xr.z), wr, acc2); + acc3 = fma(vec4(xr.w), wr, acc3); + } + workgroupBarrier(); + } + + {{IF_BIAS}} + // min-clamped reads: OOB cols load a dummy lane that storeRow never writes. + let nHi = dims.N - 1u; + let bv = vec4( + f32(B[min(colBase + tCol, nHi)]), + f32(B[min(colBase + tCol + 1u, nHi)]), + f32(B[min(colBase + tCol + 2u, nHi)]), + f32(B[min(colBase + tCol + 3u, nHi)])); + acc0 = acc0 + bv; + acc1 = acc1 + bv; + acc2 = acc2 + bv; + acc3 = acc3 + bv; + {{/IF_BIAS}} + {{IF_SILU}} + acc0 = acc0 / (vec4(1.0) + exp(-acc0)); + acc1 = acc1 / (vec4(1.0) + exp(-acc1)); + acc2 = acc2 / (vec4(1.0) + exp(-acc2)); + acc3 = acc3 / (vec4(1.0) + exp(-acc3)); + {{/IF_SILU}} + + storeRow(rowBase + tRow, colBase + tCol, acc0); + storeRow(rowBase + tRow + 1u, colBase + tCol, acc1); + storeRow(rowBase + tRow + 2u, colBase + tCol, acc2); + storeRow(rowBase + tRow + 3u, colBase + tCol, acc3); +} +`,$r=`// Tiled GEMM v2 — gemm_tiled.wgsl with the memory system fixed: +// * X and W are bound as vec4 arrays: staging loads are one coalesced +// 8-byte (f16) vec4 per thread instead of v1's stride-K scalar sweeps. +// * Xs/Ws live as vec4 shared arrays: the inner loop issues 2 vec4 LDS +// reads per kk instead of v1's 8 scalar reads (v1 profiled ~7.9 TFLOPS +// on fc1 — LDS instruction bound, not FMA bound). +// * Optional 8×4 register subtile (IF_TM8): halves threads, doubles the +// FLOPs per Ws read. Static vec4 accumulators only — array accumulators +// spill under Tint/DXC (see gemm_tiled.wgsl header; that lesson cost a +// 1.4× regression before it was learned). +// +// Y[m,n] = Σ_k X[m,k]·W[k,n] (+B[n]), optional SiLU. X: [M,K] row-major, +// W: [K,N] row-major (WT: [N,K] row-major), Y: [M,N]. +// +// Eligibility (dispatchGemmTiled routes ineligible shapes to v1): +// K % 4 == 0 (vec4 X reads; also W reads when WT) +// N % 4 == 0 unless WT (vec4 W reads along N) +// BK % 4 == 0, BM % 4 == 0, BN % TN == 0 +// M tails ARE handled (OOB rows stage zero / skip store); K and N never +// straddle a vec4 thanks to the %4 constraints, so tail guards are per-vec4. +// +// W8A16 int8 mode (IF_WQ8 — lm_head, decode FFN): W is int8 [N,K] rows +// packed 4-per-u32 along K (quantizeQ8Rows layout); the per-N f32 scales get +// their own binding (5) so bias/SiLU stay available (FFN fc1). Raw q values +// are unpacked with sign-extending shifts and accumulated in f32 exactly +// like float weights; the per-column scale factors out of the dot product +// and multiplies the accumulators in the epilogue BEFORE bias — so +// dequantized results are bit-identical to a float GEMM over q·scale. +// Halves the W traffic of the f16 kernels. +// +// Fused argmax mode (IF_ARGMAX — the lm_head): the [M, 24000] f32 logits +// never touch global memory. Each thread applies final_logits_bias plus the +// repetition penalty to its register tile and keeps a per-(row, column-quad) +// running max; the workgroup folds its BN columns to ONE (val, idx) partial +// per row, and argmax_reduce.wgsl folds the ceil(N/BN) partials per row and +// runs argmax_penalty's token/done/bitmask epilogue. Same f32 ops in the same +// lexicographic (val, lowest idx) order as the unfused store+scan path, so +// the picked tokens are bit-identical — while the logits store, the argmax +// re-read, and the per-row bias re-read all disappear (~37MB/step at B=128). +// +// Template placeholders (see buildShader in pipelines.js): +// ENABLE_F16, T, OUT_T, WG as in gemm.wgsl (WG = (BM/TM)·(BN/4)) +// BM, BN, BK tile geometry (u32 literals) +// IF_TM8 8-row register subtile (else 4-row) +// IF_WQ8 / IF_WQF int8-packed W + scales / float W + bias +// IF_BIAS / IF_SILU / IF_WT as in gemm_tiled.wgsl +// IF_STORE_Y / IF_ARGMAX store Y (default) / fused-argmax epilogue +// PENALTY / MASK_WORDS argmax mode only: repetition penalty literal, +// bitmask u32 words per row +// IF_SH16 / IF_SH32 shared tiles store the native {{T}} / f32. +// SH16 halves LDS bank traffic; f16-origin +// values round-trip f32→f16 exactly so results +// are bit-identical. ARGMAX requires SH32 (the +// pVal alias needs full f32 lanes). +// IF_DBUF / IF_SBUF two tile buffers ping-ponged with ONE barrier +// per K-slice (stage k0+BK into the idle buffer +// while the live one computes) / classic +// stage-barrier-compute-barrier single buffer +// IF_SPLITK / KSL split-K for starved small-N decode sites +// (fc2 at B=128: 7×2 = 14 workgroups): grid.z +// partitions K into KSL-sized (BK-aligned) +// ranges, each workgroup stores its RAW f32 +// partial tile to PART [nz, M, N] — bias/SiLU +// move to gemm_reduce.wgsl, which folds the nz +// slices. Excludes WQ8/ARGMAX (guarded at +// dispatch; epilogues would double-apply). +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Dims { M: u32, K: u32, N: u32, _pad: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> dims: Dims; +@group(0) @binding(1) var X: array>; +{{IF_WQF}} +@group(0) @binding(2) var W: array>; +@group(0) @binding(3) var B: array<{{T}}>; +{{/IF_WQF}} +{{IF_WQ8}} +@group(0) @binding(2) var W: array; // 4×i8 along K +@group(0) @binding(3) var B: array<{{T}}>; +{{/IF_WQ8}} +{{IF_STORE_Y}} +@group(0) @binding(4) var Y: array<{{OUT_T}}>; +{{/IF_STORE_Y}} +{{IF_ARGMAX}} +// (val bitcast to u32, vocab idx) per [row, column-tile] — row-major [M, NT]. +@group(0) @binding(4) var P: array>; +{{/IF_ARGMAX}} +{{IF_SPLITK}} +// Raw f32 partials per K-partition — [nz, M, N]; gemm_reduce folds them. +@group(0) @binding(4) var PART: array; +{{/IF_SPLITK}} +{{IF_WQ8}} +@group(0) @binding(5) var S: array; // per-N scales +{{/IF_WQ8}} +{{IF_ARGMAX}} +{{IF_WQF}} +@group(0) @binding(5) var LB: array; // final_logits_bias [N] +@group(0) @binding(6) var SEEN: array; // rep bitmask [M·MASK_WORDS] +{{/IF_WQF}} +{{IF_WQ8}} +@group(0) @binding(6) var LB: array; // final_logits_bias [N] +@group(0) @binding(7) var SEEN: array; // rep bitmask [M·MASK_WORDS] +{{/IF_WQ8}} +{{/IF_ARGMAX}} + +const BM = {{BM}}u; +const BN = {{BN}}u; +const BK = {{BK}}u; +const TM = 4u{{IF_TM8}} * 2u{{/IF_TM8}}; // register subtile rows +const BM4 = BM / 4u; +const BN4 = BN / 4u; +const BK4 = BK / 4u; +const THREADS = (BM / TM) * BN4; + +// Same transposed layouts as v1 (Xs[kk][row], Ws[kk][col]), vec4-packed +// along the second axis so the inner loop reads whole quads. ST is the +// shared-lane type (see IF_SH16 header note); NBUF the tile buffer count. +{{IF_SH16}} +alias ST = {{T}}; +{{/IF_SH16}} +{{IF_SH32}} +alias ST = f32; +{{/IF_SH32}} +const NBUF = {{IF_DBUF}}2u{{/IF_DBUF}}{{IF_SBUF}}1u{{/IF_SBUF}}; +var Xs: array, NBUF * BK * BM4>; +var Ws: array, NBUF * BK * BN4>; + +{{IF_WQ8}} +// Sign-extend one packed u32 (4×i8 along K) to a vec4 of raw q values. +fn q8quad(word: i32) -> vec4 { + return vec4( + f32((word << 24u) >> 24u), + f32((word << 16u) >> 24u), + f32((word << 8u) >> 24u), + f32(word >> 24u)); +} +{{/IF_WQ8}} + +// Stage one BK-slice of X and W into tile buffer \`buf\`. +// +// PORTABILITY INVARIANT: every shared vec4 slot is written WHOLE by exactly +// one thread — each thread transposes a 4×4 register block (four vec4 global +// reads along K, four whole-vec4 shared stores). Per-component stores to one +// shared vec4 from different threads are legal WGSL (components are distinct +// memory locations) but miscompile on Metal into a whole-vector RMW: on an +// Apple M5 Max every v2 variant read back zeros/garbage (2026-07-06) until +// the staging was restructured. Do not reintroduce \`Xs[...][i] = scalar\` +// writes from multiple threads. +fn stageTiles(lid: u32, rowBase: u32, colBase: u32, k0: u32, buf: u32) { + let K4 = dims.K / 4u; // K % 4 == 0 enforced by dispatch + let xo = buf * (BK * BM4); + let wo = buf * (BK * BN4); + // Xs: 4-row × 4-k register transpose into the Xs[kk][row-quad] layout. + // i → (rq, j) over BM4 × BK4 blocks; reads stay coalesced along K. + for (var i: u32 = lid; i < BM4 * BK4; i = i + THREADS) { + let rq = i / BK4; + let j = i % BK4; + let r0 = rowBase + 4u * rq; + var v0 = vec4(0.0); + var v1 = vec4(0.0); + var v2 = vec4(0.0); + var v3 = vec4(0.0); + if (k0 + 4u * j < dims.K) { + if (r0 < dims.M) { v0 = vec4(X[r0 * K4 + k0 / 4u + j]); } + if (r0 + 1u < dims.M) { v1 = vec4(X[(r0 + 1u) * K4 + k0 / 4u + j]); } + if (r0 + 2u < dims.M) { v2 = vec4(X[(r0 + 2u) * K4 + k0 / 4u + j]); } + if (r0 + 3u < dims.M) { v3 = vec4(X[(r0 + 3u) * K4 + k0 / 4u + j]); } + } + let kk = 4u * j; + Xs[xo + kk * BM4 + rq] = vec4(vec4(v0.x, v1.x, v2.x, v3.x)); + Xs[xo + (kk + 1u) * BM4 + rq] = vec4(vec4(v0.y, v1.y, v2.y, v3.y)); + Xs[xo + (kk + 2u) * BM4 + rq] = vec4(vec4(v0.z, v1.z, v2.z, v3.z)); + Xs[xo + (kk + 3u) * BM4 + rq] = vec4(vec4(v0.w, v1.w, v2.w, v3.w)); + } + {{IF_WT}} + // W [N,K] row-major: same 4-col × 4-k register transpose as Xs. + for (var i: u32 = lid; i < BN4 * BK4; i = i + THREADS) { + let cq = i / BK4; + let j = i % BK4; + let c0 = colBase + 4u * cq; + var v0 = vec4(0.0); + var v1 = vec4(0.0); + var v2 = vec4(0.0); + var v3 = vec4(0.0); + if (k0 + 4u * j < dims.K) { + if (c0 < dims.N) { v0 = vec4(W[c0 * K4 + k0 / 4u + j]); } + if (c0 + 1u < dims.N) { v1 = vec4(W[(c0 + 1u) * K4 + k0 / 4u + j]); } + if (c0 + 2u < dims.N) { v2 = vec4(W[(c0 + 2u) * K4 + k0 / 4u + j]); } + if (c0 + 3u < dims.N) { v3 = vec4(W[(c0 + 3u) * K4 + k0 / 4u + j]); } + } + let kk = 4u * j; + Ws[wo + kk * BN4 + cq] = vec4(vec4(v0.x, v1.x, v2.x, v3.x)); + Ws[wo + (kk + 1u) * BN4 + cq] = vec4(vec4(v0.y, v1.y, v2.y, v3.y)); + Ws[wo + (kk + 2u) * BN4 + cq] = vec4(vec4(v0.z, v1.z, v2.z, v3.z)); + Ws[wo + (kk + 3u) * BN4 + cq] = vec4(vec4(v0.w, v1.w, v2.w, v3.w)); + } + {{/IF_WT}} + {{IF_WNT}} + // W [K,N] row-major: vec4 reads along N land directly on one Ws quad + // (N % 4 == 0 enforced, so quads never straddle the boundary). + let N4 = dims.N / 4u; + for (var i: u32 = lid; i < BK * BN4; i = i + THREADS) { + let kk = i / BN4; + let c = i % BN4; + var v = vec4(0.0); + if (colBase + 4u * c < dims.N && k0 + kk < dims.K) { + v = vec4(W[(k0 + kk) * N4 + colBase / 4u + c]); + } + Ws[wo + i] = vec4(v); + } + {{/IF_WNT}} + {{IF_WQ8}} + // W int8 [N,K] packed: one u32 = 4 raw q values along K, sign-extended + // by shifts and staged UNSCALED (the per-N scale is an epilogue factor). + // Same 4-col × 4-k register transpose as the WT path. + for (var i: u32 = lid; i < BN4 * BK4; i = i + THREADS) { + let cq = i / BK4; + let j = i % BK4; + let c0 = colBase + 4u * cq; + var v0 = vec4(0.0); + var v1 = vec4(0.0); + var v2 = vec4(0.0); + var v3 = vec4(0.0); + if (k0 + 4u * j < dims.K) { + if (c0 < dims.N) { v0 = q8quad(bitcast(W[c0 * K4 + k0 / 4u + j])); } + if (c0 + 1u < dims.N) { v1 = q8quad(bitcast(W[(c0 + 1u) * K4 + k0 / 4u + j])); } + if (c0 + 2u < dims.N) { v2 = q8quad(bitcast(W[(c0 + 2u) * K4 + k0 / 4u + j])); } + if (c0 + 3u < dims.N) { v3 = q8quad(bitcast(W[(c0 + 3u) * K4 + k0 / 4u + j])); } + } + let kk = 4u * j; + Ws[wo + kk * BN4 + cq] = vec4(vec4(v0.x, v1.x, v2.x, v3.x)); + Ws[wo + (kk + 1u) * BN4 + cq] = vec4(vec4(v0.y, v1.y, v2.y, v3.y)); + Ws[wo + (kk + 2u) * BN4 + cq] = vec4(vec4(v0.z, v1.z, v2.z, v3.z)); + Ws[wo + (kk + 3u) * BN4 + cq] = vec4(vec4(v0.w, v1.w, v2.w, v3.w)); + } + {{/IF_WQ8}} +} + +{{IF_STORE_Y}} +fn storeRow(m: u32, n0: u32, v: vec4) { + if (m >= dims.M) { return; } + let base = m * dims.N; + if (n0 + 3u < dims.N) { + Y[base + n0] = {{OUT_T}}(v.x); + Y[base + n0 + 1u] = {{OUT_T}}(v.y); + Y[base + n0 + 2u] = {{OUT_T}}(v.z); + Y[base + n0 + 3u] = {{OUT_T}}(v.w); + return; + } + // N tail (v1-fallback shapes never reach here, but WT allows N % 4 != 0). + if (n0 < dims.N) { Y[base + n0] = {{OUT_T}}(v.x); } + if (n0 + 1u < dims.N) { Y[base + n0 + 1u] = {{OUT_T}}(v.y); } + if (n0 + 2u < dims.N) { Y[base + n0 + 2u] = {{OUT_T}}(v.z); } +} +{{/IF_STORE_Y}} +{{IF_SPLITK}} +const KSL = {{KSL}}u; // BK-aligned K range per grid.z partition +fn storePart(z: u32, m: u32, n0: u32, v: vec4) { + if (m >= dims.M) { return; } + let base = (z * dims.M + m) * dims.N + n0; + if (n0 + 3u < dims.N) { + PART[base] = v.x; + PART[base + 1u] = v.y; + PART[base + 2u] = v.z; + PART[base + 3u] = v.w; + return; + } + if (n0 < dims.N) { PART[base] = v.x; } + if (n0 + 1u < dims.N) { PART[base + 1u] = v.y; } + if (n0 + 2u < dims.N) { PART[base + 2u] = v.z; } +} +{{/IF_SPLITK}} +{{IF_ARGMAX}} +const NEG_MAX: f32 = -3.40282e38; // finite f32 lowest (WGSL has no inf literal) +const PENALTY: f32 = {{PENALTY}}; +const MASK_WORDS: u32 = {{MASK_WORDS}}u; +// Per-row partials for the cross-thread fold: pVal aliases Xs, which is dead +// after the K loop (f32 values through f32 lanes are bit-preserving; ARGMAX +// requires SH32 so the lanes are full f32). pIdx must NOT alias Ws — small +// u32 indices are f32 denormal bit patterns, and a shared-memory round trip +// through f32 may flush them to zero. +// pVal slots are laid out [col-quad][row] (slot = c4·BM + row) so one +// thread's TM row-partials land in the SAME Xs vec4s and are stored whole — +// the staging portability invariant applies to this reuse too. +const_assert(BK * BM >= BM * BN4); // Xs lane capacity covers the BM·BN4 slots +var pIdx: array; + +fn penal(x: f32, seen: bool) -> f32 { + if (!seen) { return x; } + if (x > 0.0) { return x / PENALTY; } + return x * PENALTY; +} + +// One register row: bias + penalty + 4-lane max. Ascending lane order with +// strict > keeps the lowest-index tie-break (torch.argmax first-max) exact. +// A quad never straddles a bitmask word: n0 % 4 == 0 so (n0 & 31) <= 28. +// Writes pIdx (scalar array — per-slot stores are race-free) and RETURNS the +// partial value; the caller packs TM of them into whole Xs vec4 stores. +fn fusedPartial(m: u32, slot: u32, n0: u32, acc: vec4, lbq: vec4) -> f32 { + var val: f32 = NEG_MAX; + var idx: u32 = 0xffffffffu; + if (m < dims.M && n0 < dims.N) { + let x = acc + lbq; + let bits = SEEN[m * MASK_WORDS + (n0 >> 5u)] >> (n0 & 31u); + var v1 = NEG_MAX; + var v2 = NEG_MAX; + var v3 = NEG_MAX; + if (n0 + 1u < dims.N) { v1 = penal(x.y, ((bits >> 1u) & 1u) == 1u); } + if (n0 + 2u < dims.N) { v2 = penal(x.z, ((bits >> 2u) & 1u) == 1u); } + if (n0 + 3u < dims.N) { v3 = penal(x.w, ((bits >> 3u) & 1u) == 1u); } + val = penal(x.x, (bits & 1u) == 1u); + idx = n0; + if (v1 > val) { val = v1; idx = n0 + 1u; } + if (v2 > val) { val = v2; idx = n0 + 2u; } + if (v3 > val) { val = v3; idx = n0 + 3u; } + } + pIdx[slot] = idx; + return val; +} +{{/IF_ARGMAX}} + +fn biasQuad(n0: u32) -> vec4 { + // min-clamped reads: OOB cols load a dummy lane that storeRow never writes. + let nHi = dims.N - 1u; + return vec4( + f32(B[min(n0, nHi)]), + f32(B[min(n0 + 1u, nHi)]), + f32(B[min(n0 + 2u, nHi)]), + f32(B[min(n0 + 3u, nHi)])); +} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_index) lid: u32) { + _ = &B; // keep B in the auto layout when the bias block is compiled out + let rowBase = wid.y * BM; + let colBase = wid.x * BN; + let tRow = (lid / BN4) * TM; + let tCol = (lid % BN4) * 4u; + let r4 = tRow / 4u; + let c4 = tCol / 4u; + + var acc0 = vec4(0.0); // row tRow+0, cols tCol..tCol+3 + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + {{IF_TM8}} + var acc4 = vec4(0.0); + var acc5 = vec4(0.0); + var acc6 = vec4(0.0); + var acc7 = vec4(0.0); + {{/IF_TM8}} + + {{IF_SPLITK}} + let kBeg = wid.z * KSL; + let kEnd = min(kBeg + KSL, dims.K); + {{/IF_SPLITK}} + {{IF_NOSPLITK}} + let kBeg = 0u; + let kEnd = dims.K; + {{/IF_NOSPLITK}} + {{IF_DBUF}} + stageTiles(lid, rowBase, colBase, kBeg, 0u); + workgroupBarrier(); + {{/IF_DBUF}} + var buf = 0u; + for (var k0: u32 = kBeg; k0 < kEnd; k0 = k0 + BK) { + {{IF_DBUF}} + // Prefetch the NEXT slice into the idle buffer while this one computes. + // The single trailing barrier both publishes this prefetch for the next + // iteration and retires the live buffer's reads before iteration + // k0 + 2·BK overwrites it — safe with one barrier per slice. + if (k0 + BK < kEnd) { stageTiles(lid, rowBase, colBase, k0 + BK, 1u - buf); } + {{/IF_DBUF}} + {{IF_SBUF}} + stageTiles(lid, rowBase, colBase, k0, 0u); + workgroupBarrier(); + {{/IF_SBUF}} + let xo = buf * (BK * BM4); + let wo = buf * (BK * BN4); + for (var kk: u32 = 0u; kk < BK; kk = kk + 1u) { + let wr = vec4(Ws[wo + kk * BN4 + c4]); + let xr = vec4(Xs[xo + kk * BM4 + r4]); + acc0 = fma(vec4(xr.x), wr, acc0); + acc1 = fma(vec4(xr.y), wr, acc1); + acc2 = fma(vec4(xr.z), wr, acc2); + acc3 = fma(vec4(xr.w), wr, acc3); + {{IF_TM8}} + let xr2 = vec4(Xs[xo + kk * BM4 + r4 + 1u]); + acc4 = fma(vec4(xr2.x), wr, acc4); + acc5 = fma(vec4(xr2.y), wr, acc5); + acc6 = fma(vec4(xr2.z), wr, acc6); + acc7 = fma(vec4(xr2.w), wr, acc7); + {{/IF_TM8}} + } + workgroupBarrier(); + {{IF_DBUF}} + buf = 1u - buf; + {{/IF_DBUF}} + } + + {{IF_WQ8}} + // Dequantize BEFORE bias: the per-N scale multiplies the whole column. + let sHi = dims.N - 1u; + let sn0 = colBase + tCol; + let sv = vec4( + S[min(sn0, sHi)], S[min(sn0 + 1u, sHi)], S[min(sn0 + 2u, sHi)], S[min(sn0 + 3u, sHi)]); + acc0 = acc0 * sv; + acc1 = acc1 * sv; + acc2 = acc2 * sv; + acc3 = acc3 * sv; + {{IF_TM8}} + acc4 = acc4 * sv; + acc5 = acc5 * sv; + acc6 = acc6 * sv; + acc7 = acc7 * sv; + {{/IF_TM8}} + {{/IF_WQ8}} + {{IF_BIAS}} + let bv = biasQuad(colBase + tCol); + acc0 = acc0 + bv; + acc1 = acc1 + bv; + acc2 = acc2 + bv; + acc3 = acc3 + bv; + {{IF_TM8}} + acc4 = acc4 + bv; + acc5 = acc5 + bv; + acc6 = acc6 + bv; + acc7 = acc7 + bv; + {{/IF_TM8}} + {{/IF_BIAS}} + {{IF_SILU}} + acc0 = acc0 / (vec4(1.0) + exp(-acc0)); + acc1 = acc1 / (vec4(1.0) + exp(-acc1)); + acc2 = acc2 / (vec4(1.0) + exp(-acc2)); + acc3 = acc3 / (vec4(1.0) + exp(-acc3)); + {{IF_TM8}} + acc4 = acc4 / (vec4(1.0) + exp(-acc4)); + acc5 = acc5 / (vec4(1.0) + exp(-acc5)); + acc6 = acc6 / (vec4(1.0) + exp(-acc6)); + acc7 = acc7 / (vec4(1.0) + exp(-acc7)); + {{/IF_TM8}} + {{/IF_SILU}} + + {{IF_STORE_Y}} + storeRow(rowBase + tRow, colBase + tCol, acc0); + storeRow(rowBase + tRow + 1u, colBase + tCol, acc1); + storeRow(rowBase + tRow + 2u, colBase + tCol, acc2); + storeRow(rowBase + tRow + 3u, colBase + tCol, acc3); + {{IF_TM8}} + storeRow(rowBase + tRow + 4u, colBase + tCol, acc4); + storeRow(rowBase + tRow + 5u, colBase + tCol, acc5); + storeRow(rowBase + tRow + 6u, colBase + tCol, acc6); + storeRow(rowBase + tRow + 7u, colBase + tCol, acc7); + {{/IF_TM8}} + {{/IF_STORE_Y}} + {{IF_SPLITK}} + storePart(wid.z, rowBase + tRow, colBase + tCol, acc0); + storePart(wid.z, rowBase + tRow + 1u, colBase + tCol, acc1); + storePart(wid.z, rowBase + tRow + 2u, colBase + tCol, acc2); + storePart(wid.z, rowBase + tRow + 3u, colBase + tCol, acc3); + {{IF_TM8}} + storePart(wid.z, rowBase + tRow + 4u, colBase + tCol, acc4); + storePart(wid.z, rowBase + tRow + 5u, colBase + tCol, acc5); + storePart(wid.z, rowBase + tRow + 6u, colBase + tCol, acc6); + storePart(wid.z, rowBase + tRow + 7u, colBase + tCol, acc7); + {{/IF_TM8}} + {{/IF_SPLITK}} + {{IF_ARGMAX}} + // The K loop's trailing barrier already ordered every thread's Xs reads + // before this reuse — fusedPartial may write pVal into Xs immediately. + // Slot layout [col-quad][row]: this thread's TM partials are consecutive + // rows of ONE column-quad, so they pack into whole Xs vec4 stores + // (tRow % 4 == 0 by construction — the portability invariant again). + let nHiL = dims.N - 1u; + let n0 = colBase + tCol; + let lbq = vec4( + LB[min(n0, nHiL)], LB[min(n0 + 1u, nHiL)], LB[min(n0 + 2u, nHiL)], LB[min(n0 + 3u, nHiL)]); + let pBase = c4 * BM + tRow; + let p0 = fusedPartial(rowBase + tRow, pBase, n0, acc0, lbq); + let p1 = fusedPartial(rowBase + tRow + 1u, pBase + 1u, n0, acc1, lbq); + let p2 = fusedPartial(rowBase + tRow + 2u, pBase + 2u, n0, acc2, lbq); + let p3 = fusedPartial(rowBase + tRow + 3u, pBase + 3u, n0, acc3, lbq); + Xs[pBase / 4u] = vec4(vec4(p0, p1, p2, p3)); + {{IF_TM8}} + let p4 = fusedPartial(rowBase + tRow + 4u, pBase + 4u, n0, acc4, lbq); + let p5 = fusedPartial(rowBase + tRow + 5u, pBase + 5u, n0, acc5, lbq); + let p6 = fusedPartial(rowBase + tRow + 6u, pBase + 6u, n0, acc6, lbq); + let p7 = fusedPartial(rowBase + tRow + 7u, pBase + 7u, n0, acc7, lbq); + Xs[pBase / 4u + 1u] = vec4(vec4(p4, p5, p6, p7)); + {{/IF_TM8}} + workgroupBarrier(); + + // Fold each row's BN4 column-quads (ascending c = ascending vocab ids, so + // the explicit lower-idx tie clause and scan order both preserve first-max) + // and emit the workgroup's one partial per row. + let NT = (dims.N + BN - 1u) / BN; + for (var r = lid; r < BM; r = r + THREADS) { + var val: f32 = NEG_MAX; + var idx: u32 = 0xffffffffu; + for (var c = 0u; c < BN4; c = c + 1u) { + let s = c * BM + r; + let v = f32(Xs[s / 4u][s % 4u]); + let i = pIdx[s]; + if (v > val || (v == val && i < idx)) { val = v; idx = i; } + } + if (rowBase + r < dims.M) { + P[(rowBase + r) * NT + wid.x] = vec2(bitcast(val), idx); + } + } + {{/IF_ARGMAX}} +} +`,zr=`// Unified scaled-dot-product attention — the ONE kernel for encoder self, +// decoder self, and cross attention (no transposes anywhere in the engine). +// +// Q: element (row, h, d) at row*Q_STRIDE + Q_OFF + h*D + d, row = b*M + m. +// With the defaults (Q_STRIDE = H·D, Q_OFF = 0) that is a plain +// [B·M, H·D] row-major buffer. +// K, V: element (b, j, h, d) at (b*L + j)*KV_STRIDE + K_OFF|V_OFF + h*D + d. +// Defaults (KV_STRIDE = H·D, offsets 0) reproduce the [B, L, H, D] +// cache layout: ((b*L + j)*H + h)*D + d. Non-default strides/offsets let +// the kernel read q/k/v slices DIRECTLY from fused projection outputs +// (encoder self: Q=K=V= fused qkv [B·S, 3·H·D] with stride 3·H·D and +// offsets 0/H·D/2·H·D; cross: K=V= fused kv [B·S, 2·H·D]). K and V may +// be bound to the same buffer (both read-only storage). +// Y: always compact [B·M, H·D] row-major, regardless of Q layout. +// L is the K/V buffer's position CAPACITY (stride) — the valid length +// is selected by lenMode: 0 → t+1 (decoder self, uniform across +// batch), 1 → lens[b] (encoder self / cross, right-padded batches). +// +// Q/K/V are bound as vec4 arrays (Task 17: 4× fewer loads in the phase-1 dot +// and phase-3 accumulation): D and every stride/offset (Q_STRIDE, Q_OFF, +// KV_STRIDE, K_OFF, V_OFF — and h·D by D%4) must be multiples of 4; +// dispatchAttention enforces this. Y stays a scalar array. +// +// Grid: dispatchWorkgroups(B·M, H) — wid.x = query row, wid.y = head. One +// workgroup computes one (row, head) output vector of D elements: +// phase 0: stage q into shared qs4 (D4 = D/4 vec4s) +// phase 1: scores_j = ATTN_SCALE·Σ_d qs·K, tree-reduce MAX (threads stride j) +// phase 2: exponentiate in place, tree-reduce SUM → denom +// phase 3: out_dq = Σ_j scores_j·V: thread (dq, jg) owns d-quad dq and sums +// j ≡ jg (mod JT) into a shared partial; the first D4 threads fold +// the JT partials and store. (Previously one thread per d with a +// SERIAL j loop — the profiled decode hotspot of this kernel.) +// All math in f32 regardless of storage type; exp after max-subtraction; +// single division at the end. +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 the f16 enable directive when T is f16, else empty +// T storage type of Q/K/V/Y (f16|f32) +// WG workgroup size (128; must be a power of two, >= D) +// H heads (8) +// D head dim (56) +// SCORES_CAP shared scores capacity (352 — max valid length) +// ATTN_SCALE 1/sqrt(D) as a full-precision literal +// Q_STRIDE per-row element stride of Q (default H·D = 448) +// Q_OFF element offset of the q slice within a Q row (default 0) +// KV_STRIDE per-position element stride of K/V (default H·D = 448) +// K_OFF element offset of the k slice within a K position (default 0) +// V_OFF element offset of the v slice within a V position (default 0) +// IF_SG/IF_NOSG subgroup reduction variant (flags.sg): subgroupMax/ +// subgroupAdd + elect + serial fold over per-subgroup partials +// in two SEPARATE arrays (no red[] reuse). Besides the barrier +// savings, this is the CORRECTNESS route on Adreno 7xx: its +// driver miscompiles this kernel's tree-reduce idiom (wave-1 +// re-writes of red[] read as stale phase-1 values across a +// barrier — 2026-07 Android probe rounds; the naked idiom in +// isolation passes, so the trigger is contextual). Reduction +// order differs from the tree → tolerance-equal, not bit-equal. +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} +{{ENABLE_SG}} + +struct Params { + B: u32, // batch + M: u32, // query rows per batch element (encoder S, else 1) + L: u32, // K/V position capacity (stride) + lenMode: u32, // 0: len = t+1, 1: len = lens[b] + t: u32, // decode step (lenMode 0 only) + _pad0: u32, _pad1: u32, _pad2: u32, +} + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var Q: array>; +@group(0) @binding(2) var K: array>; +@group(0) @binding(3) var V: array>; +// Always bound — a 4-byte dummy when lenMode == 0 (never indexed then). +@group(0) @binding(4) var lens: array; +@group(0) @binding(5) var Y: array<{{T}}>; + +const H: u32 = {{H}}u; +const D: u32 = {{D}}u; +const WG: u32 = {{WG}}u; +const SCORES_CAP: u32 = {{SCORES_CAP}}u; +const ATTN_SCALE: f32 = {{ATTN_SCALE}}; +const Q_STRIDE: u32 = {{Q_STRIDE}}u; +const Q_OFF: u32 = {{Q_OFF}}u; +const KV_STRIDE: u32 = {{KV_STRIDE}}u; +const K_OFF: u32 = {{K_OFF}}u; +const V_OFF: u32 = {{V_OFF}}u; +const D4: u32 = D / 4u; // d-quads per head (D%4 == 0 enforced) +const JT: u32 = WG / D4; // phase-3 j-lanes per d-quad + +var qs4: array, D4>; // staged query vector +var scores: array; // raw scores, then exp values +{{IF_NOSG}} +var red: array; // tree-reduction scratch +{{/IF_NOSG}} +{{IF_SG}} +// One slot per subgroup, separate arrays per reduction (max, then sum) so +// the sum pass never overwrites slots the max readers still need. WG/4 +// covers the spec-minimum subgroup size of 4. +var sgMax: array; +var sgSum: array; +{{/IF_SG}} +var part: array, WG>; // phase-3 partial V sums + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3{{IF_SG}}, @builtin(subgroup_size) sgSize: u32{{/IF_SG}}) { + // Uniform per workgroup (whole workgroup shares wid) — safe early return + // before the first barrier. + if (wid.x >= params.B * params.M) { return; } + let row = wid.x; + let h = wid.y; + let b = row / params.M; + let tid = lid.x; + + var len: u32; + if (params.lenMode == 0u) { len = params.t + 1u; } else { len = lens[b]; } + // Engine guarantees len <= SCORES_CAP; clamp anyway — cheap insurance + // against out-of-bounds shared writes on a bad Params. + len = min(len, SCORES_CAP); + + // Phase 0: stage q as D4 vec4s (D4=14 < WG → one quad per thread). + let qoff4 = (row * Q_STRIDE + Q_OFF + h * D) / 4u; + for (var i = tid; i < D4; i = i + WG) { + qs4[i] = vec4(Q[qoff4 + i]); + } + workgroupBarrier(); // qs4 must be visible before phase-1 dot products + + // Phase 1: scores + running max. Threads stride j; idle threads (j >= len) + // still contribute the neutral element -1e30 to the max reduction. + var localMax: f32 = -1e30; + for (var j = tid; j < len; j = j + WG) { + let koff4 = ((b * params.L + j) * KV_STRIDE + K_OFF + h * D) / 4u; + var dot4 = vec4(0.0); + for (var i = 0u; i < D4; i = i + 1u) { + dot4 = dot4 + qs4[i] * vec4(K[koff4 + i]); + } + let sc = (dot4.x + dot4.y + dot4.z + dot4.w) * ATTN_SCALE; + scores[j] = sc; + localMax = max(localMax, sc); + } +{{IF_NOSG}} + red[tid] = localMax; + workgroupBarrier(); // red fully written (and scores writes made visible) + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { red[tid] = max(red[tid], red[tid + s]); } + workgroupBarrier(); + } + let rowMax = red[0]; + workgroupBarrier(); // all reads of red[0] done before phase 2 reuses red +{{/IF_NOSG}} +{{IF_SG}} + // sgId assumes linear subgroup layout (tid/sgSize) — holds on + // Metal/D3D12/Vulkan under Tint; the sg equiv gates catch a device where + // it doesn't. subgroup ops sit in uniform control flow — outside the elect. + let sgId = tid / sgSize; + let nSg = (WG + sgSize - 1u) / sgSize; + let m1 = subgroupMax(localMax); + if (subgroupElect()) { sgMax[sgId] = m1; } + workgroupBarrier(); // sgMax written (and scores writes made visible) + var rowMax = sgMax[0]; + for (var i = 1u; i < nSg; i = i + 1u) { rowMax = max(rowMax, sgMax[i]); } +{{/IF_SG}} + + // Phase 2: exponentiate in place + sum. Idle threads contribute 0. + var localSum: f32 = 0.0; + for (var j = tid; j < len; j = j + WG) { + let e = exp(scores[j] - rowMax); + scores[j] = e; + localSum = localSum + e; + } +{{IF_NOSG}} + red[tid] = localSum; + workgroupBarrier(); // red fully written (and exp'd scores made visible) + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { red[tid] = red[tid] + red[tid + s]; } + workgroupBarrier(); + } + let denom = red[0]; + // No further barrier: red is not reused, and the exp'd scores reads below + // were ordered by the reduction barriers above. +{{/IF_NOSG}} +{{IF_SG}} + let s2 = subgroupAdd(localSum); + if (subgroupElect()) { sgSum[sgId] = s2; } + workgroupBarrier(); // sgSum written (and exp'd scores made visible) + var denom = sgSum[0]; + for (var i = 1u; i < nSg; i = i + 1u) { denom = denom + sgSum[i]; } +{{/IF_SG}} + + // Phase 3: thread (dq = tid%D4, jg = tid/D4) accumulates d-quad dq over + // j ≡ jg (mod JT); the first D4 threads fold the JT partials and store. + // Y is always compact [B·M, H·D] even when Q is strided into a fused buffer. + let dq = tid % D4; + let jg = tid / D4; + var acc = vec4(0.0); + if (jg < JT) { + for (var j = jg; j < len; j = j + JT) { + let voff4 = ((b * params.L + j) * KV_STRIDE + V_OFF + h * D) / 4u; + acc = acc + scores[j] * vec4(V[voff4 + dq]); + } + } + part[tid] = acc; // threads with jg >= JT park a zero in an unread slot + workgroupBarrier(); + if (tid < D4) { + var out = vec4(0.0); + for (var g = 0u; g < JT; g = g + 1u) { out = out + part[g * D4 + tid]; } + out = out / denom; + let yoff = row * H * D + h * D + tid * 4u; + Y[yoff] = {{T}}(out.x); + Y[yoff + 1u] = {{T}}(out.y); + Y[yoff + 2u] = {{T}}(out.z); + Y[yoff + 3u] = {{T}}(out.w); + } +} +`,Vr=`// Blocked (flash-style) scaled-dot-product attention — ENCODER SELF/FUSED +// ONLY (lenMode 1, ragged lens; decode keeps attention.wgsl, whose M=1 gains +// nothing from query blocking). +// +// Why it exists: attention.wgsl runs one workgroup per (query row, head), so +// K/V of a (batch, head) pair are re-read from global memory once per query +// row — S× redundant traffic at encoder shapes (enc_profile after the tiled +// GEMM round: attention ≈ 50% of encoder time). Here one workgroup owns QB +// query rows of one (b, h): K/V tiles of JB positions are staged in shared +// memory ONCE and reused by all QB queries (global K/V traffic ÷ QB), with a +// streaming online softmax so no full score row is ever materialized. +// +// Layouts match attention.wgsl exactly (fused-QKV strides/offsets, vec4 +// bindings, compact Y): see that header. Grid is (ceil(M/QB), H, B). +// +// Threads: WG = QB·D4 — thread (q = tid/D4, dq = tid%D4) owns query q's +// d-quad dq. Its V-sum lives in ONE static vec4 register accumulator +// (array accumulators spill under Tint/DXC — see gemm_tiled.wgsl header). +// Per j-tile: +// stage Ks/Vs[JB·D4] cooperatively, native {{T}} (no precision change: +// attention.wgsl converts the same f16 values at load) +// scores QB·jn dot products over the WG threads → p[] (raw, f32) +// update one thread per query: running max m, rescale = exp(m_old − m), +// exp scores in place, running sum +// accum acc = acc·rescale[q] + Σ_j p[q][j]·Vs[j][dq] +// Store: acc / rowSum, guarded for the M tail block. All math f32. +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 f16 enable directive when T is f16, else empty +// T storage type of Q/K/V/Y (f16|f32) +// QB query rows per workgroup (workgroup size = QB·D/4) +// JB K/V tile positions staged per iteration +// H, D, ATTN_SCALE, Q_STRIDE, Q_OFF, KV_STRIDE, K_OFF, V_OFF +// as in attention.wgsl (D%4 == 0 enforced by dispatch) +// IF_PACKED / IF_NOPACKED +// encoder row-packing: with PACKED, Q/K/V/Y hold T = Σ lens +// rows (no pad rows) and sequence b occupies rows +// starts[b] .. starts[b]+len — binding 6 carries starts. +// Query blocks entirely past len early-exit (uniform: len and +// starts are read-only storage loads at workgroup-uniform +// indices). NOPACKED keeps the padded b·M / b·L bases. +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Params { + B: u32, // batch + M: u32, // query rows per batch element (encoder S) + L: u32, // K/V position capacity (stride) + lenMode: u32, // must be 1 (dispatch enforces); len = lens[b] + t: u32, // unused here + _pad0: u32, _pad1: u32, _pad2: u32, +} + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var Q: array>; +@group(0) @binding(2) var K: array>; +@group(0) @binding(3) var V: array>; +@group(0) @binding(4) var lens: array; +@group(0) @binding(5) var Y: array<{{T}}>; +{{IF_PACKED}} +@group(0) @binding(6) var starts: array; // packed row offsets [B] +{{/IF_PACKED}} + +const H: u32 = {{H}}u; +const D: u32 = {{D}}u; +const QB: u32 = {{QB}}u; +const JB: u32 = {{JB}}u; +const ATTN_SCALE: f32 = {{ATTN_SCALE}}; +const Q_STRIDE: u32 = {{Q_STRIDE}}u; +const Q_OFF: u32 = {{Q_OFF}}u; +const KV_STRIDE: u32 = {{KV_STRIDE}}u; +const K_OFF: u32 = {{K_OFF}}u; +const V_OFF: u32 = {{V_OFF}}u; +const D4: u32 = D / 4u; +const WG: u32 = QB * D4; + +var Qs: array, QB * D4>; // staged query block +var Ks: array, JB * D4>; // staged K tile (native T) +var Vs: array, JB * D4>; // staged V tile (native T) +var p: array; // tile scores, then exp values +var rowMax: array; // running max per query +var rowSum: array; // running exp-sum per query +var rescale: array; // exp(oldMax - newMax) per tile + +@compute @workgroup_size(WG) // const-expr: QB·D4 threads +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + let qb0 = wid.x * QB; // first query row of this block (within the batch elt) + let h = wid.y; + let b = wid.z; + let tid = lid.x; + let q = tid / D4; // this thread's query within the block + let dq = tid % D4; // this thread's d-quad + + // Engine guarantees len <= L; clamp anyway (bad Params must not overrun L). + let len = min(lens[b], params.L); +{{IF_PACKED}} + // Nothing to store for a query block entirely past len — uniform early exit + // (before any barrier; qb0/len are workgroup-uniform). + if (qb0 >= len) { return; } + let qBase = starts[b]; + let kvBase = starts[b]; + let mEnd = len; // valid query rows of this sequence +{{/IF_PACKED}} +{{IF_NOPACKED}} + let qBase = b * params.M; + let kvBase = b * params.L; + let mEnd = params.M; +{{/IF_NOPACKED}} + + // Stage the query block: one owned quad per thread. Tail-block rows past + // mEnd stage a duplicate of the last valid row — computed but never stored. + let qRow = min(qb0 + q, mEnd - 1u); + Qs[tid] = vec4(Q[((qBase + qRow) * Q_STRIDE + Q_OFF + h * D) / 4u + dq]); + if (dq == 0u) { + rowMax[q] = -1e30; + rowSum[q] = 0.0; + } + var acc = vec4(0.0); + workgroupBarrier(); // Qs/rowMax/rowSum visible before the first tile + + // len is workgroup-uniform (lens[b], b from workgroup_id) — barriers inside + // the loop are in uniform control flow. + for (var j0 = 0u; j0 < len; j0 = j0 + JB) { + let jn = min(JB, len - j0); + + // Stage K/V tiles: JB·D4 quads each, cooperative. Slots past jn stage a + // duplicate of the last valid position — never read below. + for (var i = tid; i < JB * D4; i = i + WG) { + let jj = i / D4; + let dd = i % D4; + let base = (kvBase + j0 + min(jj, jn - 1u)) * KV_STRIDE + h * D; + Ks[i] = K[(base + K_OFF) / 4u + dd]; + Vs[i] = V[(base + V_OFF) / 4u + dd]; + } + workgroupBarrier(); // tiles staged + + // Scores: p[sq][sj] = ATTN_SCALE · (q_sq · k_sj) for sj < jn. + for (var i = tid; i < QB * JB; i = i + WG) { + let sq = i / JB; + let sj = i % JB; + if (sj < jn) { + var dot4 = vec4(0.0); + for (var d4 = 0u; d4 < D4; d4 = d4 + 1u) { + dot4 = dot4 + Qs[sq * D4 + d4] * vec4(Ks[sj * D4 + d4]); + } + p[i] = (dot4.x + dot4.y + dot4.z + dot4.w) * ATTN_SCALE; + } + } + workgroupBarrier(); // raw scores written + + // Online softmax bookkeeping — one thread per query row (serial jn scan; + // negligible next to the QB·JB·D score FLOPs above). + if (tid < QB) { + var m = rowMax[tid]; + for (var jj = 0u; jj < jn; jj = jj + 1u) { + m = max(m, p[tid * JB + jj]); + } + rescale[tid] = exp(rowMax[tid] - m); // first tile: exp(-1e30 - m) = 0 + var s = rowSum[tid] * rescale[tid]; + for (var jj = 0u; jj < jn; jj = jj + 1u) { + let e = exp(p[tid * JB + jj] - m); + p[tid * JB + jj] = e; + s = s + e; + } + rowMax[tid] = m; + rowSum[tid] = s; + } + workgroupBarrier(); // exp'd scores + rescale visible + + // Accumulate this tile's V contribution into the register accumulator. + acc = acc * rescale[q]; + for (var jj = 0u; jj < jn; jj = jj + 1u) { + acc = acc + p[q * JB + jj] * vec4(Vs[jj * D4 + dq]); + } + workgroupBarrier(); // all reads of Ks/Vs/p done before the next staging + } + + // Store — Y rows mirror Q rows; skip tail-block rows past mEnd. + if (qb0 + q < mEnd) { + let out = acc / rowSum[q]; + let yoff = (qBase + qb0 + q) * H * D + h * D + dq * 4u; + Y[yoff] = {{T}}(out.x); + Y[yoff + 1u] = {{T}}(out.y); + Y[yoff + 2u] = {{T}}(out.z); + Y[yoff + 3u] = {{T}}(out.w); + } +} +`,Qr=`// Add + LayerNorm (post-LN): Y[r,i] = gamma[i]·(v-μ)/√(σ²+EPS) + beta[i], +// where v = X[r,i] + R[r,i], μ is the row mean and σ² the row population +// variance (Σ(v-μ)²/D — second pass over shared memory, not sumsq-μ², which +// is cancellation-prone). f32 math regardless of storage type. +// +// One workgroup per row: dispatchWorkgroups(rows). WG threads each own +// ceil(D/WG) strided elements (D=448, WG=256 → ≤2 each). +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 the f16 enable directive when T is f16, else empty +// T storage type of X/R/gamma/beta/Y (f16|f32) +// WG workgroup size (256; must be a power of two) +// D row width (d_model, 448) +// EPS layernorm epsilon literal +// IF_SG/IF_NOSG subgroup reduction variant (flags.sg): subgroupAdd folds +// each subgroup's partials without barriers → 2 barriers per +// row instead of ~18. Reduction ORDER differs from the tree, +// so results are tolerance-equal, not bit-equal. +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} +{{ENABLE_SG}} + +struct Params { rows: u32, _pad0: u32, _pad1: u32, _pad2: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var X: array<{{T}}>; +@group(0) @binding(2) var R: array<{{T}}>; +@group(0) @binding(3) var gamma: array<{{T}}>; +@group(0) @binding(4) var beta: array<{{T}}>; +@group(0) @binding(5) var Y: array<{{T}}>; + +const D: u32 = {{D}}u; +const WG: u32 = {{WG}}u; + +var vbuf: array; +{{IF_NOSG}} +var scratch: array; +{{/IF_NOSG}} +{{IF_SG}} +// One slot per subgroup; two arrays so pass 2 never overwrites slots pass 1 +// readers still need (saves a barrier). WG/4 covers the spec-minimum +// subgroup size of 4. +var sgSum: array; +var sgSq: array; +{{/IF_SG}} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3{{IF_SG}}, @builtin(subgroup_size) sgSize: u32{{/IF_SG}}) { + // Uniform per workgroup (one workgroup per row), so returning here before + // any barrier is safe: the whole workgroup exits together. + if (wid.x >= params.rows) { return; } + let base = wid.x * D; + let tid = lid.x; + + // Load v = X + R into shared memory, accumulating this thread's partial sum. + var sum: f32 = 0.0; + for (var i = tid; i < D; i = i + WG) { + let v = f32(X[base + i]) + f32(R[base + i]); + vbuf[i] = v; + sum = sum + v; + } +{{IF_NOSG}} + scratch[tid] = sum; + workgroupBarrier(); + + // Tree-reduce the partial sums → row sum in scratch[0]. + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { scratch[tid] = scratch[tid] + scratch[tid + s]; } + workgroupBarrier(); + } + let mu = scratch[0] / f32(D); + // All reads of scratch[0] must complete before phase 2 overwrites scratch. + workgroupBarrier(); + + // Second pass: sum of squared deviations → population variance. + var sq: f32 = 0.0; + for (var i = tid; i < D; i = i + WG) { + let dev = vbuf[i] - mu; + sq = sq + dev * dev; + } + scratch[tid] = sq; + workgroupBarrier(); + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { scratch[tid] = scratch[tid] + scratch[tid + s]; } + workgroupBarrier(); + } + let inv = inverseSqrt(scratch[0] / f32(D) + {{EPS}}); +{{/IF_NOSG}} +{{IF_SG}} + // sgId assumes linear subgroup layout (tid/sgSize) — holds on + // Metal/D3D12/Vulkan under Tint; the sg kernel equiv gate catches a device + // where it doesn't. + let sgId = tid / sgSize; + let nSg = WG / sgSize; + // subgroupAdd must sit in subgroup-uniform control flow — outside the elect. + let s1 = subgroupAdd(sum); + if (subgroupElect()) { sgSum[sgId] = s1; } + workgroupBarrier(); // also publishes vbuf for pass 2 + var total = 0.0; + for (var i = 0u; i < nSg; i = i + 1u) { total = total + sgSum[i]; } + let mu = total / f32(D); + + // Second pass: sum of squared deviations → population variance. + var sq: f32 = 0.0; + for (var i = tid; i < D; i = i + WG) { + let dev = vbuf[i] - mu; + sq = sq + dev * dev; + } + let s2 = subgroupAdd(sq); + if (subgroupElect()) { sgSq[sgId] = s2; } + workgroupBarrier(); + var total2 = 0.0; + for (var i = 0u; i < nSg; i = i + 1u) { total2 = total2 + sgSq[i]; } + let inv = inverseSqrt(total2 / f32(D) + {{EPS}}); +{{/IF_SG}} + + for (var i = tid; i < D; i = i + WG) { + Y[base + i] = {{T}}(f32(gamma[i]) * (vbuf[i] - mu) * inv + f32(beta[i])); + } +} +`,Xr=`// Fused projection + residual add + LayerNorm for the decode step's three +// LN-terminated sites (self_out+ln1, cross_out+ln2, fc2+ln3): +// Y[m,:] = LN( f16round(X[m,:]·W + bias) + R[m,:] ) · gamma + beta +// One workgroup per row — LN needs the whole N=448 output row, which caps the +// GEMM parallelism at M workgroups. That starves the GPU at large batch (the +// GEMV kernel launches ~100× more threads), so this kernel is routed at +// SMALL B only, where the decode step is dispatch-overhead-bound and merging +// two dispatches into one is worth more than GEMM occupancy (25 → 19 +// dispatches per step). +// +// Numerics match the unfused pair as closely as split kernels allow: the +// projection result is rounded through the storage type BEFORE the residual +// add — exactly the Y-buffer round trip the unfused path performs — and the +// LN is the same two-pass form (mean, then Σ(v−μ)²; not the +// cancellation-prone sumsq−μ²) in f32. Only reduction/accumulation ORDER +// differs (per-thread serial K here vs TK-lane trees there) — f32 ULP-level, +// gated by m3/golden like every kernel-routing change. +// +// Layouts: X [M,K] row-major, W [K,N] row-major (the original NWT .weight +// tensors — NOT the .wt copies), N % 4 == 0, K % 4 == 0 (both hold for all +// three sites: K = 448 | 1792, N = 448). The X row is staged into shared +// memory once (vec4 loads) so the K-loop never touches global X. +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 f16 enable directive when T is f16 +// T storage type of all tensors (f16|f32); math is f32 +// WG workgroup size (128; power of two) +// KDIM K as a const (shared X row size) — one pipeline per K +// D output row width N (= d_model 448) +// EPS layernorm epsilon literal +// IF_SG/IF_NOSG subgroup LN reduction (flags.sg): subgroupAdd folds each +// subgroup's partials without barriers — 2 LN barriers per +// row instead of ~16 (same shape as add_layernorm.wgsl). +// Tolerance-equal, not bit-equal (reduction order). +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} +{{ENABLE_SG}} + +struct Params { M: u32, _pad0: u32, _pad1: u32, _pad2: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var X: array>; // [M, K/4] +@group(0) @binding(2) var W: array>; // [K, N/4] +@group(0) @binding(3) var B: array<{{T}}>; // [N] +@group(0) @binding(4) var R: array<{{T}}>; // [M, N] residual +@group(0) @binding(5) var gamma: array<{{T}}>; // [N] +@group(0) @binding(6) var beta: array<{{T}}>; // [N] +@group(0) @binding(7) var Y: array<{{T}}>; // [M, N] + +const K: u32 = {{KDIM}}u; +const D: u32 = {{D}}u; +const D4: u32 = D / 4u; +const K4: u32 = K / 4u; +const WG: u32 = {{WG}}u; + +var Xs: array; // the staged input row +var vbuf: array; // v = f16round(gemm) + residual +{{IF_NOSG}} +var scratch: array; +{{/IF_NOSG}} +{{IF_SG}} +var sgSum: array; // one slot per subgroup (min size 4) +var sgSq: array; +{{/IF_SG}} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3{{IF_SG}}, @builtin(subgroup_size) sgSize: u32{{/IF_SG}}) { + // Uniform per workgroup (one workgroup per row) — safe early return. + if (wid.x >= params.M) { return; } + let m = wid.x; + let tid = lid.x; + + // Stage the X row: one vec4 global read per thread-iteration. + for (var i = tid; i < K4; i = i + WG) { + let xv = vec4(X[m * K4 + i]); + Xs[4u * i] = xv.x; + Xs[4u * i + 1u] = xv.y; + Xs[4u * i + 2u] = xv.z; + Xs[4u * i + 3u] = xv.w; + } + workgroupBarrier(); + + // Projection: each thread owns output quad(s) q, walking the full K with a + // 2-way unrolled chain (K is even at both sites). W reads along N are + // coalesced across threads within each k-row. + for (var q = tid; q < D4; q = q + WG) { + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + for (var k = 0u; k < K; k = k + 2u) { + acc0 = fma(vec4(Xs[k]), vec4(W[k * D4 + q]), acc0); + acc1 = fma(vec4(Xs[k + 1u]), vec4(W[(k + 1u) * D4 + q]), acc1); + } + let n0 = 4u * q; + let bq = vec4(f32(B[n0]), f32(B[n0 + 1u]), f32(B[n0 + 2u]), f32(B[n0 + 3u])); + // Round through the storage type BEFORE the residual add — replicating + // the unfused path's Y-buffer round trip bit-for-bit. + let g = vec4<{{T}}>(acc0 + acc1 + bq); + let rbase = m * D + n0; + vbuf[n0] = f32(g.x) + f32(R[rbase]); + vbuf[n0 + 1u] = f32(g.y) + f32(R[rbase + 1u]); + vbuf[n0 + 2u] = f32(g.z) + f32(R[rbase + 2u]); + vbuf[n0 + 3u] = f32(g.w) + f32(R[rbase + 3u]); + } + workgroupBarrier(); + + // LayerNorm — same two-pass structure as add_layernorm.wgsl. + var sum: f32 = 0.0; + for (var i = tid; i < D; i = i + WG) { sum = sum + vbuf[i]; } +{{IF_NOSG}} + scratch[tid] = sum; + workgroupBarrier(); + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { scratch[tid] = scratch[tid] + scratch[tid + s]; } + workgroupBarrier(); + } + let mu = scratch[0] / f32(D); + workgroupBarrier(); // scratch[0] reads before phase 2 overwrites + + var sq: f32 = 0.0; + for (var i = tid; i < D; i = i + WG) { + let dev = vbuf[i] - mu; + sq = sq + dev * dev; + } + scratch[tid] = sq; + workgroupBarrier(); + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { scratch[tid] = scratch[tid] + scratch[tid + s]; } + workgroupBarrier(); + } + let inv = inverseSqrt(scratch[0] / f32(D) + {{EPS}}); +{{/IF_NOSG}} +{{IF_SG}} + // sgId assumes linear subgroup layout (tid/sgSize) — see add_layernorm. + let sgId = tid / sgSize; + let nSg = WG / sgSize; + let s1 = subgroupAdd(sum); // subgroup-uniform flow: outside the elect + if (subgroupElect()) { sgSum[sgId] = s1; } + workgroupBarrier(); + var total = 0.0; + for (var i = 0u; i < nSg; i = i + 1u) { total = total + sgSum[i]; } + let mu = total / f32(D); + + var sq: f32 = 0.0; + for (var i = tid; i < D; i = i + WG) { + let dev = vbuf[i] - mu; + sq = sq + dev * dev; + } + let s2 = subgroupAdd(sq); + if (subgroupElect()) { sgSq[sgId] = s2; } + workgroupBarrier(); + var total2 = 0.0; + for (var i = 0u; i < nSg; i = i + 1u) { total2 = total2 + sgSq[i]; } + let inv = inverseSqrt(total2 / f32(D) + {{EPS}}); +{{/IF_SG}} + + let ybase = m * D; + for (var i = tid; i < D; i = i + WG) { + Y[ybase + i] = {{T}}(f32(gamma[i]) * (vbuf[i] - mu) * inv + f32(beta[i])); + } +} +`,Yr=`// Split-K epilogue: fold the nz raw f32 partial slices a split-K +// gemm_tiled2 dispatch left in PART [nz, M, N] into Y [M, N], applying the +// bias/SiLU epilogue the GEMM skipped (it must run on the FULL sum, once). +// Ascending-z summation is deterministic; the only numerical difference vs +// the unsplit kernel is the f32 re-association at the nz seams. +// +// IF_STORE_KV (split-K self_qkv): Y is the fused q|k|v row [M, 3·HD]; the +// k and v slices additionally scatter into the decode caches +// [M, Lmax, HD] at position dims.t — the SAME {{OUT_T}} value written to Y, +// so the caches stay bit-identical to Y's slices (the kv_append contract, +// exactly like the GEMV epilogue it replaces). +// +// Template placeholders: ENABLE_F16, T (bias storage type), OUT_T, WG, +// IF_BIAS, IF_SILU, IF_STORE_KV. +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Dims { M: u32, N: u32, NZ: u32, t: u32, Lmax: u32, _p0: u32, _p1: u32, _p2: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> dims: Dims; +@group(0) @binding(1) var PART: array; +@group(0) @binding(2) var B: array<{{T}}>; +@group(0) @binding(3) var Y: array<{{OUT_T}}>; +{{IF_STORE_KV}} +@group(0) @binding(4) var Kc: array<{{OUT_T}}>; +@group(0) @binding(5) var Vc: array<{{OUT_T}}>; +{{/IF_STORE_KV}} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(global_invocation_id) gid: vec3) { + _ = &B; // keep B in the auto layout when the bias block is compiled out + let i = gid.x; + let total = dims.M * dims.N; + if (i >= total) { return; } + var acc: f32 = 0.0; + for (var z: u32 = 0u; z < dims.NZ; z = z + 1u) { + acc = acc + PART[z * total + i]; + } + {{IF_BIAS}} + acc = acc + f32(B[i % dims.N]); + {{/IF_BIAS}} + {{IF_SILU}} + acc = acc / (1.0 + exp(-acc)); + {{/IF_SILU}} + let v = {{OUT_T}}(acc); + Y[i] = v; + {{IF_STORE_KV}} + let HD = dims.N / 3u; + let m = i / dims.N; + let n = i % dims.N; + if (n >= HD) { + let dstBase = (m * dims.Lmax + dims.t) * HD; + if (n < 2u * HD) { + Kc[dstBase + (n - HD)] = v; + } else { + Vc[dstBase + (n - 2u * HD)] = v; + } + } + {{/IF_STORE_KV}} +} +`,Jr=`// Token + position embedding: Y[r,i] = f32(table[id,i])·EMBED_SCALE +// + f32(posEmbed[pos,i]) (pos NOT scaled) +// One workgroup per row: dispatchWorkgroups(nRows). Two modes, selected at +// build time (exactly one of SRC_IDS / DECODE): +// SRC_IDS (encoder): ids is src token ids [B*S]; id = ids[r], pos = r % s +// (batch rows are consecutive: r = b*S + m). PAD rows +// embed the pad token normally — masking happens in +// attention, not here. With PACKED (encoder row-packing: +// pad rows dropped, T = Σ lens rows) each word carries +// its own position: ids[r] = (pos << 16) | id — id fits +// (VOCAB 24000 < 2^16) and the dispatch enforces +// S < 2^16. +// DECODE: ids is the token ring [T_max*B]; row r = batch index b; +// id = DECODER_START when t == 0, else ids[(t-1)*batch + r]; +// pos = t. +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 the f16 enable directive when T is f16, else empty +// T storage type of table/posEmbed/Y (f16|f32) +// WG workgroup size (224 → 2 elements per thread at D=448) +// D row width (d_model, 448) +// EMBED_SCALE √d_model as a full-precision literal +// IF_SRC_IDS / IF_DECODE mode blocks +// DECODER_START decoder start token id (DECODE mode only) +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Params { nRows: u32, t: u32, batch: u32, s: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var ids: array; +@group(0) @binding(2) var table: array<{{T}}>; +@group(0) @binding(3) var posEmbed: array<{{T}}>; +@group(0) @binding(4) var Y: array<{{T}}>; + +const D: u32 = {{D}}u; +const WG: u32 = {{WG}}u; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + // Uniform per workgroup (one workgroup per row) — safe early return. + if (wid.x >= params.nRows) { return; } + let r = wid.x; +{{IF_SRC_IDS}} +{{IF_PACKED}} + let id = ids[r] & 0xffffu; + let pos = ids[r] >> 16u; +{{/IF_PACKED}} +{{IF_NOPACKED}} + let id = ids[r]; + let pos = r % params.s; +{{/IF_NOPACKED}} +{{/IF_SRC_IDS}} +{{IF_DECODE}} + var id: u32 = {{DECODER_START}}u; + if (params.t != 0u) { id = ids[(params.t - 1u) * params.batch + r]; } + let pos = params.t; +{{/IF_DECODE}} + let toff = id * D; + let poff = pos * D; + let base = r * D; + for (var i = lid.x; i < D; i = i + WG) { + Y[base + i] = {{T}}(f32(table[toff + i]) * {{EMBED_SCALE}} + f32(posEmbed[poff + i])); + } +} +`,Zr=`// Row scatter for encoder row-packing: packed activations [T, N] → padded +// [B·S, N], where T = Σ lens[b] and padded row (b, m < lens[b]) comes from +// packed row starts[b] + m. Padding rows are left untouched — the destination +// arena buffer is zero-initialized by WebGPU, so they read as zeros +// downstream (decoder cross-attention masks j ≥ len and never reads them). +// +// One workgroup per padded row, vec4 element copies (N % 4 == 0 enforced by +// the dispatch). Early returns are uniform: wid-derived plus read-only +// storage loads at workgroup-uniform indices; there are no barriers. +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16 f16 enable directive when T is f16, else empty +// T storage type of X/Y (f16|f32) +// WG workgroup size +{{ENABLE_IMMEDIATE}} +{{ENABLE_F16}} + +struct Params { B: u32, S: u32, N4: u32, _pad: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var starts: array; // packed row offsets [B] +@group(0) @binding(2) var lens: array; // sequence lengths [B] +@group(0) @binding(3) var X: array>; +@group(0) @binding(4) var Y: array>; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + if (wid.x >= params.B * params.S) { return; } + let b = wid.x / params.S; + let m = wid.x % params.S; + if (m >= lens[b]) { return; } + let src = (starts[b] + m) * params.N4; + let dst = wid.x * params.N4; + for (var i = lid.x; i < params.N4; i = i + {{WG}}u) { + Y[dst + i] = X[src + i]; + } +} +`,et=`// In-place live-row gather for decode compaction: row map[i] moves to row i +// for i < rows, inside the SAME storage buffer. copyBufferToBuffer cannot use +// one GPUBuffer as both source and destination, so this compute pass performs +// the move. +// +// The compaction invariant makes the ordering safe: map is strictly ascending, +// hence map[i] >= i. Rows move only toward lower indices. One workgroup walks +// destination rows in ascending order and a storage barrier separates rows; +// writing row i can therefore never clobber a source needed by a later row. +// +// Offsets and lengths are u32 counts rather than vec4 counts because repetition +// mask rows are not necessarily 16-byte aligned. +{{ENABLE_IMMEDIATE}} + +struct Params { rows: u32, rowStride: u32, copyLen: u32, _pad: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var map: array; +@group(0) @binding(2) var data: array; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(local_invocation_id) lid: vec3) { + for (var i = 0u; i < params.rows; i = i + 1u) { + let src = map[i]; + if (src != i) { + let srcBase = src * params.rowStride; + let dstBase = i * params.rowStride; + for (var k = lid.x; k < params.copyLen; k = k + {{WG}}u) { + data[dstBase + k] = data[srcBase + k]; + } + } + storageBarrier(); + } +} +`,nt=`// Decode-step MEGAKERNEL: one workgroup computes one batch row's ENTIRE +// decoder layer — 8 dispatches collapse into 1 (small-B latency: the b1 step +// is dispatch-overhead-bound, ~15µs of fixed cost per dispatch against +// ~5-10µs kernels). With IF_EMBED (layer 0) the decode embedding folds in +// too, so a b1 step becomes: mega L0 → mega L1 → lm_head → argmax. +// +// Stages (workgroupBarrier between each; all math f32; every stage boundary +// value is ROUNDED THROUGH THE STORAGE TYPE first — replicating the unfused +// path's activation-buffer round trips, the gemm_row_ln precedent): +// S0 x ← embed (L0: ring/table/pos, matching embed.wgsl DECODE) or the +// global hidden buffer X (L1) +// S1 qkv = x·Wqkvᵀ + b; q stays in shared, k|v quads store to the caches +// at position t (the same rounded value — the kv_append contract) +// S2 self-attention over positions 0..t (phase structure and score/exp/ +// fold order mirror attention.wgsl; heads loop serially) +// S3 x = LN1(x + self_out(attn)) +// S4 q = cross_q(x) +// S5 cross-attention over lens[b] encoder positions (fused crossKV k|v) +// S6 x = LN2(x + cross_out(attn)) +// S7 ffn = SiLU(fc1(x)) (SiLU in f32 BEFORE the f16 round, as gemv) +// S8 X ← LN3(x + fc2(ffn)) (written back to the global hidden buffer) +// +// All projections read the ORIGINAL [K, N] row-major \`.weight\` tensors with +// the gemm_row_ln access pattern: thread q owns output quad q, and at each k +// the threads read CONSECUTIVE quads of W's k-row — fully coalesced. (The +// first version walked the transposed [N,K] copies, one row per thread — +// every load touched 32 distinct lines and the whole kernel ran at ~4.5GB/s, +// 4.3× SLOWER than the chain it replaced. One workgroup has only ~8 warps of +// latency-hiding; coalescing is everything here.) Four independent +// accumulators per k-quad keep 4 loads in flight per thread. Every tensor is +// addressed inside the ONE weights buffer via compile-time vec4-element +// offsets (…4 defines = byteOffset/8; manifest offsets are 256-aligned so /8 +// is exact). One pipeline per layer. +// +// NOT bit-exact vs the unfused chain (accumulation/reduction ORDER differs +// at every site) — routed like every kernel change: m3/golden gates + the +// step-0 + divergence-rate equiv (mega_equiv), e2e A/B decides the batch +// threshold. +// +// Shared budget (16KB): (2·HD4 + TMP4)·16 + SCORES_CAP·4 + WG·4 — 13,184B +// for MoxhiMT-30 (448/1792), 16,256B for HachimiMT-60 (576/2304); checked at +// dispatch. tmp4 is max(FFN4, HD4 + WG) quads: the attention phase-3 partial +// scratch tmp4[HD4 .. HD4+WG) must fit even when the model's FFN is small +// (q lives in [0..HD4) for self, out4 for cross; the fold result lands in +// [0..HD4) only after all partial reads). +// +// Template placeholders (buildShader in pipelines.js): +// ENABLE_F16, T (must be f16 — the .wt copies only exist as f16), WG (256) +// ENABLE_SG + IF_SG/IF_NOSG subgroup wgMax/wgSum (flags.sg) — see below +// IF_EMBED / IF_NOEMBED layer-0 embedding fold (TABLE4/POS4/EMBED_SCALE/ +// DECODER_START live inside IF_EMBED) +// H, D, FFN4, LMAX, SCORES_CAP, ATTN_SCALE, EPS +// QKVW4 QKVB4 OUTW4 OUTB4 LN1G4 LN1B4 CQW4 CQB4 COW4 COB4 LN2G4 LN2B4 +// FC1W4 FC1B4 FC2W4 FC2B4 LN3G4 LN3B4 per-tensor vec4 offsets into W +// +// SG mode (flags.sg): wgMax/wgSum are where this kernel's barriers live — +// each tree call is 10 workgroupBarriers, and one layer makes 38 of them +// (2 per attention head × 2 sides × H, 2 per LN × 3), ~400 barriers per +// step per layer. That is exactly what the megakernel pays on Apple GPUs +// (Metal mega_sweep: mega LOSES b1 there while winning −12% on NVIDIA). +// With sg each call is subgroupMax/Add → one partial per subgroup → ONE +// barrier → serial fold over ≤ WG/4 partials, ~5× fewer barriers overall. +{{ENABLE_IMMEDIATE}} +{{ENABLE_SG}} +{{ENABLE_F16}} + +struct Params { + B: u32, // batch rows (grid.x) + t: u32, // decode step (cache position; embed pos; self len = t+1) + S: u32, // encoder crossKV position capacity (padded S) + _pad: u32, +} + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var W: array>; // whole weights buffer +@group(0) @binding(2) var ring: array; // token ring +@group(0) @binding(3) var Kc: array>; // [B, LMAX, H·D] +@group(0) @binding(4) var Vc: array>; +@group(0) @binding(5) var CKV: array>; // [B·S, 2·H·D] fused k|v +@group(0) @binding(6) var lens: array; +@group(0) @binding(7) var X: array>; // hidden [B, H·D] + +const H: u32 = {{H}}u; +const D: u32 = {{D}}u; +const D4: u32 = D / 4u; // quads per head +const HD4: u32 = H * D4; // quads per d_model row +const QKV4: u32 = 3u * HD4; // fused q|k|v quads +const FFN4: u32 = {{FFN4}}u; // ffn quads (FFN/4) +const KQ_FFN: u32 = FFN4; // fc2 K quads (K = FFN) +const LMAX: u32 = {{LMAX}}u; +const SCORES_MAX: u32 = {{SCORES_CAP}}u; +const ATTN_SCALE: f32 = {{ATTN_SCALE}}; +const WG: u32 = {{WG}}u; +const JT: u32 = WG / D4; // attention phase-3 j-lanes per d-quad +const TMP4: u32 = max(FFN4, HD4 + WG); // ffn row AND attn partial scratch fit + +var xs4: array, HD4>; // hidden state (residual base) +var tmp4: array, TMP4>; // q / vbuf / ffn / attn partials +var out4: array, HD4>; // stage outputs +var scores: array; +{{IF_NOSG}} +var red: array; +{{/IF_NOSG}} +{{IF_SG}} +// One partial per subgroup, in TWO alternating slots of NSG_CAP (WG/4 covers +// the spec-minimum subgroup size 4). The alternation is what buys the single +// barrier per call: call N's fold reads slot A strictly before every thread +// passes call N+1's barrier (slot B), and call N+2's elect-writes to slot A +// happen strictly after it — so no trailing barrier is needed to protect +// reuse. sgId = tid/sgSize assumes the linear tid→subgroup layout (same bet +// as add_layernorm.wgsl; the equiv gates catch a violating backend). +const NSG_CAP: u32 = WG / 4u; +var red: array; +var sgId: u32; +var nSg: u32; +var redSlot: u32 = 0u; +{{/IF_SG}} + +fn wgMax(tid: u32, v: f32) -> f32 { +{{IF_SG}} + let s1 = subgroupMax(v); + let base = redSlot * NSG_CAP; + if (subgroupElect()) { red[base + sgId] = s1; } + workgroupBarrier(); + var r = red[base]; + for (var i = 1u; i < nSg; i = i + 1u) { r = max(r, red[base + i]); } + redSlot = 1u - redSlot; + return r; +{{/IF_SG}} +{{IF_NOSG}} + red[tid] = v; + workgroupBarrier(); + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { red[tid] = max(red[tid], red[tid + s]); } + workgroupBarrier(); + } + let r = red[0]; + workgroupBarrier(); // red[0] reads done before the next reduction reuses red + return r; +{{/IF_NOSG}} +} + +fn wgSum(tid: u32, v: f32) -> f32 { +{{IF_SG}} + let s1 = subgroupAdd(v); + let base = redSlot * NSG_CAP; + if (subgroupElect()) { red[base + sgId] = s1; } + workgroupBarrier(); + var r = red[base]; + for (var i = 1u; i < nSg; i = i + 1u) { r = r + red[base + i]; } + redSlot = 1u - redSlot; + return r; +{{/IF_SG}} +{{IF_NOSG}} + red[tid] = v; + workgroupBarrier(); + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { red[tid] = red[tid] + red[tid + s]; } + workgroupBarrier(); + } + let r = red[0]; + workgroupBarrier(); + return r; +{{/IF_NOSG}} +} + +// One GEMV output quad, gemm_row_ln-style: thread computes outputs +// 4·n4 .. 4·n4+3 from the [K, N] row-major W — at each k, threads read +// consecutive quads of the k-row (coalesced across the workgroup). srcSel +// picks the shared source (0 = xs4, 1 = out4, 2 = tmp4); kq = K/4 source +// quads, nq = N/4 output quads (the W row stride). Four independent +// accumulators keep 4 loads in flight; the fold order is fixed +// (a0+a1)+(a2+a3). Returns f32 WITHOUT rounding — the caller rounds/routes. +fn gemvQuad(wOff: u32, bOff: u32, n4: u32, kq: u32, nq: u32, srcSel: u32) -> vec4 { + var a0 = vec4(0.0); + var a1 = vec4(0.0); + var a2 = vec4(0.0); + var a3 = vec4(0.0); + for (var k4 = 0u; k4 < kq; k4 = k4 + 1u) { + var xq: vec4; + if (srcSel == 0u) { xq = xs4[k4]; } + else if (srcSel == 1u) { xq = out4[k4]; } + else { xq = tmp4[k4]; } + let kBase = wOff + (k4 << 2u) * nq + n4; + a0 = fma(vec4(xq.x), vec4(W[kBase]), a0); + a1 = fma(vec4(xq.y), vec4(W[kBase + nq]), a1); + a2 = fma(vec4(xq.z), vec4(W[kBase + 2u * nq]), a2); + a3 = fma(vec4(xq.w), vec4(W[kBase + 3u * nq]), a3); + } + return (a0 + a1) + (a2 + a3) + vec4(W[bOff + n4]); +} + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3{{IF_SG}}, @builtin(subgroup_size) sgSize: u32{{/IF_SG}}) { + // Uniform per workgroup — safe early return before the first barrier. + if (wid.x >= params.B) { return; } + let b = wid.x; + let tid = lid.x; + let t = params.t; +{{IF_SG}} + sgId = tid / sgSize; + nSg = (WG + sgSize - 1u) / sgSize; +{{/IF_SG}} + + // ---- S0: hidden state into xs4 ---- +{{IF_EMBED}} + // embed.wgsl DECODE semantics: id = DECODER_START at t=0, else the ring + // token; pos = t; y = f16round(table·EMBED_SCALE + pos_embed). + var id: u32 = {{DECODER_START}}u; + if (t != 0u) { id = ring[(t - 1u) * params.B + b]; } + for (var i = tid; i < HD4; i = i + WG) { + let e = vec4(W[{{TABLE4}}u + id * HD4 + i]) * {{EMBED_SCALE}} + + vec4(W[{{POS4}}u + t * HD4 + i]); + xs4[i] = vec4(vec4<{{T}}>(e)); + } +{{/IF_EMBED}} +{{IF_NOEMBED}} + // Phony use: only the embed fold reads the ring, but the binding must stay + // statically used or layout 'auto' drops @binding(2) and the bind group + // (which always supplies it) fails validation — killing the whole submit. + _ = ring[0]; + for (var i = tid; i < HD4; i = i + WG) { + xs4[i] = vec4(X[b * HD4 + i]); + } +{{/IF_NOEMBED}} + workgroupBarrier(); + + // ---- S1: fused qkv projection; q → tmp4[0..HD4), k|v quads → caches ---- + let kvBase = (b * LMAX + t) * HD4; + for (var n4 = tid; n4 < QKV4; n4 = n4 + WG) { + let g = vec4<{{T}}>(gemvQuad({{QKVW4}}u, {{QKVB4}}u, n4, HD4, QKV4, 0u)); + if (n4 < HD4) { tmp4[n4] = vec4(g); } + else if (n4 < 2u * HD4) { Kc[kvBase + n4 - HD4] = g; } + else { Vc[kvBase + n4 - 2u * HD4] = g; } + } + workgroupBarrier(); + + // ---- S2: self-attention over positions 0..t (attention.wgsl phases) ---- + { + let len = min(t + 1u, LMAX); + for (var h = 0u; h < H; h = h + 1u) { + let hq = h * D4; + var lm: f32 = -1e30; + for (var j = tid; j < len; j = j + WG) { + let koff = (b * LMAX + j) * HD4 + hq; + var dot4 = vec4(0.0); + for (var i = 0u; i < D4; i = i + 1u) { + dot4 = dot4 + tmp4[hq + i] * vec4(Kc[koff + i]); + } + let sc = (dot4.x + dot4.y + dot4.z + dot4.w) * ATTN_SCALE; + scores[j] = sc; + lm = max(lm, sc); + } + let rowMax = wgMax(tid, lm); + var ls: f32 = 0.0; + for (var j = tid; j < len; j = j + WG) { + let e = exp(scores[j] - rowMax); + scores[j] = e; + ls = ls + e; + } + let denom = wgSum(tid, ls); + let dq = tid % D4; + let jg = tid / D4; + var acc = vec4(0.0); + if (jg < JT) { + for (var j = jg; j < len; j = j + JT) { + acc = acc + scores[j] * vec4(Vc[(b * LMAX + j) * HD4 + hq + dq]); + } + } + tmp4[HD4 + tid] = acc; // partial scratch; q region [0..HD4) untouched + workgroupBarrier(); + if (tid < D4) { + var o = vec4(0.0); + for (var g = 0u; g < JT; g = g + 1u) { o = o + tmp4[HD4 + g * D4 + tid]; } + out4[hq + tid] = vec4(vec4<{{T}}>(o / denom)); + } + workgroupBarrier(); // out4 + partial reads done before the next head + } + } + + // ---- S3: x = LN1(x + self_out(attn)); vbuf = tmp4 ---- + for (var n4 = tid; n4 < HD4; n4 = n4 + WG) { + let g = vec4<{{T}}>(gemvQuad({{OUTW4}}u, {{OUTB4}}u, n4, HD4, HD4, 1u)); + tmp4[n4] = vec4(g) + xs4[n4]; + } + workgroupBarrier(); + { + var s: f32 = 0.0; + for (var i = tid; i < HD4; i = i + WG) { + let v = tmp4[i]; + s = s + v.x + v.y + v.z + v.w; + } + let mu = wgSum(tid, s) / f32(H * D); + var sq: f32 = 0.0; + for (var i = tid; i < HD4; i = i + WG) { + let dv = tmp4[i] - vec4(mu); + sq = sq + dot(dv, dv); + } + let inv = inverseSqrt(wgSum(tid, sq) / f32(H * D) + {{EPS}}); + for (var i = tid; i < HD4; i = i + WG) { + let o = vec4(W[{{LN1G4}}u + i]) * (tmp4[i] - vec4(mu)) * inv + + vec4(W[{{LN1B4}}u + i]); + xs4[i] = vec4(vec4<{{T}}>(o)); + } + } + workgroupBarrier(); + + // ---- S4: q = cross_q(x) → out4 ---- + for (var n4 = tid; n4 < HD4; n4 = n4 + WG) { + out4[n4] = vec4(vec4<{{T}}>(gemvQuad({{CQW4}}u, {{CQB4}}u, n4, HD4, HD4, 0u))); + } + workgroupBarrier(); + + // ---- S5: cross-attention over lens[b] encoder positions → tmp4[0..HD4) ---- + { + let len = min(lens[b], SCORES_MAX); + for (var h = 0u; h < H; h = h + 1u) { + let hq = h * D4; + var lm: f32 = -1e30; + for (var j = tid; j < len; j = j + WG) { + let koff = (b * params.S + j) * 2u * HD4 + hq; // k slice at offset 0 + var dot4 = vec4(0.0); + for (var i = 0u; i < D4; i = i + 1u) { + dot4 = dot4 + out4[hq + i] * vec4(CKV[koff + i]); + } + let sc = (dot4.x + dot4.y + dot4.z + dot4.w) * ATTN_SCALE; + scores[j] = sc; + lm = max(lm, sc); + } + let rowMax = wgMax(tid, lm); + var ls: f32 = 0.0; + for (var j = tid; j < len; j = j + WG) { + let e = exp(scores[j] - rowMax); + scores[j] = e; + ls = ls + e; + } + let denom = wgSum(tid, ls); + let dq = tid % D4; + let jg = tid / D4; + var acc = vec4(0.0); + if (jg < JT) { + for (var j = jg; j < len; j = j + JT) { + // v slice at element offset H·D within the fused k|v position + acc = acc + scores[j] * vec4(CKV[(b * params.S + j) * 2u * HD4 + HD4 + hq + dq]); + } + } + tmp4[HD4 + tid] = acc; // fold results land in [0..HD4) only afterwards + workgroupBarrier(); + if (tid < D4) { + var o = vec4(0.0); + for (var g = 0u; g < JT; g = g + 1u) { o = o + tmp4[HD4 + g * D4 + tid]; } + tmp4[hq + tid] = vec4(vec4<{{T}}>(o / denom)); + } + workgroupBarrier(); + } + } + + // ---- S6: x = LN2(x + cross_out(attn)); vbuf = out4 ---- + for (var n4 = tid; n4 < HD4; n4 = n4 + WG) { + let g = vec4<{{T}}>(gemvQuad({{COW4}}u, {{COB4}}u, n4, HD4, HD4, 2u)); + out4[n4] = vec4(g) + xs4[n4]; + } + workgroupBarrier(); + { + var s: f32 = 0.0; + for (var i = tid; i < HD4; i = i + WG) { + let v = out4[i]; + s = s + v.x + v.y + v.z + v.w; + } + let mu = wgSum(tid, s) / f32(H * D); + var sq: f32 = 0.0; + for (var i = tid; i < HD4; i = i + WG) { + let dv = out4[i] - vec4(mu); + sq = sq + dot(dv, dv); + } + let inv = inverseSqrt(wgSum(tid, sq) / f32(H * D) + {{EPS}}); + for (var i = tid; i < HD4; i = i + WG) { + let o = vec4(W[{{LN2G4}}u + i]) * (out4[i] - vec4(mu)) * inv + + vec4(W[{{LN2B4}}u + i]); + xs4[i] = vec4(vec4<{{T}}>(o)); + } + } + workgroupBarrier(); + + // ---- S7: ffn = SiLU(fc1(x)) → tmp4[0..FFN4) (SiLU in f32, then round) ---- + for (var n4 = tid; n4 < FFN4; n4 = n4 + WG) { + var v = gemvQuad({{FC1W4}}u, {{FC1B4}}u, n4, HD4, FFN4, 0u); + v = v / (vec4(1.0) + exp(-v)); + tmp4[n4] = vec4(vec4<{{T}}>(v)); + } + workgroupBarrier(); + + // ---- S8: X ← LN3(x + fc2(ffn)); vbuf = out4 ---- + for (var n4 = tid; n4 < HD4; n4 = n4 + WG) { + let g = vec4<{{T}}>(gemvQuad({{FC2W4}}u, {{FC2B4}}u, n4, KQ_FFN, HD4, 2u)); + out4[n4] = vec4(g) + xs4[n4]; + } + workgroupBarrier(); + { + var s: f32 = 0.0; + for (var i = tid; i < HD4; i = i + WG) { + let v = out4[i]; + s = s + v.x + v.y + v.z + v.w; + } + let mu = wgSum(tid, s) / f32(H * D); + var sq: f32 = 0.0; + for (var i = tid; i < HD4; i = i + WG) { + let dv = out4[i] - vec4(mu); + sq = sq + dot(dv, dv); + } + let inv = inverseSqrt(wgSum(tid, sq) / f32(H * D) + {{EPS}}); + for (var i = tid; i < HD4; i = i + WG) { + let o = vec4(W[{{LN3G4}}u + i]) * (out4[i] - vec4(mu)) * inv + + vec4(W[{{LN3B4}}u + i]); + X[b * HD4 + i] = vec4<{{T}}>(o); + } + } +} +`,rt=`// Repetition-penalty + greedy argmax + token writeback — the kernel that keeps +// the decode loop GPU-resident: logits → next token with no CPU roundtrip. +// +// One workgroup per batch row b: dispatchWorkgroups(B). +// scan: thread strides v = tid, tid+WG, …: x = logits[b·V+v] + bias[v] +// (final_logits_bias is part of the model's logits BEFORE the +// penalty runs, matching HF); if v's bit is set in the row's seen +// bitmask: x = x>0 ? x/PENALTY : x·PENALTY. Strict > keeps the +// EARLIEST max within a thread (v ascends per thread), so the +// lowest-index tie-break (torch.argmax first-max) holds locally. +// reduce: tree reduction over (val, idx) pairs; comparator prefers higher +// val, then LOWER idx on exact equality — global lowest-index ties. +// epilogue (thread 0 only): +// wasDone → token = PAD, no done/bitmask update; +// else: token = argmax; token == EOS → done[b] = 1 (the EOS token +// itself IS written to the ring); the picked token's bit is ALWAYS +// set for a not-already-done row — including the EOS pick (HF +// appends eos to input_ids too; harmless, subsequent picks are PAD +// and skip the mask). Ring slot: t·B + b (embed at step t+1 reads +// ((t+1)-1)·B + b — the same slot). +// +// Races: thread 0 is the ONLY writer of done/tokens/bitmask for its row, and +// other rows' workgroups touch disjoint slots — no atomics needed. The +// reduction loop's final workgroupBarrier() orders the epilogue after every +// thread's scan, so thread 0's bitmask |= cannot race this dispatch's own +// bitmask reads (within one dispatch the |= wouldn't be visible to sibling +// workgroups anyway, but the barrier makes the intra-workgroup ordering +// explicit). Cross-dispatch visibility (step t's bit seen by step t+1) is +// guaranteed by WebGPU dispatch/submission ordering. +// +// Template placeholders (buildShader in pipelines.js): +// WG workgroup size (256; must be a power of two) +// V vocab size (24000) +// EOS / PAD special token ids (2 / 0) +// PENALTY repetition penalty as a literal (1.2) +// MASK_WORDS bitmask u32 words per row (V/32 = 750) + +{{ENABLE_IMMEDIATE}} + +struct Params { B: u32, t: u32, _pad0: u32, _pad1: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var logits: array; // [B·V] +@group(0) @binding(2) var bias: array; // [V] +@group(0) @binding(3) var bitmask: array; // [B·MASK_WORDS] +@group(0) @binding(4) var done: array; // [B] +@group(0) @binding(5) var tokens: array; // ring [T_max·B] +{{IF_SHORT}} +// Shortlisted lm_head (tiled-unfused arm): logits/bias/bitmask are all in +// LOCAL column space (V = shortlist length); IDMAP maps the winning local +// index to a vocab id, and gmask keeps the vocab-space mask in sync for any +// full-vocab state that inherits this row (see argmax_reduce.wgsl). +@group(0) @binding(6) var IDMAP: array; // [N_short] +@group(0) @binding(7) var gmask: array; // [B·GMASK_WORDS] +{{/IF_SHORT}} + +const V: u32 = {{V}}u; +const WG: u32 = {{WG}}u; +const EOS: u32 = {{EOS}}u; +const PAD: u32 = {{PAD}}u; +const PENALTY: f32 = {{PENALTY}}; +const MASK_WORDS: u32 = {{MASK_WORDS}}u; +{{IF_SHORT}} +const GMASK_WORDS: u32 = {{GMASK_WORDS}}u; +{{/IF_SHORT}} +// Finite f32 lowest — WGSL has no infinity literal (and -1.0/0.0 is a +// const-expr error). Logits + bias are finite reals, so this is safely below. +const NEG_MAX: f32 = -3.40282e38; + +var bestVal: array; +var bestIdx: array; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + // Uniform per workgroup (one workgroup per row) — safe early return. + if (wid.x >= params.B) { return; } + let b = wid.x; + let tid = lid.x; + let lbase = b * V; + let mbase = b * MASK_WORDS; + + // Scan: strided over v, ascending per thread. + var val: f32 = NEG_MAX; + var idx: u32 = 0xffffffffu; + for (var v = tid; v < V; v = v + WG) { + var x = logits[lbase + v] + bias[v]; + let seen = (bitmask[mbase + (v >> 5u)] >> (v & 31u)) & 1u; + if (seen == 1u) { + if (x > 0.0) { x = x / PENALTY; } else { x = x * PENALTY; } + } + // Strict >: earliest (lowest-v) max wins within this thread. + if (x > val) { val = x; idx = v; } + } + bestVal[tid] = val; + bestIdx[tid] = idx; + workgroupBarrier(); + + // Tree reduce: higher val wins; exact-equal vals break to the LOWER idx. + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { + let ov = bestVal[tid + s]; + let oi = bestIdx[tid + s]; + if (ov > bestVal[tid] || (ov == bestVal[tid] && oi < bestIdx[tid])) { + bestVal[tid] = ov; + bestIdx[tid] = oi; + } + } + // Final iteration's barrier also orders the epilogue below after ALL + // threads' scans/reduction writes (see races note in the header). + workgroupBarrier(); + } + + // Epilogue: single writer per row. + if (tid == 0u) { + let wasDone = done[b]; + var tok = bestIdx[0]; +{{IF_SHORT}} + let lidx = tok; + tok = IDMAP[lidx]; +{{/IF_SHORT}} + if (wasDone == 1u) { + tok = PAD; + } else { + if (tok == EOS) { done[b] = 1u; } +{{IF_SHORT}} + bitmask[mbase + (lidx >> 5u)] = bitmask[mbase + (lidx >> 5u)] | (1u << (lidx & 31u)); + let gbase = b * GMASK_WORDS; + gmask[gbase + (tok >> 5u)] = gmask[gbase + (tok >> 5u)] | (1u << (tok & 31u)); +{{/IF_SHORT}} +{{IF_NOSHORT}} + bitmask[mbase + (tok >> 5u)] = bitmask[mbase + (tok >> 5u)] | (1u << (tok & 31u)); +{{/IF_NOSHORT}} + } + tokens[params.t * params.B + b] = tok; + } +} +`,tt=`// Stage 2 of the FUSED lm_head argmax (gemm_tiled2.wgsl IF_ARGMAX): fold the +// per-workgroup (val, idx) partials — final_logits_bias and the repetition +// penalty were already applied in the GEMM epilogue — into the row's token, +// then run argmax_penalty.wgsl's done/EOS/bitmask/ring epilogue VERBATIM. +// One workgroup per batch row; NT = ceil(V/BN) partials per row (375 at +// BN=64), so this dispatch reads ~3KB per row where argmax_penalty re-read +// 96KB of logits plus 96KB of bias. +// +// Tie-breaking matches argmax_penalty exactly: partial j covers vocab ids +// [j·BN, (j+1)·BN) — ascending j = ascending, disjoint id ranges — and each +// partial already carries its tile's lowest-index max, so a strict > over the +// per-thread ascending stride plus the (higher val, then LOWER idx) tree +// comparator reproduce torch.argmax first-max semantics globally. +// +// Template placeholders (buildShader in pipelines.js): +// WG workgroup size (256; must be a power of two) +// EOS / PAD special token ids (2 / 0) +// MASK_WORDS bitmask u32 words per row (V/32 = 750) + +{{ENABLE_IMMEDIATE}} + +struct Params { B: u32, t: u32, NT: u32, _pad0: u32 } + +{{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; +@group(0) @binding(1) var P: array>; // [B·NT] (bitcast f32 val, idx) +@group(0) @binding(2) var bitmask: array; // [B·MASK_WORDS] +@group(0) @binding(3) var done: array; // [B] +@group(0) @binding(4) var tokens: array; // ring [T_max·B] +{{IF_SHORT}} +// Shortlisted lm_head: P carries LOCAL column indices; IDMAP maps them to +// vocab ids. Two masks stay in sync — \`bitmask\` (binding 2) is the +// LOCAL-space mask the fused gemm epilogue reads next step, \`gmask\` the +// vocab-space mask that survives routing changes (GEMV/full-vocab states +// created by compaction keep reading it). +@group(0) @binding(5) var IDMAP: array; // [N_short] +@group(0) @binding(6) var gmask: array; // [B·GMASK_WORDS] +{{/IF_SHORT}} + +const WG: u32 = {{WG}}u; +const EOS: u32 = {{EOS}}u; +const PAD: u32 = {{PAD}}u; +const MASK_WORDS: u32 = {{MASK_WORDS}}u; +{{IF_SHORT}} +const GMASK_WORDS: u32 = {{GMASK_WORDS}}u; +{{/IF_SHORT}} +// Finite f32 lowest (WGSL has no infinity literal). Real partials hold finite +// logits, so only all-OOB tile partials carry this value — and never win. +const NEG_MAX: f32 = -3.40282e38; + +var bestVal: array; +var bestIdx: array; + +@compute @workgroup_size({{WG}}) +fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { + // Uniform per workgroup (one workgroup per row) — safe early return. + if (wid.x >= params.B) { return; } + let b = wid.x; + let tid = lid.x; + let base = b * params.NT; + + var val: f32 = NEG_MAX; + var idx: u32 = 0xffffffffu; + for (var j = tid; j < params.NT; j = j + WG) { + let p = P[base + j]; + let v = bitcast(p.x); + // Strict >: ascending j = ascending vocab ranges, so equal-val ties keep + // the earliest partial = the lower vocab id. + if (v > val) { val = v; idx = p.y; } + } + bestVal[tid] = val; + bestIdx[tid] = idx; + workgroupBarrier(); + + // Tree reduce: higher val wins; exact-equal vals break to the LOWER idx. + for (var s = WG / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { + let ov = bestVal[tid + s]; + let oi = bestIdx[tid + s]; + if (ov > bestVal[tid] || (ov == bestVal[tid] && oi < bestIdx[tid])) { + bestVal[tid] = ov; + bestIdx[tid] = oi; + } + } + // Final iteration's barrier also orders the epilogue below after ALL + // threads' scans/reduction writes (see argmax_penalty.wgsl races note). + workgroupBarrier(); + } + + // Epilogue: single writer per row — identical to argmax_penalty.wgsl. + if (tid == 0u) { + let mbase = b * MASK_WORDS; + let wasDone = done[b]; + var tok = bestIdx[0]; +{{IF_SHORT}} + let lidx = tok; + tok = IDMAP[lidx]; +{{/IF_SHORT}} + if (wasDone == 1u) { + tok = PAD; + } else { + if (tok == EOS) { done[b] = 1u; } +{{IF_SHORT}} + bitmask[mbase + (lidx >> 5u)] = bitmask[mbase + (lidx >> 5u)] | (1u << (lidx & 31u)); + let gbase = b * GMASK_WORDS; + gmask[gbase + (tok >> 5u)] = gmask[gbase + (tok >> 5u)] | (1u << (tok & 31u)); +{{/IF_SHORT}} +{{IF_NOSHORT}} + bitmask[mbase + (tok >> 5u)] = bitmask[mbase + (tok >> 5u)] | (1u << (tok & 31u)); +{{/IF_NOSHORT}} + } + tokens[params.t * params.B + b] = tok; + } +} +`;function ar(e={}){const n=e.defines??{};return{t:e.t??"f32",outT:e.outT??e.t??"f32",wg:e.wg??64,bias:!!e.bias,silu:!!e.silu,wt:!!e.wt,sg:!!e.sg,immediate:!!e.immediate,defines:Object.fromEntries(Object.keys(n).sort().map(r=>[r,n[r]]))}}function ot(e,n={}){const{t:r,outT:t,wg:o,bias:a,silu:s,wt:u,sg:l,immediate:c,defines:f}=ar(n),i={T:r,OUT_T:t,WG:String(o),ENABLE_F16:r==="f16"||t==="f16"?"enable f16;":"",ENABLE_SG:l?"enable subgroups;":"",ENABLE_IMMEDIATE:c?"requires immediate_address_space;":"",PARAM_BINDING:c?"":"@group(0) @binding(0) ",PARAM_ADDRESS:c?"immediate":"uniform"},p={BIAS:!!a,SILU:!!s,WT:!!u,SG:!!l,NOSG:!l};for(const[w,d]of Object.entries(f))typeof d=="boolean"?p[w.toUpperCase()]=d:i[w.toUpperCase()]=String(d);let g=e;for(let w=null;w!==g;)w=g,g=g.replace(/\{\{IF_([A-Z0-9_]+)\}\}([\s\S]*?)\{\{\/IF_\1\}\}/g,(d,m,h)=>{if(!(m in p))throw new Error(`buildShader: unknown conditional {{IF_${m}}}`);return p[m]?h:""});return g=g.replace(/\{\{([A-Z0-9_/]+)\}\}/g,(w,d)=>{if(!(d in i))throw new Error(`buildShader: unresolved placeholder {{${d}}}`);return i[d]}),g}const Wn=new WeakMap,at=256,st=256*1024,it=64,Pn=new WeakMap,yn=new WeakMap,kn=new WeakSet;let ct=1;function _n(e){let n=yn.get(e);return n||(n=ct++,yn.set(e,n)),n}function Ke(e){let n=Pn.get(e);return n||(n={bindGroups:new Map,bindGroupLimit:at,activeUniformPools:0,uniformPoolOriginalBindGroupLimit:null,dummyStorage:null,uniformFrame:null,stats:{uniformBuffersCreated:0,uniformPoolBuffersCreated:0,uniformPoolBuffersDestroyed:0,uniformPoolFramesBegun:0,uniformPoolFramesFlushed:0,uniformPoolBlocks:0,uniformPoolBytes:0,uniformPoolBindGroupCacheHits:0,uniformPoolWarmBindGroupLookups:0,uniformPoolWarmBindGroupCacheHits:0,uniformPoolWarmBindGroupResets:0,uniformPoolGenerationInvalidations:0,uniformPoolCachePurges:0,dummyBuffersCreated:0,bindGroupsCreated:0,bindGroupCacheHits:0,bindGroupEvictions:0,bindGroupTargetedPurgeCalls:0,bindGroupTargetedPurges:0,immediateSets:0}},Pn.set(e,n)),n}function un(e){const n=Ke(e);return{...n.stats,bindGroupCacheSize:n.bindGroups.size,bindGroupCacheLimit:n.bindGroupLimit}}function sr(e,{B:n,immediate:r}={}){return!!e&&r===!1&&Number.isInteger(n)&&n>=1&&n<=it}function ir(e,{banks:n=2,bankBytes:r=st,alignment:t=e?.limits?.minUniformBufferOffsetAlignment??256}={}){if(!Number.isInteger(n)||n<1)throw new Error(`uniform pool banks must be a positive integer, got ${n}`);if(!Number.isInteger(t)||t<16||t%16!==0)throw new Error(`uniform pool alignment must be a positive multiple of 16, got ${t}`);if(!Number.isInteger(r)||r= alignment, got ${r}`);r=Math.ceil(r/t)*t;const o=Ke(e),a={uniformBuffersCreated:o.stats.uniformBuffersCreated,bindGroupsCreated:o.stats.bindGroupsCreated,bindGroupCacheHits:o.stats.bindGroupCacheHits,pooledBindGroupCacheHits:o.stats.uniformPoolBindGroupCacheHits,warmBindGroupLookups:o.stats.uniformPoolWarmBindGroupLookups,warmBindGroupCacheHits:o.stats.uniformPoolWarmBindGroupCacheHits,warmBindGroupResets:o.stats.uniformPoolWarmBindGroupResets,generationInvalidations:o.stats.uniformPoolGenerationInvalidations,cachePurges:o.stats.uniformPoolCachePurges},s={highWater:o.bindGroupLimit},u=Array.from({length:n},(d,m)=>{const h=e.createBuffer({label:`uniform params bank ${m}`,size:r,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});return kn.add(h),{buffer:h,cpu:new Uint32Array(r/4),cursor:0,blocks:0,busy:!1,flushed:!1,highWaterBytes:0,highWaterBlocks:0}});o.activeUniformPools===0&&(o.uniformPoolOriginalBindGroupLimit=o.bindGroupLimit),o.activeUniformPools++;const l=o.uniformPoolOriginalBindGroupLimit>0;o.stats.uniformPoolBuffersCreated+=u.length;const c=u.map(d=>_n(d.buffer)),f=()=>{let d=0;for(const m of[...o.bindGroups.keys()])c.some(h=>m.includes(`|${h}@`))&&(o.bindGroups.delete(m),o.stats.uniformPoolCachePurges++,d++);return d};let i=!1;const p=()=>{if(i)throw new Error("uniform pool is destroyed")},g=d=>{if(!Number.isInteger(d)||d<0||d>=u.length)throw new Error(`uniform pool bank ${d} out of range 0..${u.length-1}`);return u[d]},w={begin(d){if(p(),o.uniformFrame)throw new Error("uniform pool frame already active");const m=g(d);if(m.busy)throw new Error(`uniform pool bank ${d} reused while busy`);return m.cursor=0,m.blocks=0,m.busy=!0,m.flushed=!1,o.uniformFrame={pool:w,bank:m,index:d,alignment:t,bankBytes:r,cacheBanks:u.length,cacheEnabled:l,cacheLimitStats:s,warm:m.highWaterBlocks>0},o.stats.uniformPoolFramesBegun++,d},flush(){p();const d=o.uniformFrame;if(!d||d.pool!==w)throw new Error("uniform pool has no active frame to flush");const{bank:m,index:h}=d,b=Math.ceil(m.cursor/4)*4;return b>0&&e.queue.writeBuffer(m.buffer,0,m.cpu.buffer,m.cpu.byteOffset,b),m.flushed=!0,m.highWaterBytes=Math.max(m.highWaterBytes,b),m.highWaterBlocks=Math.max(m.highWaterBlocks,m.blocks),o.uniformFrame=null,o.stats.uniformPoolFramesFlushed++,o.stats.uniformPoolBlocks+=m.blocks,o.stats.uniformPoolBytes+=b,{bank:h,blocks:m.blocks,usedBytes:b}},abort(){p();const d=o.uniformFrame;if(!d||d.pool!==w)throw new Error("uniform pool has no active frame to abort");d.bank.cursor=0,d.bank.blocks=0,d.bank.busy=!1,d.bank.flushed=!1,o.uniformFrame=null},release(d){p();const m=g(d);if(o.uniformFrame?.bank===m)throw new Error(`uniform pool bank ${d} released while its frame is active`);if(!m.busy||!m.flushed)throw new Error(`uniform pool bank ${d} released before a flushed submission`);m.busy=!1,m.flushed=!1},invalidateBindings(){if(p(),o.uniformFrame?.pool===w)throw new Error("uniform pool generation invalidated while its frame is active");const d=u.findIndex(m=>m.busy);if(d>=0)throw new Error(`uniform pool bank ${d} is busy during generation invalidation`);return o.stats.uniformPoolGenerationInvalidations++,f()},snapshot(){return{banks:u.length,bankBytes:r,alignment:t,bindGroupCacheLimit:o.bindGroupLimit,bindGroupCacheLimitHighWater:s.highWater,busyBanks:u.filter(d=>d.busy).length,highWaterBytes:Math.max(0,...u.map(d=>d.highWaterBytes)),highWaterBlocks:Math.max(0,...u.map(d=>d.highWaterBlocks)),transientUniformBuffersCreated:o.stats.uniformBuffersCreated-a.uniformBuffersCreated,bindGroupsCreated:o.stats.bindGroupsCreated-a.bindGroupsCreated,bindGroupCacheHits:o.stats.bindGroupCacheHits-a.bindGroupCacheHits,pooledBindGroupCacheHits:o.stats.uniformPoolBindGroupCacheHits-a.pooledBindGroupCacheHits,warmBindGroupLookups:o.stats.uniformPoolWarmBindGroupLookups-a.warmBindGroupLookups,warmBindGroupCacheHits:o.stats.uniformPoolWarmBindGroupCacheHits-a.warmBindGroupCacheHits,warmBindGroupResets:o.stats.uniformPoolWarmBindGroupResets-a.warmBindGroupResets,generationInvalidations:o.stats.uniformPoolGenerationInvalidations-a.generationInvalidations,bindGroupsPurged:o.stats.uniformPoolCachePurges-a.cachePurges}},destroy(){if(!i){if(o.uniformFrame?.pool===w){const d=o.uniformFrame.bank;d.busy=!1,d.flushed=!1,o.uniformFrame=null}f();for(const d of u)kn.delete(d.buffer),d.busy=!1,d.flushed=!1,d.buffer.destroy();if(o.activeUniformPools--,o.activeUniformPools===0)for(o.bindGroupLimit=o.uniformPoolOriginalBindGroupLimit,o.uniformPoolOriginalBindGroupLimit=null;o.bindGroups.size>o.bindGroupLimit;)o.bindGroups.delete(o.bindGroups.keys().next().value),o.stats.bindGroupEvictions++;o.stats.uniformPoolBuffersDestroyed+=u.length,i=!0}}};return w}function le(e,n,r,t={}){let o=Wn.get(e);o||(o=new Map,Wn.set(e,o));const a=`${n}:${JSON.stringify(ar(t))}`;let s=o.get(a);if(s){if(s.source!==r)throw new Error(`pipeline cache key collision: ${a}`);return s.pipeline}const u=e.createShaderModule({label:a,code:ot(r,t)}),l=e.createComputePipeline({label:a,layout:"auto",compute:{module:u,entryPoint:"main"}});return o.set(a,{pipeline:l,source:r}),l}function Bn(e){return e.buffer?e:{buffer:e}}function ut(e,n,r){Ke(e).stats.uniformBuffersCreated++;const t=Math.max(16,Math.ceil(r.length*4/16)*16),o=e.createBuffer({label:n,size:t,usage:GPUBufferUsage.UNIFORM,mappedAtCreation:!0});return new Uint32Array(o.getMappedRange()).set(r),o.unmap(),o}function nn(e){const n=Ke(e);return n.dummyStorage||(n.dummyStorage=e.createBuffer({label:"shared dummy storage",size:4,usage:GPUBufferUsage.STORAGE}),n.stats.dummyBuffersCreated++),n.dummyStorage}function we(e,n,r,t){if(t)return{resource:null,values:Uint32Array.from(r),scratch:[]};const o=Ke(e),a=o.uniformFrame;if(a){const u=Math.max(16,Math.ceil(r.length*4/16)*16),l=Math.ceil(a.bank.cursor/a.alignment)*a.alignment,c=l+u;if(c>a.bankBytes)throw new Error(`uniform pool bank ${a.index} overflow: need ${c} bytes, cap ${a.bankBytes}`);return a.bank.cpu.fill(0,l/4,c/4),a.bank.cpu.set(r,l/4),a.bank.cursor=c,a.bank.blocks++,a.cacheEnabled&&(o.bindGroupLimit=Math.max(o.bindGroupLimit,a.bank.blocks*a.cacheBanks),a.cacheLimitStats.highWater=Math.max(a.cacheLimitStats.highWater,o.bindGroupLimit)),{resource:{buffer:a.bank.buffer,offset:l,size:u},values:null,scratch:[]}}const s=ut(e,n,r);return{resource:s,values:null,scratch:[s]}}function be(e,n){return e.resource?[e.resource,...n]:n}function lt(e,n,r){const t=n.map(o=>{const a=Bn(o);return`${_n(a.buffer)}@${a.offset??0}:${a.size??"*"}`});return`${r?"i":"u"}|p${_n(e)}|${t.join("|")}`}function pe(e,n,r,t,o,a=1,s=1,u=null){const l=u?1:0,c=Ke(r),f=!u&&t.length>0&&kn.has(Bn(t[0]).buffer);let i=null,p=null;if((u||f)&&c.bindGroupLimit>0){p=lt(n,t,!!u),i=c.bindGroups.get(p)??null;const g=f&&c.uniformFrame?.warm;g&&c.stats.uniformPoolWarmBindGroupLookups++,i?(c.bindGroups.delete(p),c.bindGroups.set(p,i),c.stats.bindGroupCacheHits++,f&&(c.stats.uniformPoolBindGroupCacheHits++,c.uniformFrame?.warm&&c.stats.uniformPoolWarmBindGroupCacheHits++)):g&&(c.uniformFrame.warm=!1,c.stats.uniformPoolWarmBindGroupResets++)}if(i||(i=r.createBindGroup({layout:n.getBindGroupLayout(0),entries:t.map((g,w)=>({binding:w+l,resource:Bn(g)}))}),c.stats.bindGroupsCreated++,p&&(c.bindGroups.set(p,i),c.bindGroups.size>c.bindGroupLimit&&(c.bindGroups.delete(c.bindGroups.keys().next().value),c.stats.bindGroupEvictions++))),e.setPipeline(n),u){if(typeof e.setImmediates!="function")throw new Error("WebGPU immediate shader selected but pass.setImmediates is unavailable");e.setImmediates(0,u),c.stats.immediateSets++}e.setBindGroup(0,i),e.dispatchWorkgroups(o,a,s)}function Ee(e,n,{x:r,w:t,b:o=null,y:a,M:s,K:u,N:l,storeKV:c=null,scales:f=null,fusedArgmax:i=null,splitK:p=null,flags:g={}}){if(g.tiled&&g.gemv)throw new Error("flags.tiled and flags.gemv are exclusive");if(g.tiled){if(c)throw new Error("storeKV requires flags.gemv");return ft(e,n,{x:r,w:t,b:o,y:a,M:s,K:u,N:l,scales:f,fusedArgmax:i,splitK:p,flags:g})}if(i)throw new Error("fusedArgmax requires flags.tiled");if(p)throw new Error("splitK requires flags.tiled");if(g.gemv)return dt(e,n,{x:r,w:t,b:o,y:a,M:s,K:u,N:l,storeKV:c,scales:f,flags:g});if(f)throw new Error("scales (wq8) requires flags.tiled or flags.gemv");if(c)throw new Error("storeKV requires flags.gemv");const w=g.wg??64,d=le(e,"gemm",Ur,{...g,bias:!!o}),m=we(e,"gemm dims",[s,u,l,0],!!g.immediate),h=[...m.scratch];let b=o;return b||(b=nn(e)),pe(n,d,e,be(m,[r,t,b,a]),Math.ceil(l/w),s,1,m.values),{pipeline:d,scratch:h}}function dt(e,n,{x:r,w:t,b:o,y:a,M:s,K:u,N:l,storeKV:c=null,scales:f=null,flags:i}){const p=!!i.wt,g=!!i.wq8;if(g&&(!p||!f||c))throw new Error("gemv wq8: needs wt layout and scales, excludes storeKV");if(p&&u%4!==0)throw new Error(`gemv wt requires K%4==0, got K=${u}`);if(!p&&l%4!==0)throw new Error(`gemv requires N%4==0, got N=${l}`);if(c&&l%3!==0)throw new Error("storeKV requires fused QKV (N=3·H·D)");const w=i.tk??16,d=i.tn??(p?8:4),m=p?i.mt??8:1,h=le(e,"gemm_gemv",Hr,{t:i.t,outT:i.outT,wg:w*d,bias:!!o,silu:i.silu,wt:p,immediate:!!i.immediate,sg:!!i.sg&&p,defines:{TK:w,TN:d,NWT:!p,STORE_KV:!!c,WQ8:g,WQF:!g,...p?{MT:m}:{}}}),b=we(e,"gemv dims",c?[s,u,l,0,c.t,c.Lmax,0,0]:[s,u,l,0],!!i.immediate),y=[...b.scratch];let _=o;_||(_=nn(e));const F=Math.ceil(p?l/d:l/(4*d)),A=p?Math.ceil(s/m):s,B=[r,t,_,a];return c&&B.push(c.kCache,c.vCache),g&&B.push(f),pe(n,h,e,be(b,B),F,A,1,b.values),{pipeline:h,scratch:y}}function Re(e,n,r=16){const t=Math.ceil(e/n/r)*r;return{KSL:t,nz:Math.ceil(e/t)}}function ft(e,n,{x:r,w:t,b:o,y:a,M:s,K:u,N:l,scales:c=null,fusedArgmax:f=null,splitK:i=null,flags:p}){const g=p.bm??64,w=p.bn??64,d=p.bkk??16;if(g%4!==0||w%4!==0)throw new Error(`gemm_tiled: BM/BN must be multiples of 4 (${g}, ${w})`);const m=!!p.sh16,h=!!p.dbuf;if(f&&(m||h))throw new Error("gemm_tiled2 fused argmax: sh16/dbuf unsupported (pVal aliases f32 Xs)");if(i){if(f)throw new Error("gemm_tiled2 splitK: exclusive with fusedArgmax");if(!(i.sk>=2))throw new Error(`gemm_tiled2 splitK: sk must be >= 2, got ${i.sk}`);if(o||p.silu)throw new Error("gemm_tiled2 splitK: pass bias/silu to dispatchGemmReduce, not the GEMM")}const b=f?g*(w/4)*4:0,y=m&&p.t==="f16"?2:4,_=(g+w)*d*y*(h?2:1)+b;if(_>16384)throw new Error(`gemm_tiled: shared memory ${_} bytes > 16384 limit`);if(f&&d256)throw new Error(`gemm_tiled: ${v} threads > 256 workgroup limit`);const k=B?le(e,"gemm_tiled2",$r,{t:p.t,outT:p.outT,wg:v,bias:!!o,silu:p.silu,immediate:!!p.immediate,wt:p.wt&&!F,defines:{BM:g,BN:w,BK:d,TM8:S===8,WNT:!p.wt&&!F,WQ8:F,WQF:!F,STORE_Y:!f&&!i,ARGMAX:!!f,SPLITK:!!i,NOSPLITK:!i,SH16:m,SH32:!m,DBUF:h,SBUF:!h,...i?{KSL:T.KSL}:{},...f?{PENALTY:Yn,MASK_WORDS:f.maskWords??ke}:{}}}):le(e,"gemm_tiled",jr,{t:p.t,outT:p.outT,wg:v,bias:!!o,silu:p.silu,wt:p.wt,immediate:!!p.immediate,defines:{BM:g,BN:w,BK:d}}),N=we(e,"gemm_tiled dims",[s,u,l,0],!!p.immediate),M=[...N.scratch];let E=o;E||(E=nn(e));const D=[r,t,E,f?.partials??i?.parts??a];return F&&D.push(c),f&&D.push(f.lbias,f.seen),pe(n,k,e,be(N,D),Math.ceil(l/w),Math.ceil(s/g),T?.nz??1,N.values),{pipeline:k,scratch:M}}function Rn(e,n,{parts:r,b:t=null,y:o,M:a,N:s,nz:u,storeKV:l=null,flags:c={}}){if(l&&s%3!==0)throw new Error("gemm_reduce storeKV requires fused QKV (N=3·H·D)");const f=c.wg??128,i=le(e,"gemm_reduce",Yr,{t:c.t,outT:c.outT,wg:f,bias:!!t,silu:!!c.silu,immediate:!!c.immediate,defines:{STORE_KV:!!l}}),p=we(e,"gemm_reduce dims",[a,s,u,l?.t??0,l?.Lmax??0,0,0,0],!!c.immediate),g=[...p.scratch];let w=t;w||(w=nn(e));const d=[r,w,o];return l&&d.push(l.kCache,l.vCache),pe(n,i,e,be(p,d),Math.ceil(a*s/f),1,1,p.values),{pipeline:i,scratch:g}}function pt({D4:e,t:n,qb:r=null,jb:t=null,qbAlign8:o=!1}){let a=r??Math.max(1,Math.min(16,Math.floor(256/e)));if(o&&r==null&&a*e%16!==0){for(;a>1&&a*e%16!==0;)a--;if(a*e%16!==0)throw new Error(`attention block: no QB <= 256 threads aligns D4=${e} to 16 threads`)}const s=n==="f16"?8:16,u=c=>a*e*16+2*c*e*s+a*c*4+3*a*4;let l=t??32;if(t==null)for(;l>8&&u(l)>16384;)l>>=1;return{QB:a,JB:l,shared:u(l)}}function Sn(e,n,{q:r,k:t,v:o,lens:a=null,y:s,B:u,M:l,L:c,lenMode:f,step:i=0,starts:p=null,qStride:g=ae*se,qOff:w=0,kvStride:d=ae*se,kOff:m=0,vOff:h=0,flags:b={}}){for(const[B,T]of[["qStride",g],["qOff",w],["kvStride",d],["kOff",m],["vOff",h],["HEAD_DIM",se]])if(T%4!==0)throw new Error(`attention: ${B}=${T} not vec4-aligned`);if(b.block){if(f!==1)throw new Error("attention block: lenMode must be 1");if(!a)throw new Error("attention block: lens buffer required");if(b.packed&&!p)throw new Error("attention block: packed needs a starts buffer");const B=se/4,{QB:T,JB:S,shared:v}=pt({D4:B,t:b.t,qb:b.qb??null,jb:b.jb??null,qbAlign8:!!b.qbAlign8}),k=T*B;if(k>256)throw new Error(`attention block: QB=${T} needs ${k} > 256 threads`);if(v>16384)throw new Error(`attention block: QB=${T} JB=${S} needs ${v}B shared > 16384`);const N=le(e,"attention_block",Vr,{t:b.t,immediate:!!b.immediate,defines:{H:ae,D:se,QB:T,JB:S,ATTN_SCALE:cn,Q_STRIDE:g,Q_OFF:w,KV_STRIDE:d,K_OFF:m,V_OFF:h,PACKED:!!b.packed,NOPACKED:!b.packed}}),M=we(e,"attn params",[u,l,c,f,i,0,0,0],!!b.immediate),E=[r,t,o,a,s];return b.packed&&E.push(p),pe(n,N,e,be(M,E),Math.ceil(l/T),ae,u,M.values),{pipeline:N,scratch:M.scratch}}const y=le(e,"attention",zr,{t:b.t,wg:b.wg??128,sg:!!b.sg,immediate:!!b.immediate,defines:{H:ae,D:se,SCORES_CAP:fn,ATTN_SCALE:cn,Q_STRIDE:g,Q_OFF:w,KV_STRIDE:d,K_OFF:m,V_OFF:h}}),_=we(e,"attn params",[u,l,c,f,i,0,0,0],!!b.immediate),F=[..._.scratch];let A=a;return A||(A=nn(e)),pe(n,y,e,be(_,[r,t,o,A,s]),u*l,ae,1,_.values),{pipeline:y,scratch:F}}function Tn(e,n,{x:r,r:t,gamma:o,beta:a,y:s,rows:u,flags:l={}}){const c=le(e,"add_ln",Qr,{t:l.t,wg:l.wg??256,sg:!!l.sg,immediate:!!l.immediate,defines:{D:R,EPS:pn}}),f=we(e,"add_ln params",[u,0,0,0],!!l.immediate);return pe(n,c,e,be(f,[r,t,o,a,s]),u,1,1,f.values),{pipeline:c,scratch:f.scratch}}function mt(e,n,{x:r,w:t,b:o,r:a,gamma:s,beta:u,y:l,M:c,K:f,N:i,flags:p={}}){if(f%4!==0||i%4!==0)throw new Error(`gemm_row_ln: K=${f}/N=${i} must be vec4-aligned`);if((f+i)*4+(p.wg??128)*4>16384)throw new Error(`gemm_row_ln: shared memory over budget at K=${f}, N=${i}`);const g=le(e,"gemm_row_ln",Xr,{t:p.t,wg:p.wg??128,sg:!!p.sg,immediate:!!p.immediate,defines:{KDIM:f,D:i,EPS:pn}}),w=we(e,"gemm_row_ln params",[c,0,0,0],!!p.immediate);return pe(n,g,e,be(w,[r,t,o,a,s,u,l]),c,1,1,w.values),{pipeline:g,scratch:w.scratch}}function cr(e,n,{ids:r,table:t,posEmbed:o,y:a,mode:s,nRows:u,step:l=0,batch:c,s:f=0,packed:i=!1,flags:p={}}){if(i&&(s!=="src"||f>65535))throw new Error(`embed: packed needs mode 'src' and s < 65536 (got ${s}, s=${f})`);const g=le(e,"embed",Jr,{t:p.t,wg:p.wg??224,immediate:!!p.immediate,defines:{D:R,EMBED_SCALE:Mn,SRC_IDS:s==="src",DECODE:s==="decode",DECODER_START:xe,PACKED:!!i,NOPACKED:!i}}),w=we(e,"embed params",[u,l,c,f],!!p.immediate);return pe(n,g,e,be(w,[r,t,o,a]),u,1,1,w.values),{pipeline:g,scratch:w.scratch}}function ur(e=256){const n=ae*se/4,r=Math.max(fe/4,n+e);return(2*n+r)*16+fn*4+e*4}function ht(e,n,{weights:r,layer:t,embed:o=!1,ring:a,kCache:s,vCache:u,crossKV:l,lens:c,x:f,B:i,t:p,S:g,kvCapacity:w=ne,flags:d={}}){if(r.dtype!=="f16")throw new Error("decoder mega: needs f16 weights");const m=ur(d.wg??256);if(m>16384)throw new Error(`decoder mega: shared memory ${m} bytes > 16384 limit at these dims`);const h=F=>{const A=r.tensors.get(F);if(!A)throw new Error(`decoder mega: missing tensor ${F}`);if(A.byteOffset%8!==0)throw new Error(`decoder mega: ${F} offset not vec4-aligned`);return A.byteOffset/8},b=F=>`dec.${t}.${F}`,y=le(e,"decoder_mega",nt,{t:"f16",wg:d.wg??256,sg:!!d.sg,immediate:!!d.immediate,defines:{EMBED:!!o,NOEMBED:!o,...o?{TABLE4:h("shared.weight"),POS4:h("pos_embed"),EMBED_SCALE:Mn,DECODER_START:xe}:{},H:ae,D:se,FFN4:fe/4,LMAX:w,SCORES_CAP:fn,ATTN_SCALE:cn,EPS:pn,QKVW4:h(b("self_qkv.weight")),QKVB4:h(b("self_qkv.bias")),OUTW4:h(b("self_out.weight")),OUTB4:h(b("self_out.bias")),LN1G4:h(b("ln1.weight")),LN1B4:h(b("ln1.bias")),CQW4:h(b("cross_q.weight")),CQB4:h(b("cross_q.bias")),COW4:h(b("cross_out.weight")),COB4:h(b("cross_out.bias")),LN2G4:h(b("ln2.weight")),LN2B4:h(b("ln2.bias")),FC1W4:h(b("fc1.weight")),FC1B4:h(b("fc1.bias")),FC2W4:h(b("fc2.weight")),FC2B4:h(b("fc2.bias")),LN3G4:h(b("ln3.weight")),LN3B4:h(b("ln3.bias"))}}),_=we(e,"mega params",[i,p,g,0],!!d.immediate);return pe(n,y,e,be(_,[r.buffer,a,s,u,l,c,f]),i,1,1,_.values),{pipeline:y,scratch:_.scratch}}function qn(e,n,{x:r,y:t,starts:o,lens:a,B:s,S:u,N:l,flags:c={}}){if(l%4!==0)throw new Error(`scatter_rows: N=${l} not vec4-aligned`);const f=le(e,"scatter_rows",Zr,{t:c.t,wg:c.wg??128,immediate:!!c.immediate}),i=we(e,"scatter_rows params",[s,u,l/4,0],!!c.immediate);return pe(n,f,e,be(i,[o,a,r,t]),s*u,1,1,i.values),{pipeline:f,scratch:i.scratch}}function gt(e,n){if(!Array.isArray(n))throw new Error("bind-group targeted purge needs an array of buffers");const r=Ke(e);r.stats.bindGroupTargetedPurgeCalls++;const t=new Set;for(const s of n){if(!s)continue;const u=s.buffer??s,l=yn.get(u);l&&t.add(l)}if(t.size===0)return 0;const o=[...t].map(s=>`|${s}@`);let a=0;for(const s of[...r.bindGroups.keys()])o.some(u=>s.includes(u))&&(r.bindGroups.delete(s),a++);return r.stats.bindGroupTargetedPurges+=a,a}function wt(e,n,{data:r,map:t,params:o,rowStrideU32:a,copyLenU32:s,flags:u={}}){if(!Number.isInteger(a)||a<1||!Number.isInteger(s)||s<1||s>a)throw new Error(`compact_gather: bad row shape stride=${a} copy=${s}`);const l=le(e,"compact_gather",et,{wg:u.wg??256});return pe(n,l,e,[o,t,r],1,1,1,null),{pipeline:l,scratch:[]}}function bt(e,n,{logits:r,bias:t,bitmask:o,done:a,tokens:s,B:u,t:l,short:c=null,flags:f={}}){const i=le(e,"argmax_penalty",rt,{wg:f.wg??256,immediate:!!f.immediate,defines:{V:c?.n??qe,EOS:Oe,PAD:en,PENALTY:Yn,MASK_WORDS:c?.maskWords??ke,SHORT:!!c,NOSHORT:!c,...c?{GMASK_WORDS:ke}:{}}}),p=we(e,"argmax params",[u,l,0,0],!!f.immediate),g=[r,t,o,a,s];return c&&g.push(c.idmap,c.gmask),pe(n,i,e,be(p,g),u,1,1,p.values),{pipeline:i,scratch:p.scratch}}function vt(e,n,{partials:r,bitmask:t,done:o,tokens:a,B:s,t:u,NT:l,short:c=null,flags:f={}}){const i=le(e,"argmax_reduce",tt,{wg:f.wg??256,immediate:!!f.immediate,defines:{EOS:Oe,PAD:en,MASK_WORDS:c?.maskWords??ke,SHORT:!!c,NOSHORT:!c,...c?{GMASK_WORDS:ke}:{}}}),p=we(e,"argmax_reduce params",[s,u,l,0],!!f.immediate),g=[r,t,o,a];return c&&g.push(c.idmap,c.gmask),pe(n,i,e,be(p,g),s,1,1,p.values),{pipeline:i,scratch:p.scratch}}async function Nn(e,n,{ids:r,lens:t,B:o,S:a},{recordInto:s=null,gemmOverrides:u=null,attnOverrides:l=null,packed:c="auto",sg:f=!1,attnQbAlign8:i=!1,retainEncOut:p=!0,splitSubmits:g=!1}={}){Dn(n.model,"runEncoder weights");const w=ae*se,d=3*w,m=2*w,{device:h}=e,b=n.dtype,y=b==="f16"?2:4,_=o*a;let F=0;for(let E=0;E