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!==1 || manifest.codec!=='chess-square-triple-v1' || manifest.context!==1024 || 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(history) {
31    const inputs=this.inputs([history],0);
32    for(let i=0;i<24;i++)inputs[`past_${i}`]=new ort.Tensor('float32',new Float32Array(0),[1,12,0,64]);
33    let out;
34    try {out=await this.session.run(inputs);} finally {Object.values(inputs).forEach(t=>t.dispose());}
35    const logits=Array.from(out.logits.data); out.logits.dispose();
36    return {logits,length:history.length,cache:Array.from({length:24},(_,i)=>out[`present_${i}`]),dispose(){Object.values(out).filter(t=>t!==out.logits).forEach(t=>t.dispose());}};
37  }
38  inputs(rows,start) {
39    const shape=[rows.length,rows[0].length];
40    return {input_ids:new ort.Tensor('int64',BigInt64Array.from(rows.flat(),BigInt),shape),
41      position_ids:new ort.Tensor('int64',BigInt64Array.from(rows.flatMap(r=>r.map((_,i)=>start+i)),BigInt),shape)};
42  }
43  async branch(rows,base) {
44    this.peakOwnedCacheBytes=Math.max(this.peakOwnedCacheBytes,24*12*base.length*64*4*(rows.length+1));
45    const inputs=this.inputs(rows,base.length),buffers=[];const device=ort.env.webgpu.device;
46    try {
47      const encoder=device.createCommandEncoder();
48      for(let i=0;i<24;i++) {
49        const size=12*base.length*64*4;
50        const buffer=device.createBuffer({size:size*rows.length,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_DST|GPUBufferUsage.COPY_SRC});buffers.push(buffer);
51        for(let b=0;b<rows.length;b++)encoder.copyBufferToBuffer(base.cache[i].gpuBuffer,0,buffer,b*size,size);
52        inputs[`past_${i}`]=ort.Tensor.fromGpuBuffer(buffer,{dataType:'float32',dims:[rows.length,12,base.length,64]});
53      }
54      device.queue.submit([encoder.finish()]);
55      const out=await this.session.run(inputs,['logits']);
56      try {const data=out.logits.data;return rows.map((_,i)=>Array.from(data.slice(i*128,(i+1)*128)));} finally {out.logits.dispose();}
57    } finally {Object.values(inputs).forEach(t=>t.dispose());buffers.forEach(b=>b.destroy());}
58  }
59  async dispose(){await this.session.release();}
60}
61