runner.mjs

1import * as ort from 'onnxruntime-web/webgpu';
2export class Runner {
3  static async create(manifest, progress) {
4    if(!navigator.gpu) throw Error('Desktop WebGPU is required. Try a current desktop Chrome or Edge browser.');
5    const adapter=await navigator.gpu.requestAdapter();
6    if(!adapter) throw Error('No WebGPU adapter is available. Enable hardware acceleration and restart your browser.');
7    if(manifest.schema!==2 || manifest.codec!=='chess-state-history16-v2' || manifest.context!==122 || manifest.vocabularySize!==512 || manifest.promptPositions!==120 || manifest.dtype!=='float32') throw Error('Unsupported checkpoint manifest');
8    async function download(files,size,onProgress=()=>{}) {
9      const bytes=new Uint8Array(size);let offset=0;
10      for(const file of files){const response=await fetch(new URL(file,manifest.url));if(!response.ok)throw Error(`Download failed (${response.status})`);
11        const reader=response.body.getReader();for(;;){const {done,value}=await reader.read();if(done)break;if(offset+value.length>size)throw Error('Download size mismatch');bytes.set(value,offset);offset+=value.length;onProgress(offset/size);}}
12      if(offset!==size)throw Error('Incomplete download');return bytes;
13    }
14    const bytes=await download(manifest.chunks||[manifest.model],manifest.bytes,progress);
15    const hash=Array.from(new Uint8Array(await crypto.subtle.digest('SHA-256',bytes)),b=>b.toString(16).padStart(2,'0')).join('');
16    if(hash!==manifest.sha256)throw Error('Model integrity check failed');
17    ort.env.wasm.numThreads=1;
18    ort.env.wasm.wasmPaths=new URL('./',location.href).href;
19    if(manifest.runtime){
20      const runtime=await download(manifest.runtime.chunks,manifest.runtime.bytes);
21      const runtimeHash=Array.from(new Uint8Array(await crypto.subtle.digest('SHA-256',runtime)),b=>b.toString(16).padStart(2,'0')).join('');
22      if(runtimeHash!==manifest.runtime.sha256)throw Error('Runtime integrity check failed');
23      ort.env.wasm.wasmBinary=runtime;
24    }
25    const locations=Object.fromEntries(Array.from({length:24},(_,i)=>[`present_${i}`,'gpu-buffer']));
26    const session=await ort.InferenceSession.create(bytes,{executionProviders:['webgpu'],preferredOutputLocation:locations});
27    const runner=new Runner(session);runner.deviceInfo={vendor:adapter.info?.vendor,architecture:adapter.info?.architecture,maxBufferSize:adapter.limits.maxBufferSize,maxStorageBufferBindingSize:adapter.limits.maxStorageBufferBindingSize};return runner;
28  }
29  constructor(session) {this.session=session;this.peakOwnedCacheBytes=0;}
30  async prefill(prompt) {
31    const history=prompt.input_ids;
32    const inputs=this.inputs([history],0);
33    inputs.attention_mask=new ort.Tensor('int64',BigInt64Array.from(prompt.attention_mask,BigInt),[1,120]);
34    for(let i=0;i<24;i++)inputs[`past_${i}`]=new ort.Tensor('float32',new Float32Array(0),[1,12,0,64]);
35    let out;
36    try {out=await this.session.run(inputs);} finally {Object.values(inputs).forEach(t=>t.dispose());}
37    const logits=Array.from(out.logits.data); out.logits.dispose();
38    return {logits,length:history.length,mask:prompt.attention_mask,cache:Array.from({length:24},(_,i)=>out[`present_${i}`]),dispose(){Object.values(out).filter(t=>t!==out.logits).forEach(t=>t.dispose());}};
39  }
40  inputs(rows,start) {
41    const shape=[rows.length,rows[0].length];
42    return {input_ids:new ort.Tensor('int64',BigInt64Array.from(rows.flat(),BigInt),shape),
43      position_ids:new ort.Tensor('int64',BigInt64Array.from(rows.flatMap(r=>r.map((_,i)=>start+i)),BigInt),shape)};
44  }
45  async branch(rows,base) {
46    this.peakOwnedCacheBytes=Math.max(this.peakOwnedCacheBytes,24*12*base.length*64*4*(rows.length+1));
47    const inputs=this.inputs(rows,base.length),buffers=[];const device=ort.env.webgpu.device;
48    try {
49      inputs.attention_mask=new ort.Tensor('int64',BigInt64Array.from(rows.flatMap(r=>[...base.mask,...r.map(()=>1)]),BigInt),[rows.length,base.length+rows[0].length]);
50      const encoder=device.createCommandEncoder();
51      for(let i=0;i<24;i++) {
52        const size=12*base.length*64*4;
53        const buffer=device.createBuffer({size:size*rows.length,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_DST|GPUBufferUsage.COPY_SRC});buffers.push(buffer);
54        for(let b=0;b<rows.length;b++)encoder.copyBufferToBuffer(base.cache[i].gpuBuffer,0,buffer,b*size,size);
55        inputs[`past_${i}`]=ort.Tensor.fromGpuBuffer(buffer,{dataType:'float32',dims:[rows.length,12,base.length,64]});
56      }
57      device.queue.submit([encoder.finish()]);
58      const out=await this.session.run(inputs,['logits']);
59      try {const data=out.logits.data;return rows.map((_,i)=>Array.from(data.slice(i*512,(i+1)*512)));} finally {out.logits.dispose();}
60    } finally {Object.values(inputs).forEach(t=>t.dispose());buffers.forEach(b=>b.destroy());}
61  }
62  async dispose(){await this.session.release();}
63}
64