| export const H = 2048, |
| F = 6144, |
| VOCAB = 130560, |
| LAYERS = 42, |
| KV = 256, |
| HEADS = 16, |
| D = 128, |
| MAX_CHUNKS = 32, |
| PREFILL_BATCH = 32; |
| const STEP = `struct Step { position:u32, count:u32, chunk:u32, chunks:u32 }; @group(0) @binding(0) var<uniform> step:Step;\n`; |
| const bind = (index, type, access = "read") => |
| `@group(1) @binding(${index}) var<storage,${access}> ${type};\n`; |
| const Q8_CODE = ` |
| fn weightAt(i:u32)->f32 { |
| let block=i>>5u;let b=block*34u;let scale=unpack2x16float(weights[b>>2u])[(b&2u)>>1u]; |
| let at=b+2u+(i&31u);let word=weights[at>>2u];let q=bitcast<i32>(word<<((3u-(at&3u))*8u))>>24; |
| return f32(q)*scale; |
| } |
| fn weight4(block:u32,lane:u32)->vec4<f32>{ |
| let b=block*34u;let scale=unpack2x16float(weights[b>>2u])[(b&2u)>>1u]; |
| let a=b+2u+lane*4u;let shift=(a&3u)*8u;var word=weights[a>>2u]; |
| if(shift!=0u){word=(word>>shift)|(weights[(a>>2u)+1u]<<(32u-shift));} |
| let q=vec4<i32>(bitcast<i32>(word<<24u)>>24,bitcast<i32>(word<<16u)>>24,bitcast<i32>(word<<8u)>>24,bitcast<i32>(word)>>24); |
| return vec4<f32>(q)*scale; |
| } |
| `; |
| function sumDecl(subgroups, n = 256) { |
| return subgroups |
| ? `fn teamSum(v:f32,tid:u32)->f32{var x=v;for(var s=1u;s<32u;s<<=1u){x+=subgroupShuffleXor(x,s);}return x;}` |
| : `var<workgroup> red:array<f32,${n}>;fn teamSum(v:f32,tid:u32)->f32{red[tid]=v;workgroupBarrier();let lane=tid&31u;let base=tid-lane;for(var s=16u;s>0u;s>>=1u){if(lane<s){red[tid]+=red[tid+s];}workgroupBarrier();}let x=red[base];workgroupBarrier();return x;}`; |
| } |
| export function buildMiniCPMKernels({ |
| subgroups = false, |
| kv = "f16", |
| weightFormat = "q8", |
| draft = false, |
| smallWorkgroups = false, |
| } = {}) { |
| const Q8 = |
| weightFormat === "bf16" |
| ? `fn weightAt(i:u32)->f32{let w=weights[i>>1u];return bitcast<f32>(select(w<<16u,w&0xffff0000u,(i&1u)==1u));}fn weight4(block:u32,lane:u32)->vec4<f32>{let i=block*32u+lane*4u;let a=weights[i>>1u];let b=weights[(i>>1u)+1u];return vec4<f32>(bitcast<f32>(a<<16u),bitcast<f32>(a&0xffff0000u),bitcast<f32>(b<<16u),bitcast<f32>(b&0xffff0000u));}` |
| : Q8_CODE; |
| const prefix = (subgroups ? "enable subgroups;\n" : "") + STEP, |
| K = {}, |
| sum = sumDecl(subgroups); |
| const storageType = kv === "f32" ? "f32" : "u32"; |
| const readKV = (name) => |
| kv === "f32" |
| ? `return vec4<f32>(${name}[i],${name}[i+1u],${name}[i+2u],${name}[i+3u]);` |
| : kv === "f16" |
| ? `return vec4<f32>(unpack2x16float(${name}[i>>1u]),unpack2x16float(${name}[(i>>1u)+1u]));` |
| : kv === "q8_0" |
| ? `let b=i>>5u;let base=b*9u;let w=${name}[base+1u+((i&31u)>>2u)];return unpack4x8snorm(w)*bitcast<f32>(${name}[base]);` |
| : `let b=i>>5u;let base=b*5u;let w=${name}[base+1u+((i&31u)>>3u)]>>((i&4u)*4u);return (vec4<f32>(f32(w&15u),f32((w>>4u)&15u),f32((w>>8u)&15u),f32((w>>12u)&15u))-7.0)*bitcast<f32>(${name}[base]);`; |
| const readers = `fn readK(i:u32)->vec4<f32>{${readKV("keys")}} fn readV(i:u32)->vec4<f32>{${readKV("values")}}`; |
| K.embed = |
| prefix + |
| bind(0, "weights:array<u32>") + |
| bind(1, "tokens:array<u32>") + |
| bind(2, "x:array<f32>", "read_write") + |
| Q8 + |
| ` |
| @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g:vec3<u32>){if(g.x>=${H}u||g.y>=step.count){return;}x[g.y*${H}u+g.x]=weightAt(tokens[step.position+g.y]*${H}u+g.x);}`; |
| K.norm = |
| prefix + |
| bind(0, "x:array<f32>", "read_write") + |
| bind(1, "residual:array<f32>") + |
| bind(2, "weights:array<f32>") + |
| bind(3, "out:array<f32>", "read_write") + |
| ` |
| var<workgroup> squares:array<f32,256>; |
| @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l:vec3<u32>,@builtin(workgroup_id) g:vec3<u32>){ |
| let base=g.x*${H}u;var ss=0.0;for(var i=l.x;i<${H}u;i+=256u){let v=x[base+i]+residual[base+i];x[base+i]=v;ss+=v*v;}squares[l.x]=ss;workgroupBarrier(); |
| for(var s=128u;s>0u;s>>=1u){if(l.x<s){squares[l.x]+=squares[l.x+s];}workgroupBarrier();} |
| let scale=inverseSqrt(squares[0]/${H}.0+0.000001);for(var i=l.x;i<${H}u;i+=256u){out[base+i]=x[base+i]*scale*weights[i];}} |
| `; |
| const gemv = (columns, rows, gate = false) => |
| prefix + |
| bind(0, "input:array<f32>") + |
| bind(1, "weights:array<u32>") + |
| bind(2, "out:array<f32>", "read_write") + |
| Q8 + |
| sum + |
| ` |
| var<workgroup> sharedX:array<vec4<f32>,${columns / 4}>; |
| @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l:vec3<u32>,@builtin(workgroup_id) g:vec3<u32>){ |
| for(var i=l.x;i<${columns / 4}u;i+=256u){let b=i*4u;sharedX[i]=vec4<f32>(input[b],input[b+1u],input[b+2u],input[b+3u]);}workgroupBarrier(); |
| let row=g.x*8u+(l.x>>5u);let lane=l.x&31u;var acc=0.0;${gate ? "var other=0.0;" : ""} |
| for(var b=lane;b<${columns / 32}u;b+=32u){for(var w=0u;w<8u;w++){let v=sharedX[b*8u+w];acc+=dot(v,weight4(row*${columns / 32}u+b,w));${gate ? `other+=dot(v,weight4((row+${rows}u)*${columns / 32}u+b,w));` : ""}}} |
| let total=teamSum(acc,l.x);${gate ? "let up=teamSum(other,l.x);" : ""} |
| if(lane==0u&&row<${rows}u){out[row]=${gate ? "(total/(1.0+exp(-total)))*up" : "total"};}} |
| `; |
| const gemm = (columns, rows) => |
| prefix + |
| bind(0, "input:array<f32>") + |
| bind(1, "weights:array<u32>") + |
| bind(2, "out:array<f32>", "read_write") + |
| Q8 + |
| ` |
| var<workgroup> a:array<f32,1024>;var<workgroup> b:array<f32,1024>; |
| @compute @workgroup_size(16,16) fn main(@builtin(local_invocation_id) l:vec3<u32>,@builtin(workgroup_id) g:vec3<u32>){ |
| let tid=l.y*16u+l.x;let row=g.x*32u+l.x*2u;let token=g.y*32u+l.y*2u;var acc=vec4<f32>(0.0); |
| for(var kk=0u;kk<${columns}u;kk+=32u){ |
| for(var j=tid;j<1024u;j+=256u){let r=j>>5u;let c=j&31u;let t=g.y*32u+r;var av=0.0;if(t<step.count){av=input[t*${columns}u+kk+c];}a[j]=av;b[j]=weightAt((g.x*32u+r)*${columns}u+kk+c);}workgroupBarrier(); |
| for(var k=0u;k<32u;k++){let aa=vec2<f32>(a[(l.y*2u)*32u+k],a[(l.y*2u+1u)*32u+k]);let bb=vec2<f32>(b[(l.x*2u)*32u+k],b[(l.x*2u+1u)*32u+k]);acc+=vec4<f32>(aa.x*bb.x,aa.x*bb.y,aa.y*bb.x,aa.y*bb.y);}workgroupBarrier();} |
| if(token<step.count){out[token*${rows}u+row]=acc.x;out[token*${rows}u+row+1u]=acc.y;} |
| if(token+1u<step.count){out[(token+1u)*${rows}u+row]=acc.z;out[(token+1u)*${rows}u+row+1u]=acc.w;}} |
| `; |
| const block = (columns, rows) => |
| prefix + |
| bind(0, "input:array<vec4<f32>>") + |
| bind(1, "weights:array<u32>") + |
| bind(2, "out:array<f32>", "read_write") + |
| Q8 + |
| (subgroups |
| ? `fn teamSum4(v:vec4<f32>,tid:u32)->vec4<f32>{var x=v;for(var s=1u;s<32u;s<<=1u){x+=subgroupShuffleXor(x,s);}return x;}` |
| : `var<workgroup> red4:array<vec4<f32>,256>;fn teamSum4(v:vec4<f32>,tid:u32)->vec4<f32>{red4[tid]=v;workgroupBarrier();let lane=tid&31u;for(var s=16u;s>0u;s>>=1u){if(lane<s){red4[tid]+=red4[tid+s];}workgroupBarrier();}let r=red4[tid-lane];workgroupBarrier();return r;}`) + |
| ` |
| @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let row=g.x*8u+(l.x>>5u);let lane=l.x&31u;var a=vec4<f32>(0.0);var b=vec4<f32>(0.0); |
| for(var j=lane;j<${columns / 32}u;j+=32u){for(var w=0u;w<8u;w++){let v=weight4(row*${columns / 32}u+j,w);let at=j*8u+w; |
| ${Array.from({ length: 8 }, (_, i) => `if(${i}u<step.count){${i < 4 ? "a" : "b"}[${i % 4}]+=dot(v,input[${(i * columns) / 4}u+at]);}`).join("\n")} |
| }}let aa=teamSum4(a,l.x);let bb=teamSum4(b,l.x);if(lane==0u){${Array.from({ length: 8 }, (_, i) => `if(${i}u<step.count){out[${i * rows}u+row]=${i < 4 ? "aa" : "bb"}[${i % 4}];}`).join("\n")}}} |
| `; |
| for (const [name, cols, rows] of [ |
| ["qkv", H, H + 2 * KV], |
| ["output", H, H], |
| ["down", F, H], |
| ["head", H, VOCAB], |
| ]) { |
| K[name] = gemv(cols, rows); |
| if (name !== "head") K[name + "Batch"] = gemm(cols, rows); |
| K[name + "Block"] = block(cols, rows); |
| } |
| K.upGate = gemv(H, F, true); |
| K.upGateBatch = gemm(H, F * 2); |
| K.upGateBlock = block(H, F * 2); |
| if (draft) { |
| K.projectBatch = gemm(H * 5, H); |
| K.projectBlock = block(H * 5, H); |
| K.contextBatch = gemm(H, KV * 2); |
| K.contextBlock = block(H, KV * 2); |
| K.markov = gemv(256, VOCAB); |
| } |
| K.silu = |
| prefix + |
| bind(0, "input:array<f32>") + |
| bind(1, "out:array<f32>", "read_write") + |
| `@compute @workgroup_size(256)fn main(@builtin(global_invocation_id)g:vec3<u32>){if(g.x>=${F}u||g.y>=step.count){return;}let x=input[g.y*${F * 2}u+g.x];out[g.y*${F}u+g.x]=(x/(1.0+exp(-x)))*input[g.y*${F * 2}u+${F}u+g.x];}`; |
| K.rope = |
| prefix + |
| bind(0, "qkv:array<f32>") + |
| bind(1, "queries:array<f32>", "read_write") + |
| bind(2, "keyStage:array<f32>", "read_write") + |
| bind(3, "valueStage:array<f32>", "read_write") + |
| ` |
| @compute @workgroup_size(128)fn main(@builtin(global_invocation_id)g:vec3<u32>){let i=g.x*2u;let t=g.y;if(i>=${H + KV}u||t>=step.count){return;} |
| let base=t*${H + 2 * KV}u;let angle=f32(step.position+t)*pow(5000000.0,-f32(i%128u)/128.0);let c=cos(angle);let s=sin(angle);let a=qkv[base+i];let b=qkv[base+i+1u]; |
| if(i<${H}u){queries[t*${H}u+i]=a*c-b*s;queries[t*${H}u+i+1u]=a*s+b*c;} |
| else{let p=t*${KV}u+i-${H}u;keyStage[p]=a*c-b*s;keyStage[p+1u]=a*s+b*c;valueStage[p]=qkv[base+i+${KV}u];valueStage[p+1u]=qkv[base+i+${KV}u+1u];}} |
| `; |
| K.store = |
| prefix + |
| bind(0, "keyStage:array<f32>") + |
| bind(1, "valueStage:array<f32>") + |
| bind(2, `keys:array<${storageType}>`, "read_write") + |
| bind(3, `values:array<${storageType}>`, "read_write") + |
| ` |
| var<workgroup> ranges:array<f32,8>; |
| @compute @workgroup_size(8)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let i=g.x*32u+l.x*4u;let t=g.y;let which=g.z;let b=t*${KV}u+i;var v:vec4<f32>; |
| if(which==0u){v=vec4<f32>(keyStage[b],keyStage[b+1u],keyStage[b+2u],keyStage[b+3u]);}else{v=vec4<f32>(valueStage[b],valueStage[b+1u],valueStage[b+2u],valueStage[b+3u]);} |
| let slot=(step.position+t)*${KV}u+i; |
| ${ |
| kv === "f32" |
| ? `if(which==0u){keys[slot]=v.x;keys[slot+1u]=v.y;keys[slot+2u]=v.z;keys[slot+3u]=v.w;}else{values[slot]=v.x;values[slot+1u]=v.y;values[slot+2u]=v.z;values[slot+3u]=v.w;}` |
| : kv === "f16" |
| ? `if(which==0u){keys[slot>>1u]=pack2x16float(v.xy);keys[(slot>>1u)+1u]=pack2x16float(v.zw);}else{values[slot>>1u]=pack2x16float(v.xy);values[(slot>>1u)+1u]=pack2x16float(v.zw);}` |
| : ` |
| ranges[l.x]=max(max(abs(v.x),abs(v.y)),max(abs(v.z),abs(v.w)));workgroupBarrier();for(var s=4u;s>0u;s>>=1u){if(l.x<s){ranges[l.x]=max(ranges[l.x],ranges[l.x+s]);}workgroupBarrier();} |
| ${kv === "q8_0" ? `let scale=max(ranges[0],1e-12);let at=(slot>>5u)*9u;if(l.x==0u){if(which==0u){keys[at]=bitcast<u32>(scale);}else{values[at]=bitcast<u32>(scale);}}let word=pack4x8snorm(v/scale);if(which==0u){keys[at+1u+l.x]=word;}else{values[at+1u+l.x]=word;}` : `let scale=max(ranges[0]/7.0,1e-12);let q=vec4<u32>(clamp(round(v/scale)+7.0,vec4<f32>(0.0),vec4<f32>(14.0)));workgroupBarrier();ranges[l.x]=bitcast<f32>(q.x|(q.y<<4u)|(q.z<<8u)|(q.w<<12u));workgroupBarrier();let at=(slot>>5u)*5u;if(l.x==0u){if(which==0u){keys[at]=bitcast<u32>(scale);}else{values[at]=bitcast<u32>(scale);}}if((l.x&1u)==0u){let word=bitcast<u32>(ranges[l.x])|(bitcast<u32>(ranges[l.x+1u])<<16u);if(which==0u){keys[at+1u+(l.x>>1u)]=word;}else{values[at+1u+(l.x>>1u)]=word;}}`}` |
| }} |
| `; |
| |
| K.attention = |
| prefix + |
| bind(0, "queries:array<f32>") + |
| bind(1, `keys:array<${storageType}>`) + |
| bind(2, `values:array<${storageType}>`) + |
| bind(3, "partial:array<f32>", "read_write") + |
| bind(4, "md:array<vec2<f32>>", "read_write") + |
| readers + |
| sumDecl(subgroups, 128) + |
| ` |
| var<workgroup> accs:array<vec4<f32>,128>;var<workgroup> stats:array<vec2<f32>,4>; |
| @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let team=l.x>>5u;let lane=l.x&31u;let head=g.y;let token=g.z;let pos=step.position+token;let begin=g.x*step.chunk;let end=min(pos+1u,begin+step.chunk);let qb=token*${H}u+head*128u+lane*4u; |
| let q=vec4<f32>(queries[qb],queries[qb+1u],queries[qb+2u],queries[qb+3u]);var m=-3.402823e38;var den=0.0;var acc=vec4<f32>(0.0);let n=select(0u,(end-begin+3u)/4u,end>begin); |
| for(var j=0u;j<n;j++){let t=begin+j*4u+team;let valid=t<end;let i=select(0u,t*${KV}u+(head/8u)*128u+lane*4u,valid);let s=teamSum(dot(q,readK(i)),l.x)*0.08838834764831845; |
| if(valid){let nm=max(m,s);let a=exp(m-nm);let b=exp(s-nm);acc=acc*a+readV(i)*b;den=den*a+b;m=nm;}} |
| accs[l.x]=acc;if(lane==0u){stats[team]=vec2<f32>(m,den);}workgroupBarrier(); |
| let mm=max(max(stats[0].x,stats[1].x),max(stats[2].x,stats[3].x));var result=vec4<f32>(0.0);var total=0.0; |
| for(var j=0u;j<4u;j++){let scale=exp(stats[j].x-mm);result+=accs[j*32u+lane]*scale;total+=stats[j].y*scale;} |
| if(team==0u){let dst=((token*16u+head)*${MAX_CHUNKS}u+g.x)*128u+lane*4u;partial[dst]=result.x;partial[dst+1u]=result.y;partial[dst+2u]=result.z;partial[dst+3u]=result.w;if(lane==0u){md[(token*16u+head)*${MAX_CHUNKS}u+g.x]=vec2<f32>(mm,total);}}} |
| `; |
| K.combine = |
| prefix + |
| bind(0, "partial:array<f32>") + |
| bind(1, "md:array<vec2<f32>>") + |
| bind(2, "out:array<f32>", "read_write") + |
| ` |
| @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let base=(g.y*16u+g.x)*${MAX_CHUNKS}u;var m=-3.402823e38;for(var j=0u;j<step.chunks;j++){m=max(m,md[base+j].x);}var a=0.0;var den=0.0;for(var j=0u;j<step.chunks;j++){let v=md[base+j];let s=exp(v.x-m);a+=partial[(base+j)*128u+l.x]*s;den+=v.y*s;}out[g.y*${H}u+g.x*128u+l.x]=a/max(den,1e-20);} |
| `; |
| const sorter = `var<workgroup> vals:array<f32,1024>;var<workgroup> ids:array<u32,1024>; |
| fn sort(tid:u32){for(var k=2u;k<=1024u;k<<=1u){for(var j=k>>1u;j>0u;j>>=1u){for(var i=tid;i<1024u;i+=256u){let other=i^j;if(other>i){let descending=(i&k)==0u;let greater=vals[i]>vals[other]||(vals[i]==vals[other]&&ids[i]<ids[other]);if(greater!=descending){let v=vals[i];vals[i]=vals[other];vals[other]=v;let id=ids[i];ids[i]=ids[other];ids[other]=id;}}}workgroupBarrier();}}} |
| `; |
| for (const keep of [64, 256]) { |
| K["sortFirst" + keep] = |
| prefix + |
| bind(0, "logits:array<f32>") + |
| bind(1, "out:array<vec2<u32>>", "read_write") + |
| bind(2, "tokens:array<u32>") + |
| `@group(1) @binding(3) var<uniform> penalty:vec4<f32>;` + |
| sorter + |
| ` |
| @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){for(var i=l.x;i<1024u;i+=256u){let id=g.x*1024u+i;var v=-3.402823e38;if(id<${VOCAB}u){v=logits[id];if(penalty.x!=1.0){let start=select(0u,step.position-63u,step.position>63u);for(var p=start;p<=step.position;p++){if(tokens[p]==id){v=select(v/penalty.x,v*penalty.x,v<0.0);break;}}}}vals[i]=v;ids[i]=id;}workgroupBarrier();sort(l.x);if(l.x<${keep}u){out[g.x*${keep}u+l.x]=vec2<u32>(bitcast<u32>(vals[l.x]),ids[l.x]);}} |
| `; |
| K["sortMerge" + keep] = |
| prefix + |
| bind(0, "input:array<vec2<u32>>") + |
| bind(1, "out:array<vec2<u32>>", "read_write") + |
| `@group(1) @binding(2) var<uniform> size:vec4<u32>;` + |
| sorter + |
| ` |
| @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){for(var i=l.x;i<1024u;i+=256u){let at=g.x*1024u+i;var v=vec2<u32>(bitcast<u32>(-3.402823e38),${VOCAB}u);if(at<size.x){v=input[at];}vals[i]=bitcast<f32>(v.x);ids[i]=v.y;}workgroupBarrier();sort(l.x);if(l.x<${keep}u){out[g.x*${keep}u+l.x]=vec2<u32>(bitcast<u32>(vals[l.x]),ids[l.x]);}} |
| `; |
| } |
| K.saveCandidates = |
| prefix + |
| bind(0, "input:array<vec2<u32>>") + |
| bind(1, "out:array<vec2<u32>>", "read_write") + |
| `@compute @workgroup_size(64)fn main(@builtin(global_invocation_id)g:vec3<u32>){out[g.x]=input[g.x];}`; |
| K.sample = |
| prefix + |
| bind(0, "sorted:array<vec2<u32>>") + |
| bind(1, "tokens:array<u32>", "read_write") + |
| ` |
| struct Sampler{temperature:f32,topP:f32,topK:u32,seed:u32};@group(1)@binding(2)var<uniform> cfg:Sampler; |
| fn hash(x:u32)->u32{var v=x*747796405u+2891336453u;v=((v>>((v>>28u)+4u))^v)*277803737u;return(v>>22u)^v;} |
| @compute @workgroup_size(1)fn main(){var token=sorted[0].y;if(cfg.temperature>0.0){let mx=bitcast<f32>(sorted[0].x);var sum=0.0;for(var i=0u;i<cfg.topK;i++){sum+=exp((bitcast<f32>(sorted[i].x)-mx)/cfg.temperature);}var retained=0.0;var count=0u;for(var i=0u;i<cfg.topK;i++){retained+=exp((bitcast<f32>(sorted[i].x)-mx)/cfg.temperature);count++;if(retained>=cfg.topP*sum){break;}}let random=(f32(hash(cfg.seed^(step.position*2654435761u)))+0.5)/4294967296.0;let threshold=random*retained;var cum=0.0;for(var i=0u;i<count;i++){cum+=exp((bitcast<f32>(sorted[i].x)-mx)/cfg.temperature);token=sorted[i].y;if(cum>=threshold){break;}}}tokens[step.position+1u]=token;} |
| `; |
| if (smallWorkgroups) { |
| K.down = K.down |
| .replace(/var<workgroup> sharedX:[^;]+;/, "") |
| .replace(/for\(var i=l.x;i<1536u;[\s\S]*?workgroupBarrier\(\);/, "") |
| .replace( |
| "let v=sharedX[b*8u+w];", |
| "let at=(b*8u+w)*4u;let v=vec4<f32>(input[at],input[at+1u],input[at+2u],input[at+3u]);", |
| ); |
| } |
| return K; |
| } |
|
|
| |
| export function shardVocabularyKernel(code, embedding = false) { |
| const binding = embedding ? 0 : 1; |
| code = code.replace(/@binding\((\d+)\)/g, (match, n) => |
| Number(n) > binding ? `@binding(${Number(n) + 3})` : match, |
| ); |
| const declaration = `@group(1) @binding(${binding}) var<storage,read> weights:array<u32>;`; |
| code = code.replace( |
| declaration, |
| Array.from( |
| { length: 4 }, |
| (_, i) => |
| `@group(1) @binding(${binding + i}) var<storage,read> weights${i}:array<u32>;`, |
| ).join("\n") + |
| ` |
| fn weightWord(i:u32)->u32{let part=i/17756160u;let at=i%17756160u;switch part{case 0u:{return weights0[at];}case 1u:{return weights1[at];}case 2u:{return weights2[at];}default:{return weights3[at];}}}`, |
| ); |
| return code.replace(/weights\[([^\]]+)\]/g, "weightWord($1)"); |
| } |
|
|
