| <!doctype html> |
| <meta charset="utf-8" /> |
| <title>Baberu capability-preserving WebGPU memory experiment</title> |
| <style> |
| body { background: #111; color: #eee; font: 14px/1.45 ui-monospace, monospace; margin: 24px; } |
| pre { white-space: pre-wrap; } |
| </style> |
| <h1>Baberu capability-preserving WebGPU memory experiment</h1> |
| <pre id="log">Starting…</pre> |
| <canvas id="canvas" width="224" height="224" hidden></canvas> |
| <script type="module"> |
| const parameters = new URLSearchParams(location.search); |
| const tier = parameters.get("tier") ?? "121"; |
| const memoryMode = parameters.get("memory") ?? "lean"; |
| const loaderMode = parameters.get("loader") ?? "url"; |
| const scheduleMode = parameters.get("schedule") ?? "staged"; |
| const decoderMode = parameters.get("decoder") ?? "baseline"; |
| const benchmarkScope = parameters.get("scope") ?? "full"; |
| const decoderOnly = benchmarkScope === "decoder"; |
| const maxTokens = 128; |
| if (!new Set(["121", "242"]).has(tier)) throw new Error(`Unknown tier ${tier}`); |
| if (!new Set(["default", "lean"]).has(memoryMode)) throw new Error(`Unknown memory mode ${memoryMode}`); |
| if (!new Set(["url", "bytes"]).has(loaderMode)) throw new Error(`Unknown loader ${loaderMode}`); |
| if (!new Set(["resident", "staged"]).has(scheduleMode)) throw new Error(`Unknown schedule ${scheduleMode}`); |
| if (!new Set(["baseline", "gather-opt", "fp16-matmul", "fixed-kv"]).has(decoderMode)) throw new Error(`Unknown decoder ${decoderMode}`); |
| if (!new Set(["full", "decoder"]).has(benchmarkScope)) throw new Error(`Unknown scope ${benchmarkScope}`); |
| const ort = await import("./.work/ort/ort-webgpu.mjs"); |
| const output = document.querySelector("#log"); |
| const canvas = document.querySelector("#canvas"); |
| const context = canvas.getContext("2d", { willReadFrequently: true }); |
| const logLines = []; |
| const memorySamples = []; |
| const log = (message) => { |
| logLines.push(`${new Date().toISOString()} ${message}`); |
| output.textContent = logLines.join("\n"); |
| console.log(message); |
| }; |
| window.addEventListener("error", (event) => log(`FAIL ${event.error?.stack ?? event.message}`)); |
| window.addEventListener("unhandledrejection", (event) => log(`FAIL ${event.reason?.stack ?? event.reason}`)); |
| const sampleMemory = async (label) => { |
| const sample = { label, atMs: performance.now() }; |
| if (performance.memory) { |
| sample.jsHeapSizeLimit = performance.memory.jsHeapSizeLimit; |
| sample.totalJSHeapSize = performance.memory.totalJSHeapSize; |
| sample.usedJSHeapSize = performance.memory.usedJSHeapSize; |
| } |
| if (typeof performance.measureUserAgentSpecificMemory === "function") { |
| try { |
| sample.userAgentBytes = (await performance.measureUserAgentSpecificMemory()).bytes; |
| } catch (error) { |
| sample.userAgentMemoryError = String(error); |
| } |
| } |
| memorySamples.push(sample); |
| log(`MEMORY ${JSON.stringify(sample)}`); |
| return sample; |
| }; |
| const runId = `tier-${tier}-${memoryMode}-${loaderMode}-${scheduleMode}-${decoderMode}-${benchmarkScope}-${Date.now()}`; |
| await fetch("/benchmark-start", { |
| method: "POST", |
| headers: { "Content-Type": "application/json" }, |
| body: JSON.stringify({ runId }), |
| }); |
| const markPhase = async (phase) => { |
| const response = await fetch("/benchmark-mark", { |
| method: "POST", |
| headers: { "Content-Type": "application/json" }, |
| body: JSON.stringify({ runId, phase }), |
| }); |
| if (!response.ok) throw new Error(`Failed marking ${phase}: HTTP ${response.status}`); |
| }; |
| await markPhase("baseline"); |
| await sampleMemory("baseline"); |
| |
| ort.env.wasm.numThreads = 1; |
| ort.env.wasm.wasmPaths = "/.work/ort/"; |
| if (!navigator.gpu) throw new Error("WebGPU is unavailable"); |
| let gpuDevice = null; |
| const ensureMainDevice = async () => { |
| if (gpuDevice) return; |
| const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" }); |
| if (!adapter) throw new Error("No WebGPU adapter is available"); |
| gpuDevice = await adapter.requestDevice(); |
| ort.env.webgpu.device = gpuDevice; |
| }; |
| const cacheNames = [ |
| ...Array.from({ length: 6 }, (_, index) => `present_k${index}`), |
| ...Array.from({ length: 6 }, (_, index) => `present_v${index}`), |
| ]; |
| const createSession = async (path, preferredOutputLocation) => { |
| await ensureMainDevice(); |
| const options = { |
| executionProviders: ["webgpu"], |
| preferredOutputLocation, |
| logSeverityLevel: 3, |
| ...(memoryMode === "lean" ? { enableCpuMemArena: false, enableMemPattern: false } : {}), |
| }; |
| const started = performance.now(); |
| const session = loaderMode === "url" |
| ? await ort.InferenceSession.create(path, options) |
| : await fetch(path).then(async (response) => { |
| if (!response.ok) throw new Error(`${path}: HTTP ${response.status}`); |
| return ort.InferenceSession.create(new Uint8Array(await response.arrayBuffer()), options); |
| }); |
| log(`SESSION ${path} ${(performance.now() - started).toFixed(1)} ms`); |
| return session; |
| }; |
| const visionPath = tier === "121" |
| ? "./.work/models/webgpu-121/vision_int4.onnx" |
| : "./.work/models/webgpu-242/vision_fp16.onnx"; |
| const fixedKv = decoderMode === "fixed-kv"; |
| const decoderPath = ({ |
| baseline: "./.work/models/shared/decoder_unified_gather_qdq_int8.onnx", |
| "gather-opt": "./.work/models/model-opt/decoder_gather_before_dq_int8.onnx", |
| "fp16-matmul": "./.work/models/model-opt/decoder_static_fp16_matmul.onnx", |
| "fixed-kv": "./.work/models/model-opt/decoder_gather_dq_fixed_kv_int8.onnx", |
| })[decoderMode]; |
| let vision = null; |
| let decoder = null; |
| const createVision = async () => { |
| vision = await createSession(visionPath, scheduleMode === "staged" ? { vision_embeds: "cpu" } : { vision_embeds: "gpu-buffer" }); |
| if (JSON.stringify(vision.inputNames) !== JSON.stringify(["pixel_values"])) { |
| throw new Error(`Unexpected vision inputs ${vision.inputNames}`); |
| } |
| await sampleMemory("vision-session-ready"); |
| }; |
| const createDecoder = async () => { |
| decoder = await createSession(decoderPath, Object.fromEntries([ |
| ["logits", "cpu"], |
| ...cacheNames.map((name) => [name, "gpu-buffer"]), |
| ])); |
| for (const required of ["vision_embeds", "token_ids", "position_ids", "past_k0", "past_v5"]) { |
| if (!decoder.inputNames.includes(required)) throw new Error(`Decoder input missing ${required}`); |
| } |
| if (fixedKv !== decoder.inputNames.includes("past_length")) { |
| throw new Error(`Decoder past_length contract mismatch for ${decoderMode}`); |
| } |
| if (!decoder.outputNames.includes("logits") || !decoder.outputNames.includes("present_v5")) { |
| throw new Error("Decoder capability outputs are incomplete"); |
| } |
| await sampleMemory("decoder-session-ready"); |
| }; |
| if (scheduleMode === "resident" && !decoderOnly) { |
| await markPhase("vision-load"); |
| await createVision(); |
| await markPhase("decoder-load"); |
| await createDecoder(); |
| } |
| if (decoderOnly) { |
| await markPhase("decoder-load"); |
| await createDecoder(); |
| } |
| |
| const charset = await (await fetch("./tokenizer/vocab.json")).json(); |
| const idToCharacter = ["", "", "", "", ...charset]; |
| if (idToCharacter.length !== 14630) throw new Error(`Unexpected vocabulary ${idToCharacter.length}`); |
| const contentIds = new Set(); |
| for (let index = 0; index < charset.length; index += 1) { |
| const character = charset[index]; |
| if (character.length === 1 && !"ーｰ〜~".includes(character) && /[\p{Letter}\p{Number}]/u.test(character)) { |
| contentIds.add(index + 4); |
| } |
| } |
| const loadImage = async (path) => createImageBitmap(await (await fetch(path)).blob()); |
| const preprocess = (image) => { |
| context.imageSmoothingEnabled = true; |
| context.imageSmoothingQuality = "high"; |
| context.clearRect(0, 0, 224, 224); |
| context.drawImage(image, 0, 0, image.width, image.height, 0, 0, 224, 224); |
| const rgba = context.getImageData(0, 0, 224, 224).data; |
| const values = new Float32Array(3 * 224 * 224); |
| const mean = [0.485, 0.456, 0.406]; |
| const std = [0.229, 0.224, 0.225]; |
| for (let pixel = 0; pixel < 224 * 224; pixel += 1) { |
| for (let channel = 0; channel < 3; channel += 1) { |
| values[channel * 224 * 224 + pixel] = (rgba[pixel * 4 + channel] / 255 - mean[channel]) / std[channel]; |
| } |
| } |
| return values; |
| }; |
| const tokenId = (token) => new ort.Tensor("int32", Int32Array.of(token), [1, 1]); |
| const emptyVision = () => new ort.Tensor("float32", new Float32Array(0), [1, 0, 512]); |
| const emptyCache = () => fixedKv |
| ? new ort.Tensor("float32", new Float32Array(1 * 2 * 384 * 64), [1, 2, 384, 64]) |
| : new ort.Tensor("float32", new Float32Array(0), [1, 2, 0, 64]); |
| const pastLength = (length) => new ort.Tensor("int64", BigInt64Array.of(BigInt(length)), [1]); |
| const positions = () => new ort.Tensor("int32", Int32Array.from({ length: 257 }, (_, index) => index), [1, 257]); |
| const disposeResult = (result) => { |
| result.logits?.dispose(); |
| for (const name of cacheNames) result[name]?.dispose(); |
| }; |
| const chooseToken = (source, sequence, tokens) => { |
| const seen = new Set(sequence); |
| let blocked = -1; |
| const last = tokens.at(-1); |
| if (last !== undefined && contentIds.has(last)) { |
| let run = 0; |
| for (let index = tokens.length - 1; index >= 0 && tokens[index] === last; index -= 1) run += 1; |
| if (run >= 12) blocked = last; |
| } |
| const adjusted = (index) => { |
| if (index === blocked) return Number.NEGATIVE_INFINITY; |
| const value = source[index]; |
| if (!seen.has(index)) return value; |
| return value < 0 ? value * 1.2 : value / 1.2; |
| }; |
| let bestToken = 0; |
| let bestValue = adjusted(0); |
| for (let token = 1; token < 14630; token += 1) { |
| const value = adjusted(token); |
| if (value > bestValue) { |
| bestToken = token; |
| bestValue = value; |
| } |
| } |
| return bestToken; |
| }; |
| const encode = async (image) => { |
| const pixels = new ort.Tensor("float32", preprocess(image), [1, 3, 224, 224]); |
| const visionStarted = performance.now(); |
| const visionResult = await vision.run({ pixel_values: pixels }); |
| const visionMs = performance.now() - visionStarted; |
| pixels.dispose(); |
| return { visionEmbeds: visionResult.vision_embeds, visionMs }; |
| }; |
| const decode = async (visionEmbeds, visionMs) => { |
| const started = performance.now(); |
| const prefillFeeds = { |
| vision_embeds: visionEmbeds, |
| token_ids: tokenId(1), |
| position_ids: positions(), |
| ...Object.fromEntries(cacheNames.map((name) => [name.replace("present_", "past_"), emptyCache()])), |
| }; |
| if (fixedKv) prefillFeeds.past_length = pastLength(0); |
| const prefillStarted = performance.now(); |
| let cacheResult = await decoder.run(prefillFeeds); |
| const prefillMs = performance.now() - prefillStarted; |
| visionEmbeds.dispose(); |
| for (const [name, tensor] of Object.entries(prefillFeeds)) if (name !== "vision_embeds") tensor.dispose(); |
| const sequence = [1]; |
| const tokens = []; |
| const decodeStarted = performance.now(); |
| for (let iteration = 0; iteration < maxTokens; iteration += 1) { |
| |
| |
| const next = decoderOnly ? 4 : chooseToken(cacheResult.logits.data, sequence, tokens); |
| if (next === 2) break; |
| tokens.push(next); |
| sequence.push(next); |
| if (tokens.length >= maxTokens) break; |
| const feeds = { |
| vision_embeds: emptyVision(), |
| token_ids: tokenId(next), |
| position_ids: new ort.Tensor("int32", Int32Array.of(257 + iteration), [1, 1]), |
| }; |
| if (fixedKv) feeds.past_length = pastLength(257 + iteration); |
| for (let layer = 0; layer < 6; layer += 1) { |
| feeds[`past_k${layer}`] = cacheResult[`present_k${layer}`]; |
| feeds[`past_v${layer}`] = cacheResult[`present_v${layer}`]; |
| } |
| const nextResult = await decoder.run(feeds); |
| disposeResult(cacheResult); |
| feeds.vision_embeds.dispose(); |
| feeds.token_ids.dispose(); |
| feeds.position_ids.dispose(); |
| feeds.past_length?.dispose(); |
| cacheResult = nextResult; |
| } |
| const decodeMs = performance.now() - decodeStarted; |
| disposeResult(cacheResult); |
| return { |
| text: tokens.map((token) => idToCharacter[token] ?? "").join(""), |
| tokens: tokens.length, |
| visionMs, |
| prefillMs, |
| decodeMs, |
| totalMs: visionMs + performance.now() - started, |
| }; |
| }; |
| const recognize = async (image) => { |
| const encoded = await encode(image); |
| return decode(encoded.visionEmbeds, encoded.visionMs); |
| }; |
| const expected = [ |
| ["01", "知らない世界で見つけたイメージを"], |
| ["02", "カナデトモスソラ（Ｋａｎａｄｅｔｏｍｏｓｕｓｏｒａ）"], |
| ["03", "建設会社社員行方"], |
| ["04", "だとしてもこのレベルがウロつくなんて．．．おそらく２級の呪い"], |
| ["05", "パチパチパチパチ"], |
| ["06", "バビュン"], |
| ["07", "僕の過去とか未来とか"], |
| ["08", "くらべられっ子"], |
| ["09", "そうだクラス分けがあるんだった！！"], |
| ["10", "脇役よ、主役を超えよ！"], |
| ["11", "Ｅｈ～Ｉｄｏｎ＇ｔｒｅａｌｌｙｗａｎｔｔｏ～"], |
| ["12", "「Ｓｏｒｒｙｆｏｒｔｈｅｗａｉｔ～！Ｄｉｄｙｏｕｗａｉｔｌｏｎｇ？」"], |
| ["13", "ＹａｍａｔｅＡｒｅａＮｅｗｒｅｓｉｄｅｎｔｉａｌｄｉｓｔｒｉｃｔｆｏｒｆｏｒｅｉｇｎｅｒｓ"], |
| ]; |
| const normalize = (text) => text.normalize("NFKC").replace(/\s+/g, ""); |
| const editDistance = (left, right) => { |
| let previous = Array.from({ length: right.length + 1 }, (_, index) => index); |
| for (let leftIndex = 1; leftIndex <= left.length; leftIndex += 1) { |
| const current = [leftIndex]; |
| for (let rightIndex = 1; rightIndex <= right.length; rightIndex += 1) { |
| current.push(Math.min( |
| current[rightIndex - 1] + 1, |
| previous[rightIndex] + 1, |
| previous[rightIndex - 1] + (left[leftIndex - 1] === right[rightIndex - 1] ? 0 : 1), |
| )); |
| } |
| previous = current; |
| } |
| return previous.at(-1); |
| }; |
| const stagedInputs = []; |
| if (scheduleMode === "staged" && !decoderOnly) { |
| await markPhase("vision-load"); |
| const visionWorker = new Worker("./vision-worker.mjs", { type: "module" }); |
| let workerRequestId = 0; |
| const workerCalls = new Map(); |
| visionWorker.onmessage = ({ data }) => { |
| const call = workerCalls.get(data.requestId); |
| if (!call) return; |
| workerCalls.delete(data.requestId); |
| if (data.error) call.reject(new Error(data.error)); |
| else call.resolve(data.payload); |
| }; |
| visionWorker.onerror = (event) => { |
| for (const call of workerCalls.values()) call.reject(event.error ?? new Error(event.message)); |
| workerCalls.clear(); |
| }; |
| const callWorker = (type, payload, transfer = []) => new Promise((resolve, reject) => { |
| const requestId = ++workerRequestId; |
| workerCalls.set(requestId, { resolve, reject }); |
| visionWorker.postMessage({ requestId, type, payload }, transfer); |
| }); |
| await callWorker("init", { |
| modelPath: new URL(visionPath, location.href).href, |
| wasmPaths: new URL("./.work/ort/", location.href).href, |
| lean: memoryMode === "lean", |
| }); |
| await sampleMemory("vision-worker-session-ready"); |
| await markPhase("vision-encode"); |
| for (const [id, expectedText] of expected) { |
| log(`ENCODE ${id}/13`); |
| const image = await loadImage(`./.work/hayai13/${id}.png`); |
| const pixels = preprocess(image); |
| image.close(); |
| const encoded = await callWorker("encode", { pixels }, [pixels.buffer]); |
| stagedInputs.push({ id, expectedText, ...encoded }); |
| } |
| await sampleMemory("all-vision-embeddings-ready"); |
| await markPhase("vision-release"); |
| await callWorker("close", {}); |
| visionWorker.terminate(); |
| await new Promise((resolve) => setTimeout(resolve, 100)); |
| await sampleMemory("vision-worker-terminated"); |
| await markPhase("decoder-load"); |
| await createDecoder(); |
| } |
| const benchmarkCases = decoderOnly ? [["synthetic", ""]] : expected; |
| const results = []; |
| await markPhase("decoder-run"); |
| for (let caseIndex = 0; caseIndex < benchmarkCases.length; caseIndex += 1) { |
| const [id, expectedText] = benchmarkCases[caseIndex]; |
| log(`CASE ${id}/${benchmarkCases.length}`); |
| let result; |
| if (decoderOnly) { |
| const visionEmbeds = new ort.Tensor("float32", new Float32Array(256 * 512), [1, 256, 512]); |
| result = await decode(visionEmbeds, 0); |
| } else if (scheduleMode === "staged") { |
| const staged = stagedInputs[caseIndex]; |
| const visionEmbeds = new ort.Tensor("float32", staged.values, staged.dims); |
| result = await decode(visionEmbeds, staged.visionMs); |
| } else { |
| const image = await loadImage(`./.work/hayai13/${id}.png`); |
| result = await recognize(image); |
| image.close(); |
| } |
| const normalizedExpected = normalize(expectedText); |
| const normalizedActual = normalize(result.text); |
| results.push({ |
| id, |
| expected: expectedText, |
| actual: result.text, |
| distance: editDistance(normalizedExpected, normalizedActual), |
| characters: normalizedExpected.length, |
| exact: normalizedExpected === normalizedActual, |
| ...result, |
| }); |
| await sampleMemory(`after-case-${id}`); |
| } |
| const distance = results.reduce((sum, result) => sum + result.distance, 0); |
| const characters = results.reduce((sum, result) => sum + result.characters, 0); |
| const sortedTimes = results.map((result) => result.totalMs).sort((left, right) => left - right); |
| await markPhase("release"); |
| if (vision) await vision.release(); |
| if (decoder) await decoder.release(); |
| if (gpuDevice) gpuDevice.destroy(); |
| await new Promise((resolve) => setTimeout(resolve, 250)); |
| await sampleMemory("released"); |
| const summary = { |
| runId, |
| tier, |
| memoryMode, |
| loaderMode, |
| scheduleMode, |
| decoderMode, |
| benchmarkScope, |
| capability: { |
| layers: 6, |
| hiddenSize: 512, |
| intermediateSize: 1536, |
| attentionHeads: 8, |
| kvHeads: 2, |
| vocabulary: 14630, |
| maxTokens, |
| vision: tier === "121" ? "upstream INT4" : "upstream FP16", |
| decoder: ({ |
| baseline: "complete INT8-QDQ unified Gather", |
| "gather-opt": "complete INT8-QDQ Gather-before-DQ", |
| "fp16-matmul": "layers 1-3 static FP16 MatMul with Gather-before-DQ embedding", |
| "fixed-kv": "complete INT8-QDQ Gather-before-DQ with fixed 384-slot KV I/O", |
| })[decoderMode], |
| }, |
| cases: results.length, |
| nCER: characters ? distance / characters : null, |
| distance, |
| characters, |
| exact: results.filter((result) => result.exact).length, |
| medianMs: sortedTimes[Math.floor(sortedTimes.length / 2)], |
| peakCaseMs: sortedTimes.at(-1), |
| memorySamples, |
| results, |
| }; |
| const response = await fetch("/benchmark-result", { |
| method: "POST", |
| headers: { "Content-Type": "application/json" }, |
| body: JSON.stringify(summary), |
| }); |
| if (!response.ok) throw new Error(`Failed persisting result: HTTP ${response.status}`); |
| summary.processMemory = await response.json(); |
| window.__benchmarkResult = summary; |
| window.__benchmarkDone = true; |
| log(`SUITE_RESULT ${JSON.stringify(summary)}`); |
| </script> |
|
|
