File size: 4,559 Bytes
520f98d 8020c75 520f98d 8020c75 520f98d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | <!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>
|