| <!doctype html> |
| <meta charset="utf-8" /> |
| <title>Baberu FP32 decoder WebGPU smoke test</title> |
| <style> |
| body { |
| background: #111; |
| color: #eee; |
| font: 14px/1.45 ui-monospace, monospace; |
| margin: 24px; |
| } |
| pre { white-space: pre-wrap; } |
| </style> |
| <h1>Baberu FP32 decoder WebGPU smoke test</h1> |
| <pre id="log">Starting…</pre> |
| <script type="module"> |
| import * as ort from "./.cache/ort/ort.webgpu.min.mjs"; |
| |
| const output = document.querySelector("#log"); |
| const lines = []; |
| const log = (message) => { |
| lines.push(`${new Date().toISOString()} ${message}`); |
| output.textContent = lines.join("\n"); |
| console.log(message); |
| }; |
| const timed = async (label, operation) => { |
| const started = performance.now(); |
| const result = await operation(); |
| log(`PASS ${label} (${(performance.now() - started).toFixed(1)} ms)`); |
| return result; |
| }; |
| window.addEventListener("error", (event) => log(`FAIL ${event.error?.stack ?? event.message}`)); |
| window.addEventListener("unhandledrejection", (event) => log(`FAIL ${event.reason?.stack ?? event.reason}`)); |
| |
| ort.env.wasm.numThreads = 1; |
| ort.env.wasm.wasmPaths = "/.cache/ort/"; |
| log(`navigator.gpu=${Boolean(navigator.gpu)} secure=${window.isSecureContext}`); |
| if (!navigator.gpu) throw new Error("WebGPU is unavailable"); |
| |
| const cacheNames = [ |
| ...Array.from({ length: 6 }, (_, index) => `present_k${index}`), |
| ...Array.from({ length: 6 }, (_, index) => `present_v${index}`), |
| ]; |
| const parameters = new URLSearchParams(location.search); |
| const gpuCache = parameters.get("cache") !== "cpu"; |
| const stepCount = Number(parameters.get("steps") ?? 5); |
| log(`gpuCache=${gpuCache}`); |
| log(`stepCount=${stepCount}`); |
| const preferredOutputLocation = Object.fromEntries([ |
| ["logits", "cpu"], |
| ...cacheNames.map((name) => [name, gpuCache ? "gpu-buffer" : "cpu"]), |
| ]); |
| const sessionOptions = { |
| executionProviders: ["webgpu"], |
| preferredOutputLocation, |
| logSeverityLevel: 2, |
| }; |
| const createSession = async (path) => { |
| const bytes = new Uint8Array(await (await fetch(path)).arrayBuffer()); |
| return ort.InferenceSession.create(bytes, sessionOptions); |
| }; |
| |
| const prefill = await timed("prefill session creation", () => |
| createSession("./output/decoder_prefill_fp32.onnx") |
| ); |
| const visionEmbeds = new ort.Tensor( |
| "float32", |
| new Float32Array(1 * 256 * 512), |
| [1, 256, 512] |
| ); |
| const argmax = (values) => { |
| let result = 0; |
| for (let index = 1; index < values.length; index += 1) { |
| if (values[index] > values[result]) result = index; |
| } |
| return result; |
| }; |
| const prefillResult = await timed("prefill execution", () => |
| prefill.run({ vision_embeds: visionEmbeds }) |
| ); |
| log(`prefill logits=${prefillResult.logits.location} ${JSON.stringify(prefillResult.logits.dims)}`); |
| log(`prefill cache=${prefillResult.present_k0.location} ${JSON.stringify(prefillResult.present_k0.dims)}`); |
| const logits = prefillResult.logits.data; |
| log(`prefill token=${argmax(logits)} min=${Math.min(...logits)} max=${Math.max(...logits)} mean=${logits.reduce((sum, value) => sum + value, 0) / logits.length} sample=${Array.from(logits.slice(0, 10)).join(",")}`); |
| |
| const step = await timed("step session creation", () => |
| createSession("./output/decoder_step_fp32.onnx") |
| ); |
| const oneHot = (token) => { |
| const values = new Float32Array(14630); |
| values[token] = 1; |
| return new ort.Tensor("float32", values, [1, 1, 14630]); |
| }; |
| let token = argmax(prefillResult.logits.data); |
| let cacheResult = prefillResult; |
| for (let iteration = 0; iteration < stepCount; iteration += 1) { |
| const stepFeeds = { |
| token_one_hot: oneHot(token), |
| position_ids: new ort.Tensor("int32", new Int32Array([257 + iteration]), [1, 1]), |
| }; |
| for (let index = 0; index < 6; index += 1) { |
| stepFeeds[`past_k${index}`] = cacheResult[`present_k${index}`]; |
| stepFeeds[`past_v${index}`] = cacheResult[`present_v${index}`]; |
| } |
| const stepResult = await timed( |
| `step ${iteration + 1} execution with ${gpuCache ? "GPU" : "CPU"} KV cache`, |
| () => step.run(stepFeeds) |
| ); |
| token = argmax(stepResult.logits.data); |
| log(`step ${iteration + 1} token=${token} cache=${stepResult.present_k0.location} ${JSON.stringify(stepResult.present_k0.dims)}`); |
| for (const name of cacheNames) cacheResult[name].dispose(); |
| cacheResult.logits.dispose(); |
| stepFeeds.token_one_hot.dispose(); |
| stepFeeds.position_ids.dispose(); |
| cacheResult = stepResult; |
| } |
| for (const name of cacheNames) cacheResult[name].dispose(); |
| cacheResult.logits.dispose(); |
| visionEmbeds.dispose(); |
| await prefill.release(); |
| await step.release(); |
| log("DONE"); |
| </script> |
|
|
