| import { buildMiniCPMKernels, H, KV, VOCAB } from "./minicpm-kernels.js"; |
| export function buildDSparkKernels(options) { |
| const kernels = buildMiniCPMKernels({ |
| ...options, |
| weightFormat: "bf16", |
| draft: true, |
| }); |
| kernels.attention = kernels.attention.replace( |
| "let pos=step.position+token;", |
| "let pos=step.position+step.count-1u;", |
| ); |
| const pre = `${options.subgroups ? "enable subgroups;" : ""}\nstruct Step{position:u32,count:u32,chunk:u32,chunks:u32};@group(0)@binding(0)var<uniform>step:Step;\n`; |
| const bind = (i, name, type, mode = "read") => |
| `@group(1)@binding(${i})var<storage,${mode}>${name}:array<${type}>;\n`; |
| kernels.capture = |
| pre + |
| bind(0, "x", "f32") + |
| bind(1, "residual", "f32") + |
| bind(2, "features", "f32", "read_write") + |
| "@group(1)@binding(3)var<uniform>slot:vec4<u32>;" + |
| ` |
| @compute @workgroup_size(256)fn main(@builtin(global_invocation_id)g:vec3<u32>){if(g.x<2048u&&g.y<step.count){features[g.y*10240u+slot.x*2048u+g.x]=x[g.y*2048u+g.x]+residual[g.y*2048u+g.x];}}`; |
| const normRope = (context) => |
| pre + |
| bind(0, "input", "f32") + |
| bind(1, "qNorm", "f32") + |
| bind(2, "kNorm", "f32") + |
| bind(3, "queries", "f32", "read_write") + |
| bind(4, "keys", "f32", "read_write") + |
| bind(5, "values", "f32", "read_write") + |
| ` |
| var<workgroup> red:array<f32,128>;var<workgroup> normalized:array<f32,128>; |
| @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let d=l.x;let head=g.x;let t=g.y;let isQ=${context ? "false" : "head<16u"};let i=t*${context ? 512 : 2560}u+head*128u+d;let raw=input[i];red[d]=raw*raw;workgroupBarrier();for(var s=64u;s>0u;s>>=1u){if(d<s){red[d]+=red[d+s];}workgroupBarrier();}var weight=kNorm[d];if(isQ){weight=qNorm[d];}normalized[d]=raw*inverseSqrt(red[0]/128.0+0.000001)*weight;workgroupBarrier(); |
| let theta=f32(step.position+t)*pow(5000000.0,-f32(d%64u)/64.0);let pair=select(d-64u,d+64u,d<64u);let rotated=normalized[d]*cos(theta)+select(normalized[pair],-normalized[pair],d<64u)*sin(theta); |
| if(isQ){queries[t*2048u+head*128u+d]=rotated;}else{let kh=${context ? "head" : "head-16u"};let out=t*256u+kh*128u+d;keys[out]=rotated;values[out]=input[t*${context ? 512 : 2560}u+${context ? 256 : 2304}u+kh*128u+d];}} |
| `; |
| kernels.noiseNormRope = normRope(false); |
| kernels.contextNormRope = normRope(true); |
| kernels.noise = |
| pre + |
| bind(0, "source", "u32") + |
| bind(1, "tokens", "u32", "read_write") + |
| `@compute @workgroup_size(8)fn main(@builtin(local_invocation_id)l:vec3<u32>){if(l.x<step.count){tokens[step.position+l.x]=select(75982u,source[step.position],l.x==0u);}}`; |
| kernels.markovEmbed = |
| pre + |
| bind(0, "weights", "u32") + |
| bind(1, "tokens", "u32") + |
| bind(2, "out", "f32", "read_write") + |
| `@compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>){let i=tokens[step.position]*256u+l.x;let w=weights[i>>1u];out[l.x]=bitcast<f32>(select(w<<16u,w&0xffff0000u,(i&1u)==1u));}`; |
| kernels.addBias = |
| pre + |
| bind(0, "base", "f32") + |
| bind(1, "bias", "f32") + |
| bind(2, "out", "f32", "read_write") + |
| `@compute @workgroup_size(256)fn main(@builtin(global_invocation_id)g:vec3<u32>){if(g.x<${VOCAB}u){out[g.x]=base[g.x]+bias[g.x];}}`; |
| kernels.argmax = |
| pre + |
| bind(0, "logits", "f32") + |
| bind(1, "out", "vec2<u32>", "read_write") + |
| ` |
| var<workgroup> val:array<f32,256>;var<workgroup> ids:array<u32,256>; |
| @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){var v=-3.402823e38;var id=0xffffffffu;for(var i=g.x*256u+l.x;i<${VOCAB}u;i+=32768u){let x=logits[i];if(x>v||(x==v&&i<id)){v=x;id=i;}}val[l.x]=v;ids[l.x]=id;workgroupBarrier();for(var s=128u;s>0u;s>>=1u){if(l.x<s){let o=l.x+s;if(val[o]>val[l.x]||(val[o]==val[l.x]&&ids[o]<ids[l.x])){val[l.x]=val[o];ids[l.x]=ids[o];}}workgroupBarrier();}if(l.x==0u){out[g.x]=vec2<u32>(bitcast<u32>(val[0]),ids[0]);}} |
| `; |
| kernels.argmaxFinal = |
| pre + |
| bind(0, "input", "vec2<u32>") + |
| bind(1, "tokens", "u32", "read_write") + |
| ` |
| var<workgroup> val:array<f32,128>;var<workgroup> ids:array<u32,128>; |
| @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>){val[l.x]=bitcast<f32>(input[l.x].x);ids[l.x]=input[l.x].y;workgroupBarrier();for(var s=64u;s>0u;s>>=1u){if(l.x<s){let o=l.x+s;if(val[o]>val[l.x]||(val[o]==val[l.x]&&ids[o]<ids[l.x])){val[l.x]=val[o];ids[l.x]=ids[o];}}workgroupBarrier();}if(l.x==0u){tokens[step.position+1u]=ids[0];}} |
| `; |
| return kernels; |
| } |
|
|
