| |
| |
| local Encoding = game:GetService('EncodingService') |
| local Http = game:GetService('HttpService') |
| local Engine = {} |
| local D,H,KV,HD,FF,LAYERS,VOCAB = 576,9,3,64,1536,30,49152 |
| local EPS,THETA = 1e-5,100000 |
| local rd,wr,ri = buffer.readf32,buffer.writef32,buffer.readi8 |
| local function vec(n) return buffer.create(n*4) end |
| local lastYield=os.clock() |
| local cancelCheck=nil |
| local function checkpoint(force) |
| if force or os.clock()-lastYield>0.012 then |
| task.wait();lastYield=os.clock() |
| if cancelCheck and cancelCheck() then error('CANCELLED',0) end |
| end |
| end |
| local function matvec(w,x,out) |
| local n,m,b=w.rows,w.cols,w.data |
| local so=n*m |
| for r=0,n-1 do |
| local a=r*m;local s=0 |
| for g=0,m/64-1 do |
| local partial=0;local start=g*64 |
| for c=start,start+63,4 do |
| partial+=ri(b,a+c)*rd(x,c*4)+ri(b,a+c+1)*rd(x,(c+1)*4)+ri(b,a+c+2)*rd(x,(c+2)*4)+ri(b,a+c+3)*rd(x,(c+3)*4) |
| end |
| s+=partial*rd(b,so+(r*m/64+g)*4) |
| end |
| wr(out,r*4,s) |
| if r%128==127 then checkpoint(false) end |
| end |
| end |
| local function norm(x,w,out) |
| local s=0 |
| for i=0,D-1 do local v=rd(x,i*4);s+=v*v end |
| local inv=1/math.sqrt(s/D+EPS) |
| for i=0,D-1 do wr(out,i*4,rd(x,i*4)*inv*rd(w.data,i*4)) end |
| end |
| local function add(x,b) |
| for i=0,D-1 do wr(x,i*4,rd(x,i*4)+rd(b,i*4)) end |
| end |
| function Engine.load(folder,onProgress) |
| local manifest=Http:JSONDecode(folder.Manifest.Value) |
| local weights={};local names={} |
| for name in manifest.tensors do table.insert(names,name) end |
| table.sort(names) |
| for index,name in names do |
| local meta=manifest.tensors[name];local node=folder[meta.key] |
| local chunks=node:GetChildren();table.sort(chunks,function(a,b) return a.Name<b.Name end) |
| local data=buffer.create(meta.bytes);local off=0 |
| for _,chunk in chunks do |
| local decoded=Encoding:Base64Decode(buffer.fromstring(chunk.Value)) |
| buffer.copy(data,off,decoded);off+=buffer.len(decoded) |
| end |
| assert(off==meta.bytes,'Corrupt tensor '..name) |
| if meta.kind=='q4g64' then |
| local count=meta.shape[1]*meta.shape[2] |
| local expanded=buffer.create(count+count/64*4) |
| for j=0,count/2-1 do |
| local packed=buffer.readu8(data,j) |
| buffer.writei8(expanded,j*2,bit32.band(packed,15)-8) |
| buffer.writei8(expanded,j*2+1,bit32.rshift(packed,4)-8) |
| if j%131072==131071 then checkpoint(false) end |
| end |
| buffer.copy(expanded,count,data,count/2,count/64*4) |
| data=expanded |
| end |
| weights[name]={data=data,rows=meta.shape[1],cols=meta.shape[2]} |
| node:Destroy() |
| if onProgress then onProgress(index/#names) end |
| checkpoint(true) |
| end |
| Engine.weights=weights;Engine.manifest=manifest |
| Engine.scratch={x=vec(D),n=vec(D),q=vec(D),k=vec(KV*HD),v=vec(KV*HD),a=vec(D),o=vec(D),gate=vec(FF),up=vec(FF),logits=vec(VOCAB)} |
| return Engine |
| end |
| function Engine.newSession(maxContext) |
| local s={position=0,maxContext=maxContext or 256,keys={},values={}} |
| for l=1,LAYERS do s.keys[l]=vec(s.maxContext*KV*HD);s.values[l]=vec(s.maxContext*KV*HD) end |
| return s |
| end |
| function Engine.forward(session,token,project,shouldCancel) |
| cancelCheck=shouldCancel;checkpoint(false) |
| local pos=session.position |
| assert(pos<session.maxContext,'Context limit reached') |
| local w=Engine.weights;local t=Engine.scratch;local x=t.x |
| local emb=w['model.embed_tokens.weight'] |
| for i=0,D-1 do |
| local scale=rd(emb.data,emb.rows*D+(token*D/64+math.floor(i/64))*4) |
| wr(x,i*4,ri(emb.data,token*D+i)*scale) |
| end |
| for layer=0,LAYERS-1 do |
| local p='model.layers.'..layer..'.';local att=p..'self_attn.' |
| norm(x,w[p..'input_layernorm.weight'],t.n) |
| matvec(w[att..'q_proj.weight'],t.n,t.q) |
| matvec(w[att..'k_proj.weight'],t.n,t.k) |
| matvec(w[att..'v_proj.weight'],t.n,t.v) |
| for h=0,H-1 do |
| for j=0,HD/2-1 do |
| local angle=pos/(THETA^(2*j/HD));local co,si=math.cos(angle),math.sin(angle) |
| local a=(h*HD+j)*4;local b=a+HD/2*4;local v1,v2=rd(t.q,a),rd(t.q,b) |
| wr(t.q,a,v1*co-v2*si);wr(t.q,b,v2*co+v1*si) |
| if h<KV then |
| v1,v2=rd(t.k,a),rd(t.k,b);wr(t.k,a,v1*co-v2*si);wr(t.k,b,v2*co+v1*si) |
| end |
| end |
| end |
| local kc,vc=session.keys[layer+1],session.values[layer+1] |
| buffer.copy(kc,pos*KV*HD*4,t.k);buffer.copy(vc,pos*KV*HD*4,t.v) |
| for h=0,H-1 do |
| local kh=math.floor(h/(H/KV));local scores=table.create(pos+1);local peak=-math.huge |
| for step=0,pos do |
| local base=(step*KV*HD+kh*HD)*4;local qbase=h*HD*4;local dot=0 |
| for j=0,HD-1 do dot+=rd(t.q,qbase+j*4)*rd(kc,base+j*4) end |
| dot/=8;scores[step+1]=dot;peak=math.max(peak,dot) |
| end |
| local total=0 |
| for step=1,pos+1 do local e=math.exp(scores[step]-peak);scores[step]=e;total+=e end |
| for j=0,HD-1 do |
| local sum=0 |
| for step=0,pos do sum+=scores[step+1]/total*rd(vc,(step*KV*HD+kh*HD+j)*4) end |
| wr(t.a,(h*HD+j)*4,sum) |
| end |
| end |
| matvec(w[att..'o_proj.weight'],t.a,t.o);add(x,t.o) |
| norm(x,w[p..'post_attention_layernorm.weight'],t.n) |
| matvec(w[p..'mlp.gate_proj.weight'],t.n,t.gate) |
| matvec(w[p..'mlp.up_proj.weight'],t.n,t.up) |
| for i=0,FF-1 do local g=rd(t.gate,i*4);wr(t.gate,i*4,g/(1+math.exp(-g))*rd(t.up,i*4)) end |
| matvec(w[p..'mlp.down_proj.weight'],t.gate,t.o);add(x,t.o) |
| checkpoint(false) |
| end |
| session.position+=1 |
| if not project then return nil end |
| norm(x,w['model.norm.weight'],t.n);matvec(emb,t.n,t.logits) |
| return t.logits |
| end |
| function Engine.sample(logits,temperature,rng,previous) |
| local best={};local repetition={} |
| for _,id in previous or {} do repetition[id]=true end |
| for id=0,VOCAB-1 do |
| local score=rd(logits,id*4) |
| if repetition[id] then score=if score>0 then score/1.08 else score*1.08 end |
| if id==0 or id==1 or (id>=3 and id<=16) then score=-math.huge end |
| if #best<32 or score>best[#best][2] then |
| local k=#best+1 |
| while k>1 and best[k-1][2]<score do best[k]=best[k-1];k-=1 end |
| best[k]={id,score};if #best>32 then best[33]=nil end |
| end |
| end |
| if temperature==0 then return best[1][1] end |
| local sum=0 |
| for _,v in best do v[3]=math.exp((v[2]-best[1][2])/temperature);sum+=v[3] end |
| local z=rng:NextNumber()*sum |
| for _,v in best do z-=v[3];if z<=0 then return v[1] end end |
| return best[1][1] |
| end |
| return Engine |
|
|
