Spaces:
Running
Running
ternary GEMV kernel lab: V0/V1/V2 variants
Browse files- gemv-bench.html +183 -0
gemv-bench.html
ADDED
|
@@ -0,0 +1,183 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!doctype html><html><head><meta charset=utf8><meta name=viewport content="width=device-width,initial-scale=1">
|
| 2 |
+
<title>Ternary GEMV kernel lab — your GPU</title>
|
| 3 |
+
<style>
|
| 4 |
+
:root{--bg:#0a0d13;--panel:#111825;--ink:#e8ebf1;--dim:#8b95a7;--ac:#7c5cff;--ok:#48c26c;--no:#f0616d;--warn:#e0a94a;--line:#1e2836}
|
| 5 |
+
*{box-sizing:border-box}body{margin:0;background:var(--bg);color:var(--ink);font:15px/1.6 -apple-system,Segoe UI,Roboto,monospace;padding:24px;max-width:820px;margin:0 auto}
|
| 6 |
+
h1{font-size:20px;margin:0 0 4px}.sub{color:var(--dim);font-size:13px;margin:0 0 18px}
|
| 7 |
+
.card{font-family:ui-monospace,monospace;padding:14px 16px;border:1px solid var(--line);border-radius:10px;background:var(--panel);margin-bottom:12px;word-break:break-word}
|
| 8 |
+
.big{font-size:22px;font-weight:700}.k{color:var(--dim)}
|
| 9 |
+
.ok{color:var(--ok)}.no{color:var(--no)}.warn{color:var(--warn)}
|
| 10 |
+
table{width:100%;border-collapse:collapse;font-family:ui-monospace,monospace;font-size:13px}
|
| 11 |
+
td,th{padding:6px 8px;border-bottom:1px solid var(--line);text-align:left}th{color:var(--dim);font-weight:600}
|
| 12 |
+
td.n{text-align:right;font-variant-numeric:tabular-nums}
|
| 13 |
+
.verdict{font-size:16px;font-weight:600;margin-top:14px}
|
| 14 |
+
</style></head><body>
|
| 15 |
+
<h1>Ternary GEMV <span style="color:var(--ac)">kernel lab</span> — your GPU</h1>
|
| 16 |
+
<p class="sub">Runs the batch-1 BitNet decode kernel (weights read once/token) in several variants against a ~0.69 GB ternary matrix on your real GPU, and reports the decode rate each reaches vs the 152 GB/s roofline (220 tok/s).</p>
|
| 17 |
+
<div id="status" class="card">starting…</div>
|
| 18 |
+
<div id="out"></div>
|
| 19 |
+
|
| 20 |
+
<script type="module">
|
| 21 |
+
const $=s=>document.querySelector(s), st=$("#status"), out=$("#out");
|
| 22 |
+
const MODEL_GB = 0.69, ROOF_TOK = 220, CUR_TOK = 54;
|
| 23 |
+
const say=(t,cls)=>{ st.textContent=t; st.className="card"+(cls?" "+cls:""); };
|
| 24 |
+
window.addEventListener("unhandledrejection",e=>say("✗ unhandled: "+(e.reason&&(e.reason.message||e.reason)),"no"));
|
| 25 |
+
|
| 26 |
+
// ---- kernel body templates (KX=1 decode). Same reduction structure; vary the inner math. ----
|
| 27 |
+
// V0: scalar u32 load, per-code f32 convert + (-1) bias (the current engine kernel)
|
| 28 |
+
const DOT_V0 = `
|
| 29 |
+
fn dot16(word:u32, v:u32) -> f32 {
|
| 30 |
+
var s4=vec4<f32>(0.0);
|
| 31 |
+
s4=s4+x[v] *(vec4<f32>(f32(word&3u),f32((word>>2u)&3u),f32((word>>4u)&3u),f32((word>>6u)&3u))-vec4<f32>(1.0));
|
| 32 |
+
s4=s4+x[v+1u]*(vec4<f32>(f32((word>>8u)&3u),f32((word>>10u)&3u),f32((word>>12u)&3u),f32((word>>14u)&3u))-vec4<f32>(1.0));
|
| 33 |
+
s4=s4+x[v+2u]*(vec4<f32>(f32((word>>16u)&3u),f32((word>>18u)&3u),f32((word>>20u)&3u),f32((word>>22u)&3u))-vec4<f32>(1.0));
|
| 34 |
+
s4=s4+x[v+3u]*(vec4<f32>(f32((word>>24u)&3u),f32((word>>26u)&3u),f32((word>>28u)&3u),f32((word>>30u)&3u))-vec4<f32>(1.0));
|
| 35 |
+
return s4.x+s4.y+s4.z+s4.w;
|
| 36 |
+
}`;
|
| 37 |
+
const BODY_V0 = `let word=qw[rowW+w]; acc=acc+dot16(word, w<<2u);`;
|
| 38 |
+
const LOOP_V0 = { arrQw:"array<u32>", stride:"64u", nwExpr:"nw", asum:false };
|
| 39 |
+
|
| 40 |
+
// V1: wide vec4<u32> load (64 codes / 16 bytes per load), same math as V0
|
| 41 |
+
const BODY_V1 = `let qv=qw[rowW+w]; let v=(w<<2u)<<2u;
|
| 42 |
+
acc=acc+dot16(qv.x,v)+dot16(qv.y,v+4u)+dot16(qv.z,v+8u)+dot16(qv.w,v+12u);`;
|
| 43 |
+
const LOOP_V1 = { arrQw:"array<vec4<u32>>", stride:"64u", nwExpr:"(nw>>2u)", asum:false };
|
| 44 |
+
|
| 45 |
+
// V2: wide load + BIAS FOLD — drop the 16 per-code (-1); subtract sum(activations) once per row
|
| 46 |
+
const DOT_V2 = `
|
| 47 |
+
fn dot16(word:u32, v:u32) -> f32 {
|
| 48 |
+
var s4=vec4<f32>(0.0);
|
| 49 |
+
s4=s4+x[v] *vec4<f32>(f32(word&3u),f32((word>>2u)&3u),f32((word>>4u)&3u),f32((word>>6u)&3u));
|
| 50 |
+
s4=s4+x[v+1u]*vec4<f32>(f32((word>>8u)&3u),f32((word>>10u)&3u),f32((word>>12u)&3u),f32((word>>14u)&3u));
|
| 51 |
+
s4=s4+x[v+2u]*vec4<f32>(f32((word>>16u)&3u),f32((word>>18u)&3u),f32((word>>20u)&3u),f32((word>>22u)&3u));
|
| 52 |
+
s4=s4+x[v+3u]*vec4<f32>(f32((word>>24u)&3u),f32((word>>26u)&3u),f32((word>>28u)&3u),f32((word>>30u)&3u));
|
| 53 |
+
return s4.x+s4.y+s4.z+s4.w;
|
| 54 |
+
}`;
|
| 55 |
+
const BODY_V2 = `let qv=qw[rowW+w]; let v=(w<<2u)<<2u;
|
| 56 |
+
acc=acc+dot16(qv.x,v)+dot16(qv.y,v+4u)+dot16(qv.z,v+8u)+dot16(qv.w,v+12u);`;
|
| 57 |
+
const LOOP_V2 = { arrQw:"array<vec4<u32>>", stride:"64u", nwExpr:"(nw>>2u)", asum:true };
|
| 58 |
+
|
| 59 |
+
function kernelWGSL(dot, body, L){
|
| 60 |
+
return `
|
| 61 |
+
@group(0) @binding(0) var<storage,read> x: array<vec4<f32>>;
|
| 62 |
+
@group(0) @binding(1) var<storage,read> qw: ${L.arrQw};
|
| 63 |
+
@group(0) @binding(2) var<storage,read_write> o: array<f32>;
|
| 64 |
+
@group(0) @binding(3) var<uniform> P: vec4<u32>; // x=K, y=N, z=0, w=asumBits
|
| 65 |
+
var<workgroup> red: array<f32, 256>;
|
| 66 |
+
${dot}
|
| 67 |
+
@compute @workgroup_size(256)
|
| 68 |
+
fn main(@builtin(workgroup_id) wg:vec3<u32>, @builtin(local_invocation_id) lid:vec3<u32>){
|
| 69 |
+
let K=P.x; let nw=K>>4u; let rr=lid.x/64u; let t=lid.x%64u;
|
| 70 |
+
let n0=(wg.y*65535u+wg.x)*4u+rr; let n=min(n0, P.y-1u);
|
| 71 |
+
let rowW=n*${L.nwExpr}; var acc=0.0;
|
| 72 |
+
var w=t;
|
| 73 |
+
loop{ if(w>=${L.nwExpr}){break;}
|
| 74 |
+
${body}
|
| 75 |
+
w=w+${L.stride}; }
|
| 76 |
+
${L.asum? "acc = acc - bitcast<f32>(P.w);" : ""}
|
| 77 |
+
red[lid.x]=acc; workgroupBarrier();
|
| 78 |
+
var s=32u; loop{ if(s==0u){break;} if(t<s){ red[rr*64u+t]=red[rr*64u+t]+red[rr*64u+t+s]; } workgroupBarrier(); s=s/2u; }
|
| 79 |
+
if(t==0u && n0<P.y){ o[n0]=red[rr*64u]; }
|
| 80 |
+
}`;
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
// pseudo-random 2-bit codes fill (non-zero, non-compressible)
|
| 84 |
+
const FILL = `
|
| 85 |
+
@group(0) @binding(0) var<storage,read_write> d: array<u32>;
|
| 86 |
+
@group(0) @binding(1) var<uniform> P: vec4<u32>; // x=#u32, y=stride
|
| 87 |
+
@compute @workgroup_size(256)
|
| 88 |
+
fn main(@builtin(global_invocation_id) gid:vec3<u32>){
|
| 89 |
+
let n=P.x; var i=gid.x;
|
| 90 |
+
loop{ if(i>=n){break;} d[i]=(i*2654435761u+1u)^((i<<7u)*40503u); i=i+P.y; }
|
| 91 |
+
}`;
|
| 92 |
+
|
| 93 |
+
const VARIANTS=[
|
| 94 |
+
{id:"V0", name:"V0 baseline (scalar load, per-code −1)", dot:DOT_V0, body:BODY_V0, L:LOOP_V0},
|
| 95 |
+
{id:"V1", name:"V1 wide vec4 load", dot:DOT_V0, body:BODY_V1, L:LOOP_V1},
|
| 96 |
+
{id:"V2", name:"V2 wide load + bias-fold", dot:DOT_V2, body:BODY_V2, L:LOOP_V2},
|
| 97 |
+
];
|
| 98 |
+
|
| 99 |
+
(async()=>{
|
| 100 |
+
try{
|
| 101 |
+
if(!navigator.gpu){ say("✗ No WebGPU. Open in Chrome/Edge.","no"); return; }
|
| 102 |
+
say("requesting adapter…");
|
| 103 |
+
let ad=await navigator.gpu.requestAdapter({powerPreference:"high-performance"}); if(!ad) ad=await navigator.gpu.requestAdapter();
|
| 104 |
+
if(!ad){ say("✗ no GPU adapter.","no"); return; }
|
| 105 |
+
const L=ad.limits, info=ad.info||{};
|
| 106 |
+
const dev=await ad.requestDevice({requiredLimits:{
|
| 107 |
+
maxStorageBufferBindingSize:L.maxStorageBufferBindingSize, maxBufferSize:L.maxBufferSize,
|
| 108 |
+
maxComputeWorkgroupsPerDimension:L.maxComputeWorkgroupsPerDimension }});
|
| 109 |
+
dev.lost.then(i=>say("✗ device lost: "+(i&&i.message||i.reason||""),"no"));
|
| 110 |
+
|
| 111 |
+
const K=4096, nw=K/16; // 256 words/row
|
| 112 |
+
// weight bytes: as close to 0.69 GB as the device allows
|
| 113 |
+
const wantBytes=Math.min(L.maxStorageBufferBindingSize, L.maxBufferSize, Math.round(MODEL_GB*1e9));
|
| 114 |
+
const N=Math.floor((wantBytes/4)/nw); // whole rows
|
| 115 |
+
const qU32=N*nw, wBytes=qU32*4;
|
| 116 |
+
say("allocating "+(wBytes/1048576).toFixed(0)+" MB ternary weights ("+N.toLocaleString()+" rows × "+K+")…");
|
| 117 |
+
|
| 118 |
+
const qw=dev.createBuffer({size:wBytes, usage:GPUBufferUsage.STORAGE});
|
| 119 |
+
const xbuf=dev.createBuffer({size:K*4, usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_DST});
|
| 120 |
+
const obuf=dev.createBuffer({size:N*4, usage:GPUBufferUsage.STORAGE});
|
| 121 |
+
const P=dev.createBuffer({size:16, usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});
|
| 122 |
+
|
| 123 |
+
// activations: random f32, and their sum (for V2 bias fold)
|
| 124 |
+
const xa=new Float32Array(K); let asum=0; for(let i=0;i<K;i++){ xa[i]=(Math.sin(i*12.9898)*43758.5453)%1; asum+=xa[i]; }
|
| 125 |
+
dev.queue.writeBuffer(xbuf,0,xa);
|
| 126 |
+
const asumBits=new Uint32Array(new Float32Array([asum]).buffer)[0];
|
| 127 |
+
dev.queue.writeBuffer(P,0,new Uint32Array([K, N, 0, asumBits]));
|
| 128 |
+
|
| 129 |
+
// fill weights (non-zero) so reads hit real VRAM
|
| 130 |
+
say("filling weights…");
|
| 131 |
+
{ const fmod=dev.createShaderModule({code:FILL});
|
| 132 |
+
const fpipe=dev.createComputePipeline({layout:"auto",compute:{module:fmod,entryPoint:"main"}});
|
| 133 |
+
const fp=dev.createBuffer({size:16,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});
|
| 134 |
+
const wgF=Math.min(L.maxComputeWorkgroupsPerDimension,65535);
|
| 135 |
+
dev.queue.writeBuffer(fp,0,new Uint32Array([qU32, wgF*256, 0, 0]));
|
| 136 |
+
const fbg=dev.createBindGroup({layout:fpipe.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:qw}},{binding:1,resource:{buffer:fp}}]});
|
| 137 |
+
const e=dev.createCommandEncoder(); const p=e.beginComputePass(); p.setPipeline(fpipe); p.setBindGroup(0,fbg); p.dispatchWorkgroups(wgF); p.end();
|
| 138 |
+
dev.queue.submit([e.finish()]); await dev.queue.onSubmittedWorkDone(); }
|
| 139 |
+
|
| 140 |
+
// dispatch dims: workgroup = 4 rows
|
| 141 |
+
const groups=Math.ceil(N/4), wgx=Math.min(groups,65535), wgy=Math.ceil(groups/65535);
|
| 142 |
+
const PASSES=8, ITERS=5;
|
| 143 |
+
|
| 144 |
+
const rows=[];
|
| 145 |
+
for(const V of VARIANTS){
|
| 146 |
+
say("compiling "+V.id+"…");
|
| 147 |
+
dev.pushErrorScope("validation");
|
| 148 |
+
const mod=dev.createShaderModule({code:kernelWGSL(V.dot,V.body,V.L)});
|
| 149 |
+
const ci=await mod.getCompilationInfo(); const er=ci.messages.filter(m=>m.type==="error");
|
| 150 |
+
if(er.length){ await dev.popErrorScope(); rows.push({V, err:er[0].message}); continue; }
|
| 151 |
+
const pipe=dev.createComputePipeline({layout:"auto",compute:{module:mod,entryPoint:"main"}});
|
| 152 |
+
const bg=dev.createBindGroup({layout:pipe.getBindGroupLayout(0),entries:[
|
| 153 |
+
{binding:0,resource:{buffer:xbuf}},{binding:1,resource:{buffer:qw}},{binding:2,resource:{buffer:obuf}},{binding:3,resource:{buffer:P}}]});
|
| 154 |
+
const scopeErr=await dev.popErrorScope();
|
| 155 |
+
if(scopeErr){ rows.push({V, err:scopeErr.message}); continue; }
|
| 156 |
+
async function run(passes){
|
| 157 |
+
const e=dev.createCommandEncoder();
|
| 158 |
+
for(let k=0;k<passes;k++){ const p=e.beginComputePass(); p.setPipeline(pipe); p.setBindGroup(0,bg); p.dispatchWorkgroups(wgx,wgy); p.end(); }
|
| 159 |
+
const t0=performance.now(); dev.queue.submit([e.finish()]); await dev.queue.onSubmittedWorkDone(); return performance.now()-t0;
|
| 160 |
+
}
|
| 161 |
+
say("timing "+V.id+"…");
|
| 162 |
+
await run(2); // warm
|
| 163 |
+
let best=1e9; for(let k=0;k<ITERS;k++){ best=Math.min(best, await run(PASSES)); }
|
| 164 |
+
const msPer=best/PASSES, gbps=(wBytes/1073741824)/(msPer/1000), tok=gbps/MODEL_GB, pct=100*tok/ROOF_TOK;
|
| 165 |
+
rows.push({V, msPer, gbps, tok, pct});
|
| 166 |
+
}
|
| 167 |
+
|
| 168 |
+
say("done · adapter: "+((info.vendor||"?")+" "+(info.architecture||"")+" "+(info.device||"")).trim(),"ok");
|
| 169 |
+
const best=rows.filter(r=>r.tok).sort((a,b)=>b.tok-a.tok)[0];
|
| 170 |
+
out.innerHTML=`
|
| 171 |
+
<div class="card"><span class="k">Test matrix</span><br><span class="big">${N.toLocaleString()} × ${K}</span> · ${(wBytes/1073741824).toFixed(2)} GB ternary weights · roofline ${ROOF_TOK} tok/s</div>
|
| 172 |
+
<div class="card"><table>
|
| 173 |
+
<tr><th>variant</th><th class="n">ms/token</th><th class="n">GB/s</th><th class="n">tok/s</th><th class="n">% roofline</th></tr>
|
| 174 |
+
${rows.map(r=>r.err
|
| 175 |
+
? `<tr><td>${r.V.name}</td><td colspan=4 class="no">compile/validate failed: ${r.err.slice(0,60)}</td></tr>`
|
| 176 |
+
: `<tr><td>${r.V.name}</td><td class="n">${r.msPer.toFixed(2)}</td><td class="n">${r.gbps.toFixed(0)}</td><td class="n ${r===best?'ok':''}"><b>${r.tok.toFixed(0)}</b></td><td class="n ${r.pct<30?'no':r.pct<70?'warn':'ok'}">${r.pct.toFixed(0)}%</td></tr>`).join("")}
|
| 177 |
+
<tr><td class="k">current shipped kernel (measured)</td><td class="n k">18.4</td><td class="n k">38</td><td class="n k">${CUR_TOK}</td><td class="n k">${(100*CUR_TOK/ROOF_TOK).toFixed(0)}%</td></tr>
|
| 178 |
+
</table></div>
|
| 179 |
+
${best? `<div class="verdict ok">✓ Best: <b>${best.V.name}</b> → ${best.tok.toFixed(0)} tok/s (${best.pct.toFixed(0)}% of roofline, ${(best.tok/CUR_TOK).toFixed(1)}× the shipped kernel). ${best.pct>=70? "Bandwidth-bound — ship it." : "Still ALU-bound; next lever = integer dot / fewer converts."}</div>`
|
| 180 |
+
: `<div class="verdict no">All variants failed to run.</div>`}`;
|
| 181 |
+
}catch(e){ say("✗ "+(e&&(e.message||e)), "no"); }
|
| 182 |
+
})();
|
| 183 |
+
</script></body></html>
|