webgpu-memory-opt.html

<!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) {
			// Decoder-only memory runs force a stable non-EOS token so every graph
			// executes the complete 128-token cache path.
			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>