Optimize Baberu WebGPU decoder memory execution
Browse filesPublish the exact Gather-before-DequantizeLinear decoder for both WebGPU tiers, plus reproducible worker-memory and rejected-candidate reports. Model architecture and generated outputs remain unchanged.
- MEMORY_RESULTS.md +28 -0
- MODEL_EXECUTION_RESULTS.md +55 -0
- README.md +12 -3
- README_MEMORY_OPT.md +82 -0
- optimize_decoder_execution.py +345 -0
- optimize_decoder_fp16_matmul.py +242 -0
- reports/model-execution-report.json +37 -0
- requirements-model-opt.txt +3 -0
- serve_memory_experiment.py +188 -0
- summarize_memory_results.py +76 -0
- variants/webgpu-121/decoder_unified_gather_qdq_int8.onnx +2 -2
- variants/webgpu-242/decoder_unified_gather_qdq_int8.onnx +2 -2
- vision-worker.mjs +57 -0
- webgpu-memory-opt.html +446 -0
MEMORY_RESULTS.md
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Measured memory result
|
| 2 |
+
|
| 3 |
+
Measured on 2026-07-17 in the T3 Chromium WebGPU renderer on Apple unified
|
| 4 |
+
memory. The control keeps both sessions resident and loads explicit model
|
| 5 |
+
bytes. The optimized mode uses direct URL loading, lean CPU allocation options,
|
| 6 |
+
and a dedicated vision worker that is terminated before decoder startup.
|
| 7 |
+
|
| 8 |
+
| Tier | Runs (control / optimized) | Renderer peak delta | JS heap peak | OCR parity |
|
| 9 |
+
| --- | ---: | ---: | ---: | --- |
|
| 10 |
+
| WebGPU-121 | 2 / 3 | 722.8 -> 710.1 MiB (-1.8%) | 87.0 -> 44.6 MiB (-48.7%) | exact, all 13 crops |
|
| 11 |
+
| WebGPU-242 | 3 / 3 | 937.5 -> 796.6 MiB (-15.0%) | 202.1 -> 44.6 MiB (-77.9%) | exact, all 13 crops |
|
| 12 |
+
|
| 13 |
+
The 242 tier receives a material renderer-RSS reduction because terminating the
|
| 14 |
+
vision worker returns its large FP16 model runtime before the decoder starts.
|
| 15 |
+
For the 121 tier, renderer RSS improves only slightly: its smaller INT4 vision
|
| 16 |
+
session is not the dominant renderer allocation. This means there is no honest
|
| 17 |
+
runtime-only change here that dramatically reduces the 121 tier without
|
| 18 |
+
changing its model representation.
|
| 19 |
+
|
| 20 |
+
`performance.memory` measures the page's JavaScript heap, while renderer RSS
|
| 21 |
+
also includes WebAssembly, browser internals, and other renderer allocations.
|
| 22 |
+
The GPU process is shared by all T3 tabs, so its sampled delta is intentionally
|
| 23 |
+
excluded from the comparison. Page staging also requires retaining all crop
|
| 24 |
+
embeddings (about 0.5 MiB each) until decoding begins and adds a worker/session
|
| 25 |
+
transition; it is best suited to known page-level crop batches.
|
| 26 |
+
|
| 27 |
+
Run `python3 summarize_memory_results.py` after new benchmark runs. The summary
|
| 28 |
+
fails if any optimized OCR output differs from its current tier control.
|
MODEL_EXECUTION_RESULTS.md
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Model-execution optimization result
|
| 2 |
+
|
| 3 |
+
Measured on 2026-07-17 with ONNX Runtime WebGPU 1.27.0 in the T3 Chromium
|
| 4 |
+
renderer on Apple unified memory. Every candidate retains all six decoder
|
| 5 |
+
layers, hidden size 512, two KV heads, vocabulary 14,630, and the 128-token
|
| 6 |
+
generation limit.
|
| 7 |
+
|
| 8 |
+
## Accepted exact rewrite
|
| 9 |
+
|
| 10 |
+
`decoder=gather-opt` changes only the token embedding order:
|
| 11 |
+
|
| 12 |
+
```text
|
| 13 |
+
baseline: INT8 table -> DequantizeLinear [14630,512] -> Gather one row
|
| 14 |
+
optimized: INT8 table -> Gather one row -> DequantizeLinear [1,1,512]
|
| 15 |
+
```
|
| 16 |
+
|
| 17 |
+
This removes a 29,962,240-byte (28.6 MiB) FP32 intermediate. ONNX Runtime CPU
|
| 18 |
+
comparison is bit-exact for prefill, cache length 257, and maximum cache length
|
| 19 |
+
383. Three full WebGPU runs per tier produced exactly the same 13 OCR strings
|
| 20 |
+
as baseline.
|
| 21 |
+
|
| 22 |
+
Renderer RSS does not show a repeatable whole-run reduction because WebGPU
|
| 23 |
+
buffers, the vision high-water mark, and Chromium allocator retention dominate
|
| 24 |
+
the process measurement:
|
| 25 |
+
|
| 26 |
+
| Full staged suite | Baseline median peak | Gather median peak | Change |
|
| 27 |
+
| --- | ---: | ---: | ---: |
|
| 28 |
+
| WebGPU-121, 3 + 3 runs | 836.0 MiB | 839.2 MiB | +0.4% |
|
| 29 |
+
| WebGPU-242, 3 + 3 runs | 1,160.2 MiB | 1,165.4 MiB | +0.4% |
|
| 30 |
+
|
| 31 |
+
The rewrite is kept because it is mathematically exact and removes unnecessary
|
| 32 |
+
work, but it must not be advertised as a dramatic measured RAM reduction.
|
| 33 |
+
|
| 34 |
+
## Rejected candidates
|
| 35 |
+
|
| 36 |
+
- Fixed-capacity KV I/O is CPU bit-exact and WebGPU output-exact, but standard
|
| 37 |
+
ONNX Runtime WebGPU does not alias past/present buffers. Creating fixed
|
| 38 |
+
outputs with Slice/Pad adds work and did not produce a stable peak reduction.
|
| 39 |
+
- Static FP16 MatMul for layers 1–3 preserves the current CPU top-token checks
|
| 40 |
+
and both 13-image WebGPU outputs. Its isolated 128-token median peak changed
|
| 41 |
+
from 727.1 MiB to 755.5 MiB (+3.9%), although median execution improved from
|
| 42 |
+
964.4 ms to 886.3 ms. It is therefore rejected for the memory objective.
|
| 43 |
+
- Converting all 42 internal MatMuls, or all 43 including the language head,
|
| 44 |
+
changed case 09 from `そうだクラスわけがあるんだった!!` to
|
| 45 |
+
`そうだクラスケはあるんだった!!`. Those variants fail the explicit
|
| 46 |
+
no-capability-regression gate.
|
| 47 |
+
|
| 48 |
+
## Remaining dramatic path
|
| 49 |
+
|
| 50 |
+
The decoder has 43 INT8-QDQ MatMul weights. Avoiding their runtime FP32
|
| 51 |
+
dequantized buffers without FP16 rounding requires a fused per-channel INT8
|
| 52 |
+
MatMul WebGPU kernel (or equivalent runtime operator). That is a custom runtime
|
| 53 |
+
implementation, not a safe standard-ONNX graph rewrite. The current experiment
|
| 54 |
+
does not claim that kernel exists.
|
| 55 |
+
|
README.md
CHANGED
|
@@ -30,8 +30,9 @@ server is required after the browser downloads the files.
|
|
| 30 |
## Latest recommended models
|
| 31 |
|
| 32 |
The latest decoder uses one complete six-layer QDQ graph for both prefill and
|
| 33 |
-
cached token steps. It accepts an INT32 token ID
|
| 34 |
-
|
|
|
|
| 35 |
|
| 36 |
| Folder | Vision encoder | Latest decoder | Latest ONNX download |
|
| 37 |
| --- | --- | --- | ---: |
|
|
@@ -61,7 +62,8 @@ selection policy. A GPU ArgMax graph and fused GroupQueryAttention graph were
|
|
| 61 |
tested locally but removed because both were slower on the tested runtime.
|
| 62 |
|
| 63 |
The QDQ decoder stores INT8 weights and executes `DequantizeLinear` followed by
|
| 64 |
-
FP32 `MatMul` on WebGPU.
|
|
|
|
| 65 |
`MatMulInteger`, which would leave the intended browser WebGPU path.
|
| 66 |
|
| 67 |
## 13 difficult showcase-image check
|
|
@@ -82,6 +84,11 @@ previous unified one-hot graph. The 242 Gather graph also retained the earlier
|
|
| 82 |
242 WebGPU result of 11.76% nCER and 5/13 exact. Latency depends heavily on GPU,
|
| 83 |
shader cache, browser, and system load; these numbers are not portable.
|
| 84 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 85 |
In one same-process 121 comparison, the optimized graph sampled an 833 MB
|
| 86 |
shared GPU-process peak and 577 MB renderer peak, versus 866 MB and 613 MB for
|
| 87 |
the unified one-hot baseline. Those values must not be added as model-owned
|
|
@@ -135,6 +142,8 @@ when complete.
|
|
| 135 |
- complete unified and token-ID Gather decoder exporters;
|
| 136 |
- per-output-channel INT8-QDQ converter and CPU parity checks;
|
| 137 |
- browser WebGPU end-to-end and decoder smoke harnesses;
|
|
|
|
|
|
|
| 138 |
- native CPU and MangaOCR comparison harnesses;
|
| 139 |
- isolated Python/npm dependencies and local server;
|
| 140 |
- benchmark helpers and generated validation reports.
|
|
|
|
| 30 |
## Latest recommended models
|
| 31 |
|
| 32 |
The latest decoder uses one complete six-layer QDQ graph for both prefill and
|
| 33 |
+
cached token steps. It accepts an INT32 token ID, gathers the selected INT8
|
| 34 |
+
embedding row first, and only then runs `DequantizeLinear`. This exact rewrite
|
| 35 |
+
avoids materializing the full `[14630,512]` FP32 embedding table (28.6 MiB).
|
| 36 |
|
| 37 |
| Folder | Vision encoder | Latest decoder | Latest ONNX download |
|
| 38 |
| --- | --- | --- | ---: |
|
|
|
|
| 62 |
tested locally but removed because both were slower on the tested runtime.
|
| 63 |
|
| 64 |
The QDQ decoder stores INT8 weights and executes `DequantizeLinear` followed by
|
| 65 |
+
FP32 `MatMul` on WebGPU. The token embedding is the exception: its INT8 Gather
|
| 66 |
+
runs before dequantization. It contains neither `DynamicQuantizeLinear` nor
|
| 67 |
`MatMulInteger`, which would leave the intended browser WebGPU path.
|
| 68 |
|
| 69 |
## 13 difficult showcase-image check
|
|
|
|
| 84 |
242 WebGPU result of 11.76% nCER and 5/13 exact. Latency depends heavily on GPU,
|
| 85 |
shader cache, browser, and system load; these numbers are not portable.
|
| 86 |
|
| 87 |
+
The Gather-before-dequantize decoder is also bit-exact against the previous
|
| 88 |
+
token-Gather decoder on ONNX Runtime CPU at prefill, cache length 257, and the
|
| 89 |
+
maximum tested cache length 383. Three full WebGPU runs per vision tier retained
|
| 90 |
+
all 13 previous output strings.
|
| 91 |
+
|
| 92 |
In one same-process 121 comparison, the optimized graph sampled an 833 MB
|
| 93 |
shared GPU-process peak and 577 MB renderer peak, versus 866 MB and 613 MB for
|
| 94 |
the unified one-hot baseline. Those values must not be added as model-owned
|
|
|
|
| 142 |
- complete unified and token-ID Gather decoder exporters;
|
| 143 |
- per-output-channel INT8-QDQ converter and CPU parity checks;
|
| 144 |
- browser WebGPU end-to-end and decoder smoke harnesses;
|
| 145 |
+
- exact Gather-before-dequantize and experimental execution optimizers;
|
| 146 |
+
- staged vision-worker and decoder-only cold/warm memory harnesses;
|
| 147 |
- native CPU and MangaOCR comparison harnesses;
|
| 148 |
- isolated Python/npm dependencies and local server;
|
| 149 |
- benchmark helpers and generated validation reports.
|
README_MEMORY_OPT.md
ADDED
|
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Baberu WebGPU capability-preserving memory experiment
|
| 2 |
+
|
| 3 |
+
This experiment tests browser-memory reductions without pruning, narrowing,
|
| 4 |
+
distilling, shortening, or changing either Baberu model tier. Both variants
|
| 5 |
+
retain the complete six-layer decoder, hidden size 512, intermediate size
|
| 6 |
+
1,536, eight attention heads, two KV heads, 14,630-token vocabulary, 128-token
|
| 7 |
+
generation limit, published greedy/repetition policy, and their original INT4
|
| 8 |
+
or FP16 vision encoder.
|
| 9 |
+
|
| 10 |
+
Implemented runtime-only variables using the exact `onnxruntime-web/webgpu`
|
| 11 |
+
entry that Baberu's browser harness uses:
|
| 12 |
+
|
| 13 |
+
- `loader=url` passes the model URL directly to ONNX Runtime and avoids an
|
| 14 |
+
additional page-owned `ArrayBuffer`/`Uint8Array` load path;
|
| 15 |
+
- `loader=bytes` keeps the previous explicit `fetch -> ArrayBuffer -> Uint8Array`
|
| 16 |
+
path as the control;
|
| 17 |
+
- `memory=default` uses ONNX Runtime's default CPU arena and memory pattern;
|
| 18 |
+
- `memory=lean` disables the CPU arena and memory pattern to limit reusable
|
| 19 |
+
CPU-side allocation pools while leaving all WebGPU operators and model
|
| 20 |
+
tensors unchanged;
|
| 21 |
+
- `schedule=resident` keeps the vision and decoder sessions resident together;
|
| 22 |
+
- `schedule=staged` runs the vision phase in a dedicated worker, encodes all
|
| 23 |
+
page crops to CPU embeddings, then releases the session and terminates the
|
| 24 |
+
worker before loading the decoder. Worker termination returns the vision
|
| 25 |
+
WebAssembly heap instead of leaving its high-water allocation in the main
|
| 26 |
+
runtime. The embeddings total about 0.5 MiB per crop and retain the complete
|
| 27 |
+
vision output, so model capability is unchanged.
|
| 28 |
+
|
| 29 |
+
The exact runtime entry and its Asyncify WebGPU assets are copied under
|
| 30 |
+
`.work/`; the installed package and production OCR implementation are not
|
| 31 |
+
modified. A prior draft attempted to patch a JSEP buffer pool, but that is a
|
| 32 |
+
different build path from `onnxruntime-web/webgpu` and is intentionally not
|
| 33 |
+
part of this experiment.
|
| 34 |
+
|
| 35 |
+
The decoder query exposes model-execution candidates separately:
|
| 36 |
+
|
| 37 |
+
- `decoder=gather-opt` gathers the selected INT8 token-embedding row before
|
| 38 |
+
`DequantizeLinear`, avoiding a 28.6 MiB full FP32 embedding-table output.
|
| 39 |
+
It is bit-exact on CPU and output-exact on all 13 browser cases for both
|
| 40 |
+
vision tiers, so this is the accepted model-level rewrite.
|
| 41 |
+
- `decoder=fixed-kv` also exposes fixed 384-slot KV inputs and outputs plus
|
| 42 |
+
`past_length`. This remains a rejected experiment: ONNX Runtime WebGPU does
|
| 43 |
+
not share the external buffers in place, so dynamic Slice/Pad work offsets
|
| 44 |
+
the theoretical allocation saving.
|
| 45 |
+
- `decoder=fp16-matmul` executes the 21 MatMuls in decoder layers 1–3 with
|
| 46 |
+
static FP16 weights while retaining the other layers and language head in
|
| 47 |
+
their original INT8-QDQ form. It preserves the current browser corpus but is
|
| 48 |
+
retained as a rejected memory experiment because isolated peak RSS rose.
|
| 49 |
+
More aggressive 42/43-MatMul variants changed an OCR output.
|
| 50 |
+
|
| 51 |
+
Prepare disposable assets and start the isolated server:
|
| 52 |
+
|
| 53 |
+
```sh
|
| 54 |
+
python3 prepare_memory_experiment.py
|
| 55 |
+
uv venv .work/model-opt-venv
|
| 56 |
+
uv pip install --python .work/model-opt-venv/bin/python -r requirements-model-opt.txt
|
| 57 |
+
.work/model-opt-venv/bin/python optimize_decoder_execution.py
|
| 58 |
+
.work/model-opt-venv/bin/python optimize_decoder_fp16_matmul.py
|
| 59 |
+
python3 serve_memory_experiment.py
|
| 60 |
+
```
|
| 61 |
+
|
| 62 |
+
Example control and optimized URLs:
|
| 63 |
+
|
| 64 |
+
```text
|
| 65 |
+
http://127.0.0.1:8765/webgpu-memory-opt.html?tier=121&memory=default&loader=bytes&schedule=resident
|
| 66 |
+
http://127.0.0.1:8765/webgpu-memory-opt.html?tier=121&memory=lean&loader=url&schedule=staged
|
| 67 |
+
http://127.0.0.1:8765/webgpu-memory-opt.html?tier=242&memory=default&loader=bytes&schedule=resident
|
| 68 |
+
http://127.0.0.1:8765/webgpu-memory-opt.html?tier=242&memory=lean&loader=url&schedule=staged
|
| 69 |
+
http://127.0.0.1:8765/webgpu-memory-opt.html?tier=242&memory=lean&loader=url&schedule=staged&decoder=gather-opt
|
| 70 |
+
http://127.0.0.1:8765/webgpu-memory-opt.html?tier=121&memory=lean&loader=url&schedule=staged&decoder=gather-opt&scope=decoder
|
| 71 |
+
```
|
| 72 |
+
|
| 73 |
+
Results are written to `.work/results/`. Each result includes OCR output,
|
| 74 |
+
accuracy, latency, browser-reported heap samples, and sampled T3 renderer/GPU
|
| 75 |
+
process RSS. Process maxima remain engineering observations rather than exact
|
| 76 |
+
model-owned RAM, especially on unified-memory Apple hardware. `scope=decoder`
|
| 77 |
+
skips the vision model and forces the complete 128-token cache path for a
|
| 78 |
+
separate cold session-load and warm execution observation.
|
| 79 |
+
|
| 80 |
+
See `MEMORY_RESULTS.md` for the repeated control/optimized comparison and run
|
| 81 |
+
`python3 summarize_memory_results.py` to recompute it from local results.
|
| 82 |
+
See `MODEL_EXECUTION_RESULTS.md` for the graph-level acceptance decision.
|
optimize_decoder_execution.py
ADDED
|
@@ -0,0 +1,345 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import hashlib
|
| 5 |
+
import json
|
| 6 |
+
from collections import Counter, defaultdict
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
|
| 9 |
+
import numpy as np
|
| 10 |
+
import onnx
|
| 11 |
+
import onnxruntime as ort
|
| 12 |
+
from onnx import TensorProto, helper, numpy_helper
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
ROOT = Path(__file__).resolve().parent
|
| 16 |
+
DEFAULT_SOURCE = ROOT / ".work/models/shared/decoder_unified_gather_qdq_int8.onnx"
|
| 17 |
+
DEFAULT_DESTINATION = ROOT / ".work/models/model-opt/decoder_gather_before_dq_int8.onnx"
|
| 18 |
+
DEFAULT_FIXED_KV_DESTINATION = ROOT / ".work/models/model-opt/decoder_gather_dq_fixed_kv_int8.onnx"
|
| 19 |
+
DEFAULT_REPORT = ROOT / ".work/reports/model-execution-optimization.json"
|
| 20 |
+
NUM_LAYERS = 6
|
| 21 |
+
VISION_TOKENS = 256
|
| 22 |
+
MAX_NEW_TOKENS = 128
|
| 23 |
+
# Prefill produces 257 cache entries. The runtime stops before running a decode
|
| 24 |
+
# step for token 128, so the largest present cache produced is length 384.
|
| 25 |
+
MAX_CACHE_LENGTH = VISION_TOKENS + MAX_NEW_TOKENS
|
| 26 |
+
CACHE_NAMES = [f"{kind}_{axis}{layer}" for kind in ("past", "present") for axis in ("k", "v") for layer in range(NUM_LAYERS)]
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def sha256(path: Path) -> str:
|
| 30 |
+
digest = hashlib.sha256()
|
| 31 |
+
with path.open("rb") as source:
|
| 32 |
+
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
| 33 |
+
digest.update(chunk)
|
| 34 |
+
return digest.hexdigest()
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def set_shape(value_info: onnx.ValueInfoProto, shape: list[int | str]) -> None:
|
| 38 |
+
dimensions = value_info.type.tensor_type.shape.dim
|
| 39 |
+
del dimensions[:]
|
| 40 |
+
for value in shape:
|
| 41 |
+
dimension = dimensions.add()
|
| 42 |
+
if isinstance(value, int):
|
| 43 |
+
dimension.dim_value = value
|
| 44 |
+
else:
|
| 45 |
+
dimension.dim_param = value
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def rewrite_embedding_gather(model: onnx.ModelProto) -> dict:
|
| 49 |
+
graph = model.graph
|
| 50 |
+
initializers = {value.name: value for value in graph.initializer}
|
| 51 |
+
consumers: dict[str, list[onnx.NodeProto]] = defaultdict(list)
|
| 52 |
+
for node in graph.node:
|
| 53 |
+
for name in node.input:
|
| 54 |
+
consumers[name].append(node)
|
| 55 |
+
|
| 56 |
+
embedding_dq = None
|
| 57 |
+
embedding_gather = None
|
| 58 |
+
for node in graph.node:
|
| 59 |
+
if node.op_type != "DequantizeLinear" or len(node.input) < 3:
|
| 60 |
+
continue
|
| 61 |
+
quantized = initializers.get(node.input[0])
|
| 62 |
+
if not quantized or list(quantized.dims) != [14630, 512]:
|
| 63 |
+
continue
|
| 64 |
+
matches = [consumer for consumer in consumers[node.output[0]] if consumer.op_type == "Gather"]
|
| 65 |
+
if len(matches) != 1:
|
| 66 |
+
raise RuntimeError(f"Embedding DQ expected one Gather consumer, found {len(matches)}")
|
| 67 |
+
embedding_dq = node
|
| 68 |
+
embedding_gather = matches[0]
|
| 69 |
+
break
|
| 70 |
+
if embedding_dq is None or embedding_gather is None:
|
| 71 |
+
raise RuntimeError("Could not locate quantized 14630x512 embedding Gather")
|
| 72 |
+
|
| 73 |
+
original_output = embedding_gather.output[0]
|
| 74 |
+
quantized_output = f"{original_output}_int8"
|
| 75 |
+
embedding_gather.input[0] = embedding_dq.input[0]
|
| 76 |
+
embedding_gather.output[0] = quantized_output
|
| 77 |
+
gathered_dq = helper.make_node(
|
| 78 |
+
"DequantizeLinear",
|
| 79 |
+
[quantized_output, embedding_dq.input[1], embedding_dq.input[2]],
|
| 80 |
+
[original_output],
|
| 81 |
+
name=f"{embedding_dq.name}_after_gather",
|
| 82 |
+
# Gather axis 0 with token_ids rank 2 moves hidden axis 1 to axis 2.
|
| 83 |
+
axis=2,
|
| 84 |
+
)
|
| 85 |
+
|
| 86 |
+
rewritten = []
|
| 87 |
+
for node in graph.node:
|
| 88 |
+
if node is embedding_dq:
|
| 89 |
+
continue
|
| 90 |
+
rewritten.append(node)
|
| 91 |
+
if node is embedding_gather:
|
| 92 |
+
rewritten.append(gathered_dq)
|
| 93 |
+
del graph.node[:]
|
| 94 |
+
graph.node.extend(rewritten)
|
| 95 |
+
return {
|
| 96 |
+
"quantized_elements_selected": 14630 * 512,
|
| 97 |
+
"fp32_bytes_avoided_before_gather": 14630 * 512 * 4,
|
| 98 |
+
"gather_output": original_output,
|
| 99 |
+
}
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def rewrite_fixed_kv_io(model: onnx.ModelProto) -> dict:
|
| 103 |
+
graph = model.graph
|
| 104 |
+
input_by_name = {value.name: value for value in graph.input}
|
| 105 |
+
output_by_name = {value.name: value for value in graph.output}
|
| 106 |
+
past_names = [f"past_{axis}{layer}" for axis in ("k", "v") for layer in range(NUM_LAYERS)]
|
| 107 |
+
present_names = [f"present_{axis}{layer}" for axis in ("k", "v") for layer in range(NUM_LAYERS)]
|
| 108 |
+
|
| 109 |
+
for name in past_names:
|
| 110 |
+
if name not in input_by_name:
|
| 111 |
+
raise RuntimeError(f"Missing cache input {name}")
|
| 112 |
+
set_shape(input_by_name[name], [1, 2, MAX_CACHE_LENGTH, 64])
|
| 113 |
+
for name in present_names:
|
| 114 |
+
if name not in output_by_name:
|
| 115 |
+
raise RuntimeError(f"Missing cache output {name}")
|
| 116 |
+
set_shape(output_by_name[name], [1, 2, MAX_CACHE_LENGTH, 64])
|
| 117 |
+
|
| 118 |
+
graph.input.append(helper.make_tensor_value_info("past_length", TensorProto.INT64, [1]))
|
| 119 |
+
graph.initializer.extend(
|
| 120 |
+
[
|
| 121 |
+
numpy_helper.from_array(np.array([0], dtype=np.int64), "fixed_kv_slice_starts"),
|
| 122 |
+
numpy_helper.from_array(np.array([2], dtype=np.int64), "fixed_kv_slice_axes"),
|
| 123 |
+
numpy_helper.from_array(np.array([1], dtype=np.int64), "fixed_kv_slice_steps"),
|
| 124 |
+
numpy_helper.from_array(np.zeros(4, dtype=np.int64), "fixed_kv_pad_begin"),
|
| 125 |
+
numpy_helper.from_array(np.zeros(2, dtype=np.int64), "fixed_kv_pad_end_prefix"),
|
| 126 |
+
numpy_helper.from_array(np.zeros(1, dtype=np.int64), "fixed_kv_pad_end_suffix"),
|
| 127 |
+
numpy_helper.from_array(np.array([MAX_CACHE_LENGTH], dtype=np.int64), "fixed_kv_max_length"),
|
| 128 |
+
numpy_helper.from_array(np.array([2], dtype=np.int64), "fixed_kv_shape_axis"),
|
| 129 |
+
]
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
slice_nodes = []
|
| 133 |
+
sliced_names: dict[str, str] = {}
|
| 134 |
+
for name in past_names:
|
| 135 |
+
sliced_name = f"{name}_valid"
|
| 136 |
+
sliced_names[name] = sliced_name
|
| 137 |
+
slice_nodes.append(
|
| 138 |
+
helper.make_node(
|
| 139 |
+
"Slice",
|
| 140 |
+
[name, "fixed_kv_slice_starts", "past_length", "fixed_kv_slice_axes", "fixed_kv_slice_steps"],
|
| 141 |
+
[sliced_name],
|
| 142 |
+
name=f"fixed_kv_slice_{name}",
|
| 143 |
+
)
|
| 144 |
+
)
|
| 145 |
+
|
| 146 |
+
original_nodes = list(graph.node)
|
| 147 |
+
for node in original_nodes:
|
| 148 |
+
for index, name in enumerate(node.input):
|
| 149 |
+
if name in sliced_names:
|
| 150 |
+
node.input[index] = sliced_names[name]
|
| 151 |
+
|
| 152 |
+
pad_nodes = []
|
| 153 |
+
for name in present_names:
|
| 154 |
+
producer = next((node for node in original_nodes if name in node.output), None)
|
| 155 |
+
if producer is None:
|
| 156 |
+
raise RuntimeError(f"Missing producer for {name}")
|
| 157 |
+
valid_name = f"{name}_valid"
|
| 158 |
+
producer.output[list(producer.output).index(name)] = valid_name
|
| 159 |
+
for consumer in original_nodes:
|
| 160 |
+
for input_index, input_name in enumerate(consumer.input):
|
| 161 |
+
if input_name == name:
|
| 162 |
+
consumer.input[input_index] = valid_name
|
| 163 |
+
shape_name = f"{name}_shape"
|
| 164 |
+
length_name = f"{name}_length"
|
| 165 |
+
padding_name = f"{name}_padding"
|
| 166 |
+
pads_name = f"{name}_pads"
|
| 167 |
+
pad_nodes.extend(
|
| 168 |
+
[
|
| 169 |
+
helper.make_node("Shape", [valid_name], [shape_name], name=f"fixed_kv_shape_{name}"),
|
| 170 |
+
helper.make_node(
|
| 171 |
+
"Gather",
|
| 172 |
+
[shape_name, "fixed_kv_shape_axis"],
|
| 173 |
+
[length_name],
|
| 174 |
+
name=f"fixed_kv_length_{name}",
|
| 175 |
+
axis=0,
|
| 176 |
+
),
|
| 177 |
+
helper.make_node(
|
| 178 |
+
"Sub",
|
| 179 |
+
["fixed_kv_max_length", length_name],
|
| 180 |
+
[padding_name],
|
| 181 |
+
name=f"fixed_kv_padding_{name}",
|
| 182 |
+
),
|
| 183 |
+
helper.make_node(
|
| 184 |
+
"Concat",
|
| 185 |
+
["fixed_kv_pad_begin", "fixed_kv_pad_end_prefix", padding_name, "fixed_kv_pad_end_suffix"],
|
| 186 |
+
[pads_name],
|
| 187 |
+
name=f"fixed_kv_pads_{name}",
|
| 188 |
+
axis=0,
|
| 189 |
+
),
|
| 190 |
+
helper.make_node("Pad", [valid_name, pads_name], [name], name=f"fixed_kv_pad_{name}"),
|
| 191 |
+
]
|
| 192 |
+
)
|
| 193 |
+
|
| 194 |
+
del graph.node[:]
|
| 195 |
+
graph.node.extend(slice_nodes)
|
| 196 |
+
graph.node.extend(original_nodes)
|
| 197 |
+
graph.node.extend(pad_nodes)
|
| 198 |
+
return {
|
| 199 |
+
"max_cache_length": MAX_CACHE_LENGTH,
|
| 200 |
+
"fixed_live_cache_bytes": NUM_LAYERS * 2 * 2 * MAX_CACHE_LENGTH * 64 * 4,
|
| 201 |
+
"dynamic_present_allocation_bytes_across_max_decode": sum(
|
| 202 |
+
NUM_LAYERS * 2 * 2 * length * 64 * 4
|
| 203 |
+
for length in range(VISION_TOKENS + 1, MAX_CACHE_LENGTH + 1)
|
| 204 |
+
),
|
| 205 |
+
}
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def operator_counts(model: onnx.ModelProto) -> dict[str, int]:
|
| 209 |
+
return dict(sorted(Counter(node.op_type for node in model.graph.node).items()))
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def source_feeds(session: ort.InferenceSession, past_length: int, *, prefill: bool) -> dict[str, np.ndarray]:
|
| 213 |
+
rng = np.random.default_rng(20260717 + past_length)
|
| 214 |
+
feeds: dict[str, np.ndarray] = {}
|
| 215 |
+
for value in session.get_inputs():
|
| 216 |
+
if value.name == "vision_embeds":
|
| 217 |
+
length = VISION_TOKENS if prefill else 0
|
| 218 |
+
feeds[value.name] = rng.normal(0, 0.2, [1, length, 512]).astype(np.float32)
|
| 219 |
+
elif value.name == "token_ids":
|
| 220 |
+
feeds[value.name] = np.array([[1 if prefill else 4]], dtype=np.int32)
|
| 221 |
+
elif value.name == "position_ids":
|
| 222 |
+
feeds[value.name] = (
|
| 223 |
+
np.arange(VISION_TOKENS + 1, dtype=np.int32)[None, :]
|
| 224 |
+
if prefill
|
| 225 |
+
else np.array([[past_length]], dtype=np.int32)
|
| 226 |
+
)
|
| 227 |
+
else:
|
| 228 |
+
feeds[value.name] = rng.normal(0, 0.02, [1, 2, past_length, 64]).astype(np.float32)
|
| 229 |
+
return feeds
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def candidate_feeds(source: dict[str, np.ndarray], past_length: int) -> dict[str, np.ndarray]:
|
| 233 |
+
feeds: dict[str, np.ndarray] = {}
|
| 234 |
+
for name, value in source.items():
|
| 235 |
+
if name.startswith("past_"):
|
| 236 |
+
fixed = np.zeros([1, 2, MAX_CACHE_LENGTH, 64], dtype=np.float32)
|
| 237 |
+
fixed[:, :, :past_length, :] = value
|
| 238 |
+
feeds[name] = fixed
|
| 239 |
+
else:
|
| 240 |
+
feeds[name] = value
|
| 241 |
+
feeds["past_length"] = np.array([past_length], dtype=np.int64)
|
| 242 |
+
return feeds
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
def validate_cpu(source: Path, candidate: Path, *, fixed_kv: bool) -> dict:
|
| 246 |
+
source_session = ort.InferenceSession(str(source), providers=["CPUExecutionProvider"])
|
| 247 |
+
candidate_session = ort.InferenceSession(str(candidate), providers=["CPUExecutionProvider"])
|
| 248 |
+
checks = {}
|
| 249 |
+
for label, past_length, prefill in (("prefill", 0, True), ("step_257", 257, False), ("step_383", 383, False)):
|
| 250 |
+
source_input = source_feeds(source_session, past_length, prefill=prefill)
|
| 251 |
+
expected = source_session.run(None, source_input)
|
| 252 |
+
actual = candidate_session.run(
|
| 253 |
+
None,
|
| 254 |
+
candidate_feeds(source_input, past_length) if fixed_kv else source_input,
|
| 255 |
+
)
|
| 256 |
+
active_length = (VISION_TOKENS + 1) if prefill else past_length + 1
|
| 257 |
+
differences = [float(np.max(np.abs(expected[0] - actual[0])))]
|
| 258 |
+
padding_nonzero = 0
|
| 259 |
+
for expected_cache, actual_cache in zip(expected[1:], actual[1:]):
|
| 260 |
+
actual_valid = actual_cache[:, :, :active_length, :] if fixed_kv else actual_cache
|
| 261 |
+
differences.append(float(np.max(np.abs(expected_cache - actual_valid))))
|
| 262 |
+
if fixed_kv:
|
| 263 |
+
padding_nonzero += int(np.count_nonzero(actual_cache[:, :, active_length:, :]))
|
| 264 |
+
checks[label] = {
|
| 265 |
+
"max_abs": max(differences),
|
| 266 |
+
"logits_max_abs": differences[0],
|
| 267 |
+
"top_token_source": int(expected[0][0, -1].argmax()),
|
| 268 |
+
"top_token_candidate": int(actual[0][0, -1].argmax()),
|
| 269 |
+
"padding_nonzero": padding_nonzero,
|
| 270 |
+
}
|
| 271 |
+
if checks[label]["max_abs"] != 0 or padding_nonzero != 0:
|
| 272 |
+
raise RuntimeError(f"CPU parity failed for {label}: {checks[label]}")
|
| 273 |
+
return checks
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def main() -> None:
|
| 277 |
+
parser = argparse.ArgumentParser()
|
| 278 |
+
parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
|
| 279 |
+
parser.add_argument("--destination", type=Path, default=DEFAULT_DESTINATION)
|
| 280 |
+
parser.add_argument("--fixed-kv-destination", type=Path, default=DEFAULT_FIXED_KV_DESTINATION)
|
| 281 |
+
parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
|
| 282 |
+
arguments = parser.parse_args()
|
| 283 |
+
arguments.destination.parent.mkdir(parents=True, exist_ok=True)
|
| 284 |
+
arguments.fixed_kv_destination.parent.mkdir(parents=True, exist_ok=True)
|
| 285 |
+
arguments.report.parent.mkdir(parents=True, exist_ok=True)
|
| 286 |
+
|
| 287 |
+
model = onnx.load(arguments.source)
|
| 288 |
+
before = operator_counts(model)
|
| 289 |
+
embedding = rewrite_embedding_gather(model)
|
| 290 |
+
model.producer_name = "vibe-manga-baberu-webgpu-model-execution-opt"
|
| 291 |
+
model.producer_version = "1"
|
| 292 |
+
onnx.checker.check_model(model)
|
| 293 |
+
onnx.save(model, arguments.destination)
|
| 294 |
+
gather_parity = validate_cpu(arguments.source, arguments.destination, fixed_kv=False)
|
| 295 |
+
|
| 296 |
+
fixed_kv_model = onnx.load(arguments.destination)
|
| 297 |
+
fixed_kv = rewrite_fixed_kv_io(fixed_kv_model)
|
| 298 |
+
fixed_kv_model.producer_version = "1-fixed-kv-experiment"
|
| 299 |
+
onnx.checker.check_model(fixed_kv_model)
|
| 300 |
+
onnx.save(fixed_kv_model, arguments.fixed_kv_destination)
|
| 301 |
+
fixed_kv_parity = validate_cpu(
|
| 302 |
+
arguments.source,
|
| 303 |
+
arguments.fixed_kv_destination,
|
| 304 |
+
fixed_kv=True,
|
| 305 |
+
)
|
| 306 |
+
report = {
|
| 307 |
+
"source": {
|
| 308 |
+
"path": str(arguments.source.relative_to(ROOT)),
|
| 309 |
+
"bytes": arguments.source.stat().st_size,
|
| 310 |
+
"sha256": sha256(arguments.source),
|
| 311 |
+
"operators": before,
|
| 312 |
+
},
|
| 313 |
+
"gather_optimized": {
|
| 314 |
+
"path": str(arguments.destination.relative_to(ROOT)),
|
| 315 |
+
"bytes": arguments.destination.stat().st_size,
|
| 316 |
+
"sha256": sha256(arguments.destination),
|
| 317 |
+
"operators": operator_counts(model),
|
| 318 |
+
},
|
| 319 |
+
"fixed_kv_experiment": {
|
| 320 |
+
"path": str(arguments.fixed_kv_destination.relative_to(ROOT)),
|
| 321 |
+
"bytes": arguments.fixed_kv_destination.stat().st_size,
|
| 322 |
+
"sha256": sha256(arguments.fixed_kv_destination),
|
| 323 |
+
"operators": operator_counts(fixed_kv_model),
|
| 324 |
+
},
|
| 325 |
+
"capability": {
|
| 326 |
+
"layers": 6,
|
| 327 |
+
"hidden_size": 512,
|
| 328 |
+
"kv_heads": 2,
|
| 329 |
+
"vocabulary": 14630,
|
| 330 |
+
"max_new_tokens": MAX_NEW_TOKENS,
|
| 331 |
+
"weights_requantized": False,
|
| 332 |
+
},
|
| 333 |
+
"embedding_gather_before_dequantize": embedding,
|
| 334 |
+
"fixed_kv_io": fixed_kv,
|
| 335 |
+
"cpu_parity": {
|
| 336 |
+
"gather_optimized": gather_parity,
|
| 337 |
+
"fixed_kv_experiment": fixed_kv_parity,
|
| 338 |
+
},
|
| 339 |
+
}
|
| 340 |
+
arguments.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
| 341 |
+
print(json.dumps(report, indent=2))
|
| 342 |
+
|
| 343 |
+
|
| 344 |
+
if __name__ == "__main__":
|
| 345 |
+
main()
|
optimize_decoder_fp16_matmul.py
ADDED
|
@@ -0,0 +1,242 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import hashlib
|
| 5 |
+
import json
|
| 6 |
+
import re
|
| 7 |
+
from collections import defaultdict
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
import numpy as np
|
| 11 |
+
import onnx
|
| 12 |
+
import onnxruntime as ort
|
| 13 |
+
from onnx import TensorProto, helper, numpy_helper
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
ROOT = Path(__file__).resolve().parent
|
| 17 |
+
BASELINE = ROOT / ".work/models/shared/decoder_unified_gather_qdq_int8.onnx"
|
| 18 |
+
DEFAULT_SOURCE = ROOT / ".work/models/model-opt/decoder_gather_before_dq_int8.onnx"
|
| 19 |
+
DEFAULT_DESTINATION = ROOT / ".work/models/model-opt/decoder_static_fp16_matmul.onnx"
|
| 20 |
+
DEFAULT_REPORT = ROOT / ".work/reports/model-fp16-matmul-optimization.json"
|
| 21 |
+
VISION_TOKENS = 256
|
| 22 |
+
KEEP_QDQ_MATMUL_NAMES = {"/lm_head/MatMul"}
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def sha256(path: Path) -> str:
|
| 26 |
+
digest = hashlib.sha256()
|
| 27 |
+
with path.open("rb") as source:
|
| 28 |
+
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
| 29 |
+
digest.update(chunk)
|
| 30 |
+
return digest.hexdigest()
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
def dequantize_to_fp16(
|
| 34 |
+
quantized: onnx.TensorProto,
|
| 35 |
+
scale: onnx.TensorProto,
|
| 36 |
+
zero_point: onnx.TensorProto,
|
| 37 |
+
axis: int,
|
| 38 |
+
) -> np.ndarray:
|
| 39 |
+
values = numpy_helper.to_array(quantized).astype(np.float32)
|
| 40 |
+
scales = numpy_helper.to_array(scale).astype(np.float32)
|
| 41 |
+
zeros = numpy_helper.to_array(zero_point).astype(np.float32)
|
| 42 |
+
broadcast_shape = [1] * values.ndim
|
| 43 |
+
broadcast_shape[axis] = scales.size
|
| 44 |
+
return ((values - zeros.reshape(broadcast_shape)) * scales.reshape(broadcast_shape)).astype(np.float16)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def rewrite_matmul_weights(model: onnx.ModelProto, selected_layers: set[int]) -> dict:
|
| 48 |
+
graph = model.graph
|
| 49 |
+
initializers = {value.name: value for value in graph.initializer}
|
| 50 |
+
consumers: dict[str, list[onnx.NodeProto]] = defaultdict(list)
|
| 51 |
+
for node in graph.node:
|
| 52 |
+
for name in node.input:
|
| 53 |
+
consumers[name].append(node)
|
| 54 |
+
|
| 55 |
+
dq_by_matmul_name: dict[str, onnx.NodeProto] = {}
|
| 56 |
+
fp16_initializers: list[onnx.TensorProto] = []
|
| 57 |
+
removable_initializers: set[str] = set()
|
| 58 |
+
quantized_elements = 0
|
| 59 |
+
for node in graph.node:
|
| 60 |
+
if node.op_type != "DequantizeLinear" or len(node.input) < 3:
|
| 61 |
+
continue
|
| 62 |
+
quantized = initializers.get(node.input[0])
|
| 63 |
+
scale = initializers.get(node.input[1])
|
| 64 |
+
zero_point = initializers.get(node.input[2])
|
| 65 |
+
if quantized is None or scale is None or zero_point is None:
|
| 66 |
+
continue
|
| 67 |
+
matches = [consumer for consumer in consumers[node.output[0]] if consumer.op_type == "MatMul"]
|
| 68 |
+
if not matches:
|
| 69 |
+
continue
|
| 70 |
+
if len(matches) != 1:
|
| 71 |
+
raise RuntimeError(f"Expected one MatMul consumer for {node.name}, found {len(matches)}")
|
| 72 |
+
matmul = matches[0]
|
| 73 |
+
if matmul.name in KEEP_QDQ_MATMUL_NAMES:
|
| 74 |
+
continue
|
| 75 |
+
layer_match = re.search(r"/decoder/layers\.(\d+)/", matmul.name)
|
| 76 |
+
if layer_match is None:
|
| 77 |
+
raise RuntimeError(f"Unexpected non-layer MatMul {matmul.name}")
|
| 78 |
+
if int(layer_match.group(1)) not in selected_layers:
|
| 79 |
+
continue
|
| 80 |
+
if matmul.input[1] != node.output[0]:
|
| 81 |
+
raise RuntimeError(f"Quantized weight is not RHS input for {matmul.name}")
|
| 82 |
+
axis = next((attribute.i for attribute in node.attribute if attribute.name == "axis"), 1)
|
| 83 |
+
fp16_name = f"{quantized.name}_static_fp16"
|
| 84 |
+
fp16 = dequantize_to_fp16(quantized, scale, zero_point, axis)
|
| 85 |
+
fp16_initializers.append(numpy_helper.from_array(fp16, fp16_name))
|
| 86 |
+
matmul.input[1] = fp16_name
|
| 87 |
+
dq_by_matmul_name[matmul.name] = node
|
| 88 |
+
removable_initializers.update(node.input[:3])
|
| 89 |
+
quantized_elements += fp16.size
|
| 90 |
+
|
| 91 |
+
expected_matmuls = 7 * len(selected_layers)
|
| 92 |
+
if len(dq_by_matmul_name) != expected_matmuls:
|
| 93 |
+
raise RuntimeError(f"Expected {expected_matmuls} selected MatMuls, found {len(dq_by_matmul_name)}")
|
| 94 |
+
|
| 95 |
+
rewritten: list[onnx.NodeProto] = []
|
| 96 |
+
removed_dq_names = {node.name for node in dq_by_matmul_name.values()}
|
| 97 |
+
for node in graph.node:
|
| 98 |
+
if node.name in removed_dq_names:
|
| 99 |
+
continue
|
| 100 |
+
if node.name not in dq_by_matmul_name:
|
| 101 |
+
rewritten.append(node)
|
| 102 |
+
continue
|
| 103 |
+
input_fp16 = f"{node.input[0]}__for_{node.name.replace('/', '_')}_fp16"
|
| 104 |
+
output_float = node.output[0]
|
| 105 |
+
output_fp16 = f"{output_float}__fp16"
|
| 106 |
+
rewritten.append(
|
| 107 |
+
helper.make_node(
|
| 108 |
+
"Cast",
|
| 109 |
+
[node.input[0]],
|
| 110 |
+
[input_fp16],
|
| 111 |
+
name=f"{node.name}/CastInputToFp16",
|
| 112 |
+
to=TensorProto.FLOAT16,
|
| 113 |
+
)
|
| 114 |
+
)
|
| 115 |
+
node.input[0] = input_fp16
|
| 116 |
+
node.output[0] = output_fp16
|
| 117 |
+
rewritten.append(node)
|
| 118 |
+
rewritten.append(
|
| 119 |
+
helper.make_node(
|
| 120 |
+
"Cast",
|
| 121 |
+
[output_fp16],
|
| 122 |
+
[output_float],
|
| 123 |
+
name=f"{node.name}/CastOutputToFp32",
|
| 124 |
+
to=TensorProto.FLOAT,
|
| 125 |
+
)
|
| 126 |
+
)
|
| 127 |
+
|
| 128 |
+
retained = [value for value in graph.initializer if value.name not in removable_initializers]
|
| 129 |
+
del graph.initializer[:]
|
| 130 |
+
graph.initializer.extend(retained)
|
| 131 |
+
graph.initializer.extend(fp16_initializers)
|
| 132 |
+
del graph.node[:]
|
| 133 |
+
graph.node.extend(rewritten)
|
| 134 |
+
return {
|
| 135 |
+
"matmuls_rewritten": len(dq_by_matmul_name),
|
| 136 |
+
"layers_rewritten": sorted(selected_layers),
|
| 137 |
+
"matmuls_kept_qdq": sorted(KEEP_QDQ_MATMUL_NAMES),
|
| 138 |
+
"weight_elements": quantized_elements,
|
| 139 |
+
"runtime_fp32_weight_bytes_avoided": quantized_elements * 4,
|
| 140 |
+
"static_fp16_weight_bytes": quantized_elements * 2,
|
| 141 |
+
"source_int8_weight_bytes_removed": quantized_elements,
|
| 142 |
+
"estimated_live_weight_bytes_avoided": quantized_elements * 3,
|
| 143 |
+
}
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def feeds(session: ort.InferenceSession, past_length: int, *, prefill: bool) -> dict[str, np.ndarray]:
|
| 147 |
+
rng = np.random.default_rng(20260717 + past_length)
|
| 148 |
+
result: dict[str, np.ndarray] = {}
|
| 149 |
+
for value in session.get_inputs():
|
| 150 |
+
if value.name == "vision_embeds":
|
| 151 |
+
length = VISION_TOKENS if prefill else 0
|
| 152 |
+
result[value.name] = rng.normal(0, 0.2, [1, length, 512]).astype(np.float32)
|
| 153 |
+
elif value.name == "token_ids":
|
| 154 |
+
result[value.name] = np.array([[1 if prefill else 4]], dtype=np.int32)
|
| 155 |
+
elif value.name == "position_ids":
|
| 156 |
+
result[value.name] = (
|
| 157 |
+
np.arange(VISION_TOKENS + 1, dtype=np.int32)[None, :]
|
| 158 |
+
if prefill
|
| 159 |
+
else np.array([[past_length]], dtype=np.int32)
|
| 160 |
+
)
|
| 161 |
+
else:
|
| 162 |
+
result[value.name] = rng.normal(0, 0.02, [1, 2, past_length, 64]).astype(np.float32)
|
| 163 |
+
return result
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def validate_cpu(baseline: Path, candidate: Path) -> dict:
|
| 167 |
+
baseline_session = ort.InferenceSession(str(baseline), providers=["CPUExecutionProvider"])
|
| 168 |
+
candidate_session = ort.InferenceSession(str(candidate), providers=["CPUExecutionProvider"])
|
| 169 |
+
checks = {}
|
| 170 |
+
for label, past_length, prefill in (
|
| 171 |
+
("prefill", 0, True),
|
| 172 |
+
("step_257", 257, False),
|
| 173 |
+
("step_383", 383, False),
|
| 174 |
+
):
|
| 175 |
+
model_feeds = feeds(baseline_session, past_length, prefill=prefill)
|
| 176 |
+
expected = baseline_session.run(None, model_feeds)
|
| 177 |
+
actual = candidate_session.run(None, model_feeds)
|
| 178 |
+
max_abs = [float(np.max(np.abs(left - right))) for left, right in zip(expected, actual)]
|
| 179 |
+
checks[label] = {
|
| 180 |
+
"logits_max_abs": max_abs[0],
|
| 181 |
+
"all_outputs_max_abs": max(max_abs),
|
| 182 |
+
"top_token_baseline": int(expected[0][0, -1].argmax()),
|
| 183 |
+
"top_token_candidate": int(actual[0][0, -1].argmax()),
|
| 184 |
+
}
|
| 185 |
+
if checks[label]["top_token_baseline"] != checks[label]["top_token_candidate"]:
|
| 186 |
+
raise RuntimeError(f"CPU top-token parity failed for {label}: {checks[label]}")
|
| 187 |
+
return checks
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def main() -> None:
|
| 191 |
+
parser = argparse.ArgumentParser()
|
| 192 |
+
parser.add_argument("--baseline", type=Path, default=BASELINE)
|
| 193 |
+
parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
|
| 194 |
+
parser.add_argument("--destination", type=Path, default=DEFAULT_DESTINATION)
|
| 195 |
+
parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
|
| 196 |
+
parser.add_argument(
|
| 197 |
+
"--layers",
|
| 198 |
+
default="1,2,3",
|
| 199 |
+
help="Comma-separated decoder layers whose seven MatMuls execute in FP16",
|
| 200 |
+
)
|
| 201 |
+
arguments = parser.parse_args()
|
| 202 |
+
arguments.destination.parent.mkdir(parents=True, exist_ok=True)
|
| 203 |
+
arguments.report.parent.mkdir(parents=True, exist_ok=True)
|
| 204 |
+
|
| 205 |
+
selected_layers = {int(value) for value in arguments.layers.split(",") if value != ""}
|
| 206 |
+
if not selected_layers.issubset(set(range(6))):
|
| 207 |
+
raise ValueError(f"Invalid decoder layers {sorted(selected_layers)}")
|
| 208 |
+
model = onnx.load(arguments.source)
|
| 209 |
+
optimization = rewrite_matmul_weights(model, selected_layers)
|
| 210 |
+
model.producer_name = "vibe-manga-baberu-webgpu-static-fp16-matmul"
|
| 211 |
+
model.producer_version = "1"
|
| 212 |
+
onnx.checker.check_model(model)
|
| 213 |
+
onnx.save(model, arguments.destination)
|
| 214 |
+
parity = validate_cpu(arguments.baseline, arguments.destination)
|
| 215 |
+
report = {
|
| 216 |
+
"source": {
|
| 217 |
+
"path": str(arguments.source.relative_to(ROOT)),
|
| 218 |
+
"bytes": arguments.source.stat().st_size,
|
| 219 |
+
"sha256": sha256(arguments.source),
|
| 220 |
+
},
|
| 221 |
+
"optimized": {
|
| 222 |
+
"path": str(arguments.destination.relative_to(ROOT)),
|
| 223 |
+
"bytes": arguments.destination.stat().st_size,
|
| 224 |
+
"sha256": sha256(arguments.destination),
|
| 225 |
+
},
|
| 226 |
+
"capability": {
|
| 227 |
+
"layers": 6,
|
| 228 |
+
"hidden_size": 512,
|
| 229 |
+
"kv_heads": 2,
|
| 230 |
+
"vocabulary": 14630,
|
| 231 |
+
"max_new_tokens": 128,
|
| 232 |
+
"architecture_changed": False,
|
| 233 |
+
},
|
| 234 |
+
"static_fp16_matmul": optimization,
|
| 235 |
+
"cpu_validation": parity,
|
| 236 |
+
}
|
| 237 |
+
arguments.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
| 238 |
+
print(json.dumps(report, indent=2))
|
| 239 |
+
|
| 240 |
+
|
| 241 |
+
if __name__ == "__main__":
|
| 242 |
+
main()
|
reports/model-execution-report.json
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"date": "2026-07-17",
|
| 3 |
+
"runtime": "onnxruntime-web/webgpu 1.27.0",
|
| 4 |
+
"source": {
|
| 5 |
+
"bytes": 34647907,
|
| 6 |
+
"sha256": "3612d5787ffb1026ae09bfd4a16b3379ea75fe0dd803484f7bb41645be6b6239"
|
| 7 |
+
},
|
| 8 |
+
"gather_optimized": {
|
| 9 |
+
"bytes": 34647902,
|
| 10 |
+
"sha256": "53778e5b50d78ee55c8674d3d19f53de14da808bd1c231a9d4cebd5d44d3ad93",
|
| 11 |
+
"fp32_intermediate_bytes_avoided": 29962240,
|
| 12 |
+
"cpu_bit_exact": true,
|
| 13 |
+
"webgpu_output_exact_tier_121": true,
|
| 14 |
+
"webgpu_output_exact_tier_242": true,
|
| 15 |
+
"accepted": true
|
| 16 |
+
},
|
| 17 |
+
"fixed_kv": {
|
| 18 |
+
"bytes": 34656423,
|
| 19 |
+
"sha256": "6e3ae602fdebebe85f379ece1b24804ded32498ff94c950f7f1dda88423700f8",
|
| 20 |
+
"cpu_bit_exact": true,
|
| 21 |
+
"webgpu_output_exact": true,
|
| 22 |
+
"accepted": false,
|
| 23 |
+
"reason": "No in-place past-present buffer sharing in standard ONNX Runtime WebGPU; Slice and Pad offset the allocation benefit."
|
| 24 |
+
},
|
| 25 |
+
"fp16_matmul_layers_1_2_3": {
|
| 26 |
+
"bytes": 43622628,
|
| 27 |
+
"sha256": "eef947f0eeb23c843f02d26c6abf6bc69d4cdde396b43846ca4802ec4810f9dc",
|
| 28 |
+
"matmuls_rewritten": 21,
|
| 29 |
+
"estimated_live_weight_bytes_avoided": 27131904,
|
| 30 |
+
"webgpu_output_exact_tier_121": true,
|
| 31 |
+
"webgpu_output_exact_tier_242": true,
|
| 32 |
+
"isolated_peak_mib_baseline": 727.1,
|
| 33 |
+
"isolated_peak_mib_candidate": 755.5,
|
| 34 |
+
"accepted": false,
|
| 35 |
+
"reason": "Measured isolated peak RSS regressed by 3.9 percent."
|
| 36 |
+
}
|
| 37 |
+
}
|
requirements-model-opt.txt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
numpy==2.5.1
|
| 2 |
+
onnx==1.22.0
|
| 3 |
+
onnxruntime==1.27.0
|
serve_memory_experiment.py
ADDED
|
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import subprocess
|
| 5 |
+
import threading
|
| 6 |
+
import time
|
| 7 |
+
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parent
|
| 12 |
+
RESULTS = ROOT / ".work" / "results"
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def process_rss() -> dict[str, dict[int, int]]:
|
| 16 |
+
output = subprocess.run(
|
| 17 |
+
["ps", "-axo", "pid=,rss=,command="],
|
| 18 |
+
check=True,
|
| 19 |
+
capture_output=True,
|
| 20 |
+
text=True,
|
| 21 |
+
).stdout
|
| 22 |
+
result: dict[str, dict[int, int]] = {"renderer": {}, "gpu": {}}
|
| 23 |
+
for line in output.splitlines():
|
| 24 |
+
pieces = line.strip().split(maxsplit=2)
|
| 25 |
+
if len(pieces) != 3:
|
| 26 |
+
continue
|
| 27 |
+
pid_text, rss_text, command = pieces
|
| 28 |
+
if "T3 Code (Alpha)" not in command:
|
| 29 |
+
continue
|
| 30 |
+
if "--type=renderer" in command and "Helper (Renderer)" in command:
|
| 31 |
+
result["renderer"][int(pid_text)] = int(rss_text) * 1024
|
| 32 |
+
elif "--type=gpu-process" in command and "user-data-dir=/Users/admin/Library/Application Support/t3code" in command:
|
| 33 |
+
result["gpu"][int(pid_text)] = int(rss_text) * 1024
|
| 34 |
+
return result
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class ProcessSampler:
|
| 38 |
+
def __init__(self, run_id: str) -> None:
|
| 39 |
+
self.run_id = run_id
|
| 40 |
+
self.started_at = time.time()
|
| 41 |
+
self.phase = "startup"
|
| 42 |
+
self.samples: list[dict] = []
|
| 43 |
+
self.stop_event = threading.Event()
|
| 44 |
+
self.thread = threading.Thread(target=self._run, daemon=True)
|
| 45 |
+
self.thread.start()
|
| 46 |
+
|
| 47 |
+
def _run(self) -> None:
|
| 48 |
+
while not self.stop_event.is_set():
|
| 49 |
+
try:
|
| 50 |
+
self.samples.append(
|
| 51 |
+
{"time": time.time(), "phase": self.phase, **process_rss()}
|
| 52 |
+
)
|
| 53 |
+
except Exception as error: # diagnostic sampling must not break OCR
|
| 54 |
+
self.samples.append({"time": time.time(), "error": str(error)})
|
| 55 |
+
self.stop_event.wait(0.05)
|
| 56 |
+
|
| 57 |
+
def mark(self, phase: str) -> None:
|
| 58 |
+
self.phase = phase
|
| 59 |
+
|
| 60 |
+
def finish(self) -> dict:
|
| 61 |
+
self.stop_event.set()
|
| 62 |
+
self.thread.join(timeout=2)
|
| 63 |
+
usable = [sample for sample in self.samples if "renderer" in sample]
|
| 64 |
+
if not usable:
|
| 65 |
+
return {"samples": len(self.samples), "error": "no usable process samples"}
|
| 66 |
+
baseline = usable[0]
|
| 67 |
+
|
| 68 |
+
def summarize(kind: str) -> dict:
|
| 69 |
+
baseline_by_pid = baseline[kind]
|
| 70 |
+
peak_by_pid: dict[int, int] = {}
|
| 71 |
+
first_by_pid: dict[int, int] = dict(baseline_by_pid)
|
| 72 |
+
for sample in usable:
|
| 73 |
+
for pid, rss in sample[kind].items():
|
| 74 |
+
first_by_pid.setdefault(pid, rss)
|
| 75 |
+
peak_by_pid[pid] = max(peak_by_pid.get(pid, 0), rss)
|
| 76 |
+
deltas = {
|
| 77 |
+
pid: peak - first_by_pid.get(pid, peak)
|
| 78 |
+
for pid, peak in peak_by_pid.items()
|
| 79 |
+
}
|
| 80 |
+
selected_pid = max(deltas, key=deltas.get) if deltas else None
|
| 81 |
+
phases: dict[str, dict[str, int]] = {}
|
| 82 |
+
if selected_pid:
|
| 83 |
+
for sample in usable:
|
| 84 |
+
rss = sample[kind].get(selected_pid)
|
| 85 |
+
if rss is None:
|
| 86 |
+
continue
|
| 87 |
+
phase = sample.get("phase", "unknown")
|
| 88 |
+
phase_summary = phases.setdefault(
|
| 89 |
+
phase,
|
| 90 |
+
{"startBytes": rss, "endBytes": rss, "peakBytes": rss},
|
| 91 |
+
)
|
| 92 |
+
phase_summary["endBytes"] = rss
|
| 93 |
+
phase_summary["peakBytes"] = max(phase_summary["peakBytes"], rss)
|
| 94 |
+
for phase_summary in phases.values():
|
| 95 |
+
phase_summary["peakDeltaBytes"] = (
|
| 96 |
+
phase_summary["peakBytes"] - phase_summary["startBytes"]
|
| 97 |
+
)
|
| 98 |
+
return {
|
| 99 |
+
"selectedPid": selected_pid,
|
| 100 |
+
"baselineBytes": first_by_pid.get(selected_pid, 0) if selected_pid else 0,
|
| 101 |
+
"peakBytes": peak_by_pid.get(selected_pid, 0) if selected_pid else 0,
|
| 102 |
+
"peakDeltaBytes": deltas.get(selected_pid, 0) if selected_pid else 0,
|
| 103 |
+
"allPeakDeltas": deltas,
|
| 104 |
+
"phases": phases,
|
| 105 |
+
}
|
| 106 |
+
|
| 107 |
+
return {
|
| 108 |
+
"scope": "T3 renderer with largest RSS delta plus shared T3 GPU process",
|
| 109 |
+
"samples": len(usable),
|
| 110 |
+
"durationSeconds": time.time() - self.started_at,
|
| 111 |
+
"renderer": summarize("renderer"),
|
| 112 |
+
"gpu": summarize("gpu"),
|
| 113 |
+
}
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
active_sampler: ProcessSampler | None = None
|
| 117 |
+
sampler_lock = threading.Lock()
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
class Handler(SimpleHTTPRequestHandler):
|
| 121 |
+
def do_POST(self) -> None:
|
| 122 |
+
global active_sampler
|
| 123 |
+
length = int(self.headers.get("Content-Length", "0"))
|
| 124 |
+
if length <= 0 or length > 2_000_000:
|
| 125 |
+
self.send_error(400, "invalid request size")
|
| 126 |
+
return
|
| 127 |
+
payload = json.loads(self.rfile.read(length))
|
| 128 |
+
if self.path == "/benchmark-start":
|
| 129 |
+
run_id = str(payload.get("runId", ""))
|
| 130 |
+
if not run_id or "/" in run_id or ".." in run_id:
|
| 131 |
+
self.send_error(400, "invalid run id")
|
| 132 |
+
return
|
| 133 |
+
with sampler_lock:
|
| 134 |
+
if active_sampler:
|
| 135 |
+
active_sampler.finish()
|
| 136 |
+
active_sampler = ProcessSampler(run_id)
|
| 137 |
+
self.send_response(204)
|
| 138 |
+
self.end_headers()
|
| 139 |
+
return
|
| 140 |
+
if self.path == "/benchmark-mark":
|
| 141 |
+
run_id = str(payload.get("runId", ""))
|
| 142 |
+
phase = str(payload.get("phase", ""))
|
| 143 |
+
if not phase or len(phase) > 80:
|
| 144 |
+
self.send_error(400, "invalid benchmark phase")
|
| 145 |
+
return
|
| 146 |
+
with sampler_lock:
|
| 147 |
+
sampler = active_sampler
|
| 148 |
+
if not sampler or sampler.run_id != run_id:
|
| 149 |
+
self.send_error(409, "benchmark sampler mismatch")
|
| 150 |
+
return
|
| 151 |
+
sampler.mark(phase)
|
| 152 |
+
self.send_response(204)
|
| 153 |
+
self.end_headers()
|
| 154 |
+
return
|
| 155 |
+
if self.path == "/benchmark-result":
|
| 156 |
+
run_id = str(payload.get("runId", ""))
|
| 157 |
+
with sampler_lock:
|
| 158 |
+
sampler = active_sampler
|
| 159 |
+
active_sampler = None
|
| 160 |
+
if not sampler or sampler.run_id != run_id:
|
| 161 |
+
self.send_error(409, "benchmark sampler mismatch")
|
| 162 |
+
return
|
| 163 |
+
payload["processMemory"] = sampler.finish()
|
| 164 |
+
RESULTS.mkdir(parents=True, exist_ok=True)
|
| 165 |
+
destination = RESULTS / f"{run_id}.json"
|
| 166 |
+
destination.write_text(
|
| 167 |
+
json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
|
| 168 |
+
encoding="utf-8",
|
| 169 |
+
)
|
| 170 |
+
response = json.dumps(payload["processMemory"]).encode()
|
| 171 |
+
self.send_response(200)
|
| 172 |
+
self.send_header("Content-Type", "application/json")
|
| 173 |
+
self.send_header("Content-Length", str(len(response)))
|
| 174 |
+
self.end_headers()
|
| 175 |
+
self.wfile.write(response)
|
| 176 |
+
return
|
| 177 |
+
self.send_error(404)
|
| 178 |
+
|
| 179 |
+
def end_headers(self) -> None:
|
| 180 |
+
self.send_header("Cross-Origin-Opener-Policy", "same-origin")
|
| 181 |
+
self.send_header("Cross-Origin-Embedder-Policy", "require-corp")
|
| 182 |
+
self.send_header("Cross-Origin-Resource-Policy", "same-origin")
|
| 183 |
+
self.send_header("Cache-Control", "no-store")
|
| 184 |
+
super().end_headers()
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
if __name__ == "__main__":
|
| 188 |
+
ThreadingHTTPServer(("127.0.0.1", 8765), Handler).serve_forever()
|
summarize_memory_results.py
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import statistics
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
ROOT = Path(__file__).resolve().parent
|
| 9 |
+
RESULTS = ROOT / ".work" / "results"
|
| 10 |
+
MIB = 1024 * 1024
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def is_worker_staged(result: dict) -> bool:
|
| 14 |
+
return (
|
| 15 |
+
# Model-execution A/B runs use their own phase-level report. Keep this
|
| 16 |
+
# runtime-only summary on the original fixed cohort without a decoder
|
| 17 |
+
# query so later candidate measurements cannot silently change it.
|
| 18 |
+
result.get("decoderMode") is None
|
| 19 |
+
and result.get("benchmarkScope") in (None, "full")
|
| 20 |
+
and any(
|
| 21 |
+
sample.get("label") == "vision-worker-session-ready"
|
| 22 |
+
for sample in result.get("memorySamples", [])
|
| 23 |
+
)
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def is_resident_control(result: dict) -> bool:
|
| 28 |
+
return (
|
| 29 |
+
result.get("memoryMode") == "default"
|
| 30 |
+
and result.get("loaderMode") == "bytes"
|
| 31 |
+
and not is_worker_staged(result)
|
| 32 |
+
and result.get("scheduleMode", "resident") == "resident"
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def median(rows: list[dict], value) -> float:
|
| 37 |
+
return statistics.median(value(row) for row in rows)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def main() -> None:
|
| 41 |
+
loaded = [json.loads(path.read_text(encoding="utf-8")) for path in RESULTS.glob("*.json")]
|
| 42 |
+
for tier in ("121", "242"):
|
| 43 |
+
controls = [row for row in loaded if row.get("tier") == tier and is_resident_control(row)]
|
| 44 |
+
staged = [row for row in loaded if row.get("tier") == tier and is_worker_staged(row)]
|
| 45 |
+
if not controls or not staged:
|
| 46 |
+
print(f"tier {tier}: missing control or worker-staged results")
|
| 47 |
+
continue
|
| 48 |
+
control_outputs = [item["actual"] for item in controls[0]["results"]]
|
| 49 |
+
if any([item["actual"] for item in row["results"]] != control_outputs for row in staged):
|
| 50 |
+
raise RuntimeError(f"tier {tier}: optimized OCR output differs from control")
|
| 51 |
+
|
| 52 |
+
control_renderer = median(
|
| 53 |
+
controls, lambda row: row["processMemory"]["renderer"]["peakDeltaBytes"] / MIB
|
| 54 |
+
)
|
| 55 |
+
staged_renderer = median(
|
| 56 |
+
staged, lambda row: row["processMemory"]["renderer"]["peakDeltaBytes"] / MIB
|
| 57 |
+
)
|
| 58 |
+
control_heap = median(
|
| 59 |
+
controls,
|
| 60 |
+
lambda row: max(sample.get("usedJSHeapSize", 0) for sample in row["memorySamples"]) / MIB,
|
| 61 |
+
)
|
| 62 |
+
staged_heap = median(
|
| 63 |
+
staged,
|
| 64 |
+
lambda row: max(sample.get("usedJSHeapSize", 0) for sample in row["memorySamples"]) / MIB,
|
| 65 |
+
)
|
| 66 |
+
print(
|
| 67 |
+
f"tier {tier}: parity=PASS controls={len(controls)} staged={len(staged)} "
|
| 68 |
+
f"renderer={control_renderer:.1f}->{staged_renderer:.1f} MiB "
|
| 69 |
+
f"({(staged_renderer / control_renderer - 1) * 100:+.1f}%), "
|
| 70 |
+
f"js_heap={control_heap:.1f}->{staged_heap:.1f} MiB "
|
| 71 |
+
f"({(staged_heap / control_heap - 1) * 100:+.1f}%)"
|
| 72 |
+
)
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
if __name__ == "__main__":
|
| 76 |
+
main()
|
variants/webgpu-121/decoder_unified_gather_qdq_int8.onnx
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:53778e5b50d78ee55c8674d3d19f53de14da808bd1c231a9d4cebd5d44d3ad93
|
| 3 |
+
size 34647902
|
variants/webgpu-242/decoder_unified_gather_qdq_int8.onnx
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:53778e5b50d78ee55c8674d3d19f53de14da808bd1c231a9d4cebd5d44d3ad93
|
| 3 |
+
size 34647902
|
vision-worker.mjs
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import * as ort from "./.work/ort/ort-webgpu.mjs";
|
| 2 |
+
|
| 3 |
+
let session = null;
|
| 4 |
+
let device = null;
|
| 5 |
+
|
| 6 |
+
const reply = (requestId, payload, transfer = []) => postMessage({ requestId, payload }, transfer);
|
| 7 |
+
const fail = (requestId, error) => postMessage({ requestId, error: error?.stack ?? String(error) });
|
| 8 |
+
|
| 9 |
+
self.onmessage = async ({ data }) => {
|
| 10 |
+
const { requestId, type, payload } = data;
|
| 11 |
+
try {
|
| 12 |
+
if (type === "init") {
|
| 13 |
+
ort.env.wasm.numThreads = 1;
|
| 14 |
+
ort.env.wasm.wasmPaths = payload.wasmPaths;
|
| 15 |
+
const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
|
| 16 |
+
if (!adapter) throw new Error("No WebGPU adapter is available in the vision worker");
|
| 17 |
+
device = await adapter.requestDevice();
|
| 18 |
+
ort.env.webgpu.device = device;
|
| 19 |
+
const options = {
|
| 20 |
+
executionProviders: ["webgpu"],
|
| 21 |
+
preferredOutputLocation: { vision_embeds: "cpu" },
|
| 22 |
+
logSeverityLevel: 3,
|
| 23 |
+
...(payload.lean ? { enableCpuMemArena: false, enableMemPattern: false } : {}),
|
| 24 |
+
};
|
| 25 |
+
session = await ort.InferenceSession.create(payload.modelPath, options);
|
| 26 |
+
if (JSON.stringify(session.inputNames) !== JSON.stringify(["pixel_values"])) {
|
| 27 |
+
throw new Error(`Unexpected vision inputs ${session.inputNames}`);
|
| 28 |
+
}
|
| 29 |
+
reply(requestId, { ready: true });
|
| 30 |
+
return;
|
| 31 |
+
}
|
| 32 |
+
if (type === "encode") {
|
| 33 |
+
if (!session) throw new Error("Vision worker is not initialized");
|
| 34 |
+
const pixels = new ort.Tensor("float32", payload.pixels, [1, 3, 224, 224]);
|
| 35 |
+
const started = performance.now();
|
| 36 |
+
const result = await session.run({ pixel_values: pixels });
|
| 37 |
+
const visionMs = performance.now() - started;
|
| 38 |
+
pixels.dispose();
|
| 39 |
+
const values = new Float32Array(result.vision_embeds.data);
|
| 40 |
+
const dims = [...result.vision_embeds.dims];
|
| 41 |
+
result.vision_embeds.dispose();
|
| 42 |
+
reply(requestId, { values, dims, visionMs }, [values.buffer]);
|
| 43 |
+
return;
|
| 44 |
+
}
|
| 45 |
+
if (type === "close") {
|
| 46 |
+
if (session) await session.release();
|
| 47 |
+
session = null;
|
| 48 |
+
if (device) device.destroy();
|
| 49 |
+
device = null;
|
| 50 |
+
reply(requestId, { closed: true });
|
| 51 |
+
return;
|
| 52 |
+
}
|
| 53 |
+
throw new Error(`Unknown worker operation ${type}`);
|
| 54 |
+
} catch (error) {
|
| 55 |
+
fail(requestId, error);
|
| 56 |
+
}
|
| 57 |
+
};
|
webgpu-memory-opt.html
ADDED
|
@@ -0,0 +1,446 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!doctype html>
|
| 2 |
+
<meta charset="utf-8" />
|
| 3 |
+
<title>Baberu capability-preserving WebGPU memory experiment</title>
|
| 4 |
+
<style>
|
| 5 |
+
body { background: #111; color: #eee; font: 14px/1.45 ui-monospace, monospace; margin: 24px; }
|
| 6 |
+
pre { white-space: pre-wrap; }
|
| 7 |
+
</style>
|
| 8 |
+
<h1>Baberu capability-preserving WebGPU memory experiment</h1>
|
| 9 |
+
<pre id="log">Starting…</pre>
|
| 10 |
+
<canvas id="canvas" width="224" height="224" hidden></canvas>
|
| 11 |
+
<script type="module">
|
| 12 |
+
const parameters = new URLSearchParams(location.search);
|
| 13 |
+
const tier = parameters.get("tier") ?? "121";
|
| 14 |
+
const memoryMode = parameters.get("memory") ?? "lean";
|
| 15 |
+
const loaderMode = parameters.get("loader") ?? "url";
|
| 16 |
+
const scheduleMode = parameters.get("schedule") ?? "staged";
|
| 17 |
+
const decoderMode = parameters.get("decoder") ?? "baseline";
|
| 18 |
+
const benchmarkScope = parameters.get("scope") ?? "full";
|
| 19 |
+
const decoderOnly = benchmarkScope === "decoder";
|
| 20 |
+
const maxTokens = 128;
|
| 21 |
+
if (!new Set(["121", "242"]).has(tier)) throw new Error(`Unknown tier ${tier}`);
|
| 22 |
+
if (!new Set(["default", "lean"]).has(memoryMode)) throw new Error(`Unknown memory mode ${memoryMode}`);
|
| 23 |
+
if (!new Set(["url", "bytes"]).has(loaderMode)) throw new Error(`Unknown loader ${loaderMode}`);
|
| 24 |
+
if (!new Set(["resident", "staged"]).has(scheduleMode)) throw new Error(`Unknown schedule ${scheduleMode}`);
|
| 25 |
+
if (!new Set(["baseline", "gather-opt", "fp16-matmul", "fixed-kv"]).has(decoderMode)) throw new Error(`Unknown decoder ${decoderMode}`);
|
| 26 |
+
if (!new Set(["full", "decoder"]).has(benchmarkScope)) throw new Error(`Unknown scope ${benchmarkScope}`);
|
| 27 |
+
const ort = await import("./.work/ort/ort-webgpu.mjs");
|
| 28 |
+
const output = document.querySelector("#log");
|
| 29 |
+
const canvas = document.querySelector("#canvas");
|
| 30 |
+
const context = canvas.getContext("2d", { willReadFrequently: true });
|
| 31 |
+
const logLines = [];
|
| 32 |
+
const memorySamples = [];
|
| 33 |
+
const log = (message) => {
|
| 34 |
+
logLines.push(`${new Date().toISOString()} ${message}`);
|
| 35 |
+
output.textContent = logLines.join("\n");
|
| 36 |
+
console.log(message);
|
| 37 |
+
};
|
| 38 |
+
window.addEventListener("error", (event) => log(`FAIL ${event.error?.stack ?? event.message}`));
|
| 39 |
+
window.addEventListener("unhandledrejection", (event) => log(`FAIL ${event.reason?.stack ?? event.reason}`));
|
| 40 |
+
const sampleMemory = async (label) => {
|
| 41 |
+
const sample = { label, atMs: performance.now() };
|
| 42 |
+
if (performance.memory) {
|
| 43 |
+
sample.jsHeapSizeLimit = performance.memory.jsHeapSizeLimit;
|
| 44 |
+
sample.totalJSHeapSize = performance.memory.totalJSHeapSize;
|
| 45 |
+
sample.usedJSHeapSize = performance.memory.usedJSHeapSize;
|
| 46 |
+
}
|
| 47 |
+
if (typeof performance.measureUserAgentSpecificMemory === "function") {
|
| 48 |
+
try {
|
| 49 |
+
sample.userAgentBytes = (await performance.measureUserAgentSpecificMemory()).bytes;
|
| 50 |
+
} catch (error) {
|
| 51 |
+
sample.userAgentMemoryError = String(error);
|
| 52 |
+
}
|
| 53 |
+
}
|
| 54 |
+
memorySamples.push(sample);
|
| 55 |
+
log(`MEMORY ${JSON.stringify(sample)}`);
|
| 56 |
+
return sample;
|
| 57 |
+
};
|
| 58 |
+
const runId = `tier-${tier}-${memoryMode}-${loaderMode}-${scheduleMode}-${decoderMode}-${benchmarkScope}-${Date.now()}`;
|
| 59 |
+
await fetch("/benchmark-start", {
|
| 60 |
+
method: "POST",
|
| 61 |
+
headers: { "Content-Type": "application/json" },
|
| 62 |
+
body: JSON.stringify({ runId }),
|
| 63 |
+
});
|
| 64 |
+
const markPhase = async (phase) => {
|
| 65 |
+
const response = await fetch("/benchmark-mark", {
|
| 66 |
+
method: "POST",
|
| 67 |
+
headers: { "Content-Type": "application/json" },
|
| 68 |
+
body: JSON.stringify({ runId, phase }),
|
| 69 |
+
});
|
| 70 |
+
if (!response.ok) throw new Error(`Failed marking ${phase}: HTTP ${response.status}`);
|
| 71 |
+
};
|
| 72 |
+
await markPhase("baseline");
|
| 73 |
+
await sampleMemory("baseline");
|
| 74 |
+
|
| 75 |
+
ort.env.wasm.numThreads = 1;
|
| 76 |
+
ort.env.wasm.wasmPaths = "/.work/ort/";
|
| 77 |
+
if (!navigator.gpu) throw new Error("WebGPU is unavailable");
|
| 78 |
+
let gpuDevice = null;
|
| 79 |
+
const ensureMainDevice = async () => {
|
| 80 |
+
if (gpuDevice) return;
|
| 81 |
+
const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
|
| 82 |
+
if (!adapter) throw new Error("No WebGPU adapter is available");
|
| 83 |
+
gpuDevice = await adapter.requestDevice();
|
| 84 |
+
ort.env.webgpu.device = gpuDevice;
|
| 85 |
+
};
|
| 86 |
+
const cacheNames = [
|
| 87 |
+
...Array.from({ length: 6 }, (_, index) => `present_k${index}`),
|
| 88 |
+
...Array.from({ length: 6 }, (_, index) => `present_v${index}`),
|
| 89 |
+
];
|
| 90 |
+
const createSession = async (path, preferredOutputLocation) => {
|
| 91 |
+
await ensureMainDevice();
|
| 92 |
+
const options = {
|
| 93 |
+
executionProviders: ["webgpu"],
|
| 94 |
+
preferredOutputLocation,
|
| 95 |
+
logSeverityLevel: 3,
|
| 96 |
+
...(memoryMode === "lean" ? { enableCpuMemArena: false, enableMemPattern: false } : {}),
|
| 97 |
+
};
|
| 98 |
+
const started = performance.now();
|
| 99 |
+
const session = loaderMode === "url"
|
| 100 |
+
? await ort.InferenceSession.create(path, options)
|
| 101 |
+
: await fetch(path).then(async (response) => {
|
| 102 |
+
if (!response.ok) throw new Error(`${path}: HTTP ${response.status}`);
|
| 103 |
+
return ort.InferenceSession.create(new Uint8Array(await response.arrayBuffer()), options);
|
| 104 |
+
});
|
| 105 |
+
log(`SESSION ${path} ${(performance.now() - started).toFixed(1)} ms`);
|
| 106 |
+
return session;
|
| 107 |
+
};
|
| 108 |
+
const visionPath = tier === "121"
|
| 109 |
+
? "./.work/models/webgpu-121/vision_int4.onnx"
|
| 110 |
+
: "./.work/models/webgpu-242/vision_fp16.onnx";
|
| 111 |
+
const fixedKv = decoderMode === "fixed-kv";
|
| 112 |
+
const decoderPath = ({
|
| 113 |
+
baseline: "./.work/models/shared/decoder_unified_gather_qdq_int8.onnx",
|
| 114 |
+
"gather-opt": "./.work/models/model-opt/decoder_gather_before_dq_int8.onnx",
|
| 115 |
+
"fp16-matmul": "./.work/models/model-opt/decoder_static_fp16_matmul.onnx",
|
| 116 |
+
"fixed-kv": "./.work/models/model-opt/decoder_gather_dq_fixed_kv_int8.onnx",
|
| 117 |
+
})[decoderMode];
|
| 118 |
+
let vision = null;
|
| 119 |
+
let decoder = null;
|
| 120 |
+
const createVision = async () => {
|
| 121 |
+
vision = await createSession(visionPath, scheduleMode === "staged" ? { vision_embeds: "cpu" } : { vision_embeds: "gpu-buffer" });
|
| 122 |
+
if (JSON.stringify(vision.inputNames) !== JSON.stringify(["pixel_values"])) {
|
| 123 |
+
throw new Error(`Unexpected vision inputs ${vision.inputNames}`);
|
| 124 |
+
}
|
| 125 |
+
await sampleMemory("vision-session-ready");
|
| 126 |
+
};
|
| 127 |
+
const createDecoder = async () => {
|
| 128 |
+
decoder = await createSession(decoderPath, Object.fromEntries([
|
| 129 |
+
["logits", "cpu"],
|
| 130 |
+
...cacheNames.map((name) => [name, "gpu-buffer"]),
|
| 131 |
+
]));
|
| 132 |
+
for (const required of ["vision_embeds", "token_ids", "position_ids", "past_k0", "past_v5"]) {
|
| 133 |
+
if (!decoder.inputNames.includes(required)) throw new Error(`Decoder input missing ${required}`);
|
| 134 |
+
}
|
| 135 |
+
if (fixedKv !== decoder.inputNames.includes("past_length")) {
|
| 136 |
+
throw new Error(`Decoder past_length contract mismatch for ${decoderMode}`);
|
| 137 |
+
}
|
| 138 |
+
if (!decoder.outputNames.includes("logits") || !decoder.outputNames.includes("present_v5")) {
|
| 139 |
+
throw new Error("Decoder capability outputs are incomplete");
|
| 140 |
+
}
|
| 141 |
+
await sampleMemory("decoder-session-ready");
|
| 142 |
+
};
|
| 143 |
+
if (scheduleMode === "resident" && !decoderOnly) {
|
| 144 |
+
await markPhase("vision-load");
|
| 145 |
+
await createVision();
|
| 146 |
+
await markPhase("decoder-load");
|
| 147 |
+
await createDecoder();
|
| 148 |
+
}
|
| 149 |
+
if (decoderOnly) {
|
| 150 |
+
await markPhase("decoder-load");
|
| 151 |
+
await createDecoder();
|
| 152 |
+
}
|
| 153 |
+
|
| 154 |
+
const charset = await (await fetch("./tokenizer/vocab.json")).json();
|
| 155 |
+
const idToCharacter = ["", "", "", "", ...charset];
|
| 156 |
+
if (idToCharacter.length !== 14630) throw new Error(`Unexpected vocabulary ${idToCharacter.length}`);
|
| 157 |
+
const contentIds = new Set();
|
| 158 |
+
for (let index = 0; index < charset.length; index += 1) {
|
| 159 |
+
const character = charset[index];
|
| 160 |
+
if (character.length === 1 && !"ーー〜~".includes(character) && /[\p{Letter}\p{Number}]/u.test(character)) {
|
| 161 |
+
contentIds.add(index + 4);
|
| 162 |
+
}
|
| 163 |
+
}
|
| 164 |
+
const loadImage = async (path) => createImageBitmap(await (await fetch(path)).blob());
|
| 165 |
+
const preprocess = (image) => {
|
| 166 |
+
context.imageSmoothingEnabled = true;
|
| 167 |
+
context.imageSmoothingQuality = "high";
|
| 168 |
+
context.clearRect(0, 0, 224, 224);
|
| 169 |
+
context.drawImage(image, 0, 0, image.width, image.height, 0, 0, 224, 224);
|
| 170 |
+
const rgba = context.getImageData(0, 0, 224, 224).data;
|
| 171 |
+
const values = new Float32Array(3 * 224 * 224);
|
| 172 |
+
const mean = [0.485, 0.456, 0.406];
|
| 173 |
+
const std = [0.229, 0.224, 0.225];
|
| 174 |
+
for (let pixel = 0; pixel < 224 * 224; pixel += 1) {
|
| 175 |
+
for (let channel = 0; channel < 3; channel += 1) {
|
| 176 |
+
values[channel * 224 * 224 + pixel] = (rgba[pixel * 4 + channel] / 255 - mean[channel]) / std[channel];
|
| 177 |
+
}
|
| 178 |
+
}
|
| 179 |
+
return values;
|
| 180 |
+
};
|
| 181 |
+
const tokenId = (token) => new ort.Tensor("int32", Int32Array.of(token), [1, 1]);
|
| 182 |
+
const emptyVision = () => new ort.Tensor("float32", new Float32Array(0), [1, 0, 512]);
|
| 183 |
+
const emptyCache = () => fixedKv
|
| 184 |
+
? new ort.Tensor("float32", new Float32Array(1 * 2 * 384 * 64), [1, 2, 384, 64])
|
| 185 |
+
: new ort.Tensor("float32", new Float32Array(0), [1, 2, 0, 64]);
|
| 186 |
+
const pastLength = (length) => new ort.Tensor("int64", BigInt64Array.of(BigInt(length)), [1]);
|
| 187 |
+
const positions = () => new ort.Tensor("int32", Int32Array.from({ length: 257 }, (_, index) => index), [1, 257]);
|
| 188 |
+
const disposeResult = (result) => {
|
| 189 |
+
result.logits?.dispose();
|
| 190 |
+
for (const name of cacheNames) result[name]?.dispose();
|
| 191 |
+
};
|
| 192 |
+
const chooseToken = (source, sequence, tokens) => {
|
| 193 |
+
const seen = new Set(sequence);
|
| 194 |
+
let blocked = -1;
|
| 195 |
+
const last = tokens.at(-1);
|
| 196 |
+
if (last !== undefined && contentIds.has(last)) {
|
| 197 |
+
let run = 0;
|
| 198 |
+
for (let index = tokens.length - 1; index >= 0 && tokens[index] === last; index -= 1) run += 1;
|
| 199 |
+
if (run >= 12) blocked = last;
|
| 200 |
+
}
|
| 201 |
+
const adjusted = (index) => {
|
| 202 |
+
if (index === blocked) return Number.NEGATIVE_INFINITY;
|
| 203 |
+
const value = source[index];
|
| 204 |
+
if (!seen.has(index)) return value;
|
| 205 |
+
return value < 0 ? value * 1.2 : value / 1.2;
|
| 206 |
+
};
|
| 207 |
+
let bestToken = 0;
|
| 208 |
+
let bestValue = adjusted(0);
|
| 209 |
+
for (let token = 1; token < 14630; token += 1) {
|
| 210 |
+
const value = adjusted(token);
|
| 211 |
+
if (value > bestValue) {
|
| 212 |
+
bestToken = token;
|
| 213 |
+
bestValue = value;
|
| 214 |
+
}
|
| 215 |
+
}
|
| 216 |
+
return bestToken;
|
| 217 |
+
};
|
| 218 |
+
const encode = async (image) => {
|
| 219 |
+
const pixels = new ort.Tensor("float32", preprocess(image), [1, 3, 224, 224]);
|
| 220 |
+
const visionStarted = performance.now();
|
| 221 |
+
const visionResult = await vision.run({ pixel_values: pixels });
|
| 222 |
+
const visionMs = performance.now() - visionStarted;
|
| 223 |
+
pixels.dispose();
|
| 224 |
+
return { visionEmbeds: visionResult.vision_embeds, visionMs };
|
| 225 |
+
};
|
| 226 |
+
const decode = async (visionEmbeds, visionMs) => {
|
| 227 |
+
const started = performance.now();
|
| 228 |
+
const prefillFeeds = {
|
| 229 |
+
vision_embeds: visionEmbeds,
|
| 230 |
+
token_ids: tokenId(1),
|
| 231 |
+
position_ids: positions(),
|
| 232 |
+
...Object.fromEntries(cacheNames.map((name) => [name.replace("present_", "past_"), emptyCache()])),
|
| 233 |
+
};
|
| 234 |
+
if (fixedKv) prefillFeeds.past_length = pastLength(0);
|
| 235 |
+
const prefillStarted = performance.now();
|
| 236 |
+
let cacheResult = await decoder.run(prefillFeeds);
|
| 237 |
+
const prefillMs = performance.now() - prefillStarted;
|
| 238 |
+
visionEmbeds.dispose();
|
| 239 |
+
for (const [name, tensor] of Object.entries(prefillFeeds)) if (name !== "vision_embeds") tensor.dispose();
|
| 240 |
+
const sequence = [1];
|
| 241 |
+
const tokens = [];
|
| 242 |
+
const decodeStarted = performance.now();
|
| 243 |
+
for (let iteration = 0; iteration < maxTokens; iteration += 1) {
|
| 244 |
+
// Decoder-only memory runs force a stable non-EOS token so every graph
|
| 245 |
+
// executes the complete 128-token cache path.
|
| 246 |
+
const next = decoderOnly ? 4 : chooseToken(cacheResult.logits.data, sequence, tokens);
|
| 247 |
+
if (next === 2) break;
|
| 248 |
+
tokens.push(next);
|
| 249 |
+
sequence.push(next);
|
| 250 |
+
if (tokens.length >= maxTokens) break;
|
| 251 |
+
const feeds = {
|
| 252 |
+
vision_embeds: emptyVision(),
|
| 253 |
+
token_ids: tokenId(next),
|
| 254 |
+
position_ids: new ort.Tensor("int32", Int32Array.of(257 + iteration), [1, 1]),
|
| 255 |
+
};
|
| 256 |
+
if (fixedKv) feeds.past_length = pastLength(257 + iteration);
|
| 257 |
+
for (let layer = 0; layer < 6; layer += 1) {
|
| 258 |
+
feeds[`past_k${layer}`] = cacheResult[`present_k${layer}`];
|
| 259 |
+
feeds[`past_v${layer}`] = cacheResult[`present_v${layer}`];
|
| 260 |
+
}
|
| 261 |
+
const nextResult = await decoder.run(feeds);
|
| 262 |
+
disposeResult(cacheResult);
|
| 263 |
+
feeds.vision_embeds.dispose();
|
| 264 |
+
feeds.token_ids.dispose();
|
| 265 |
+
feeds.position_ids.dispose();
|
| 266 |
+
feeds.past_length?.dispose();
|
| 267 |
+
cacheResult = nextResult;
|
| 268 |
+
}
|
| 269 |
+
const decodeMs = performance.now() - decodeStarted;
|
| 270 |
+
disposeResult(cacheResult);
|
| 271 |
+
return {
|
| 272 |
+
text: tokens.map((token) => idToCharacter[token] ?? "").join(""),
|
| 273 |
+
tokens: tokens.length,
|
| 274 |
+
visionMs,
|
| 275 |
+
prefillMs,
|
| 276 |
+
decodeMs,
|
| 277 |
+
totalMs: visionMs + performance.now() - started,
|
| 278 |
+
};
|
| 279 |
+
};
|
| 280 |
+
const recognize = async (image) => {
|
| 281 |
+
const encoded = await encode(image);
|
| 282 |
+
return decode(encoded.visionEmbeds, encoded.visionMs);
|
| 283 |
+
};
|
| 284 |
+
const expected = [
|
| 285 |
+
["01", "知らない世界で見つけたイメージを"],
|
| 286 |
+
["02", "カナデトモスソラ(Kanadetomosusora)"],
|
| 287 |
+
["03", "建設会社社員行方"],
|
| 288 |
+
["04", "だとしてもこのレベルがウロつくなんて...おそらく2級の呪い"],
|
| 289 |
+
["05", "パチパチパチパチ"],
|
| 290 |
+
["06", "バビュン"],
|
| 291 |
+
["07", "僕の過去とか未来とか"],
|
| 292 |
+
["08", "くらべられっ子"],
|
| 293 |
+
["09", "そうだクラス分けがあるんだった!!"],
|
| 294 |
+
["10", "脇役よ、主役を超えよ!"],
|
| 295 |
+
["11", "Eh~Idon'treallywantto~"],
|
| 296 |
+
["12", "「Sorryforthewait~!Didyouwaitlong?」"],
|
| 297 |
+
["13", "YamateAreaNewresidentialdistrictforforeigners"],
|
| 298 |
+
];
|
| 299 |
+
const normalize = (text) => text.normalize("NFKC").replace(/\s+/g, "");
|
| 300 |
+
const editDistance = (left, right) => {
|
| 301 |
+
let previous = Array.from({ length: right.length + 1 }, (_, index) => index);
|
| 302 |
+
for (let leftIndex = 1; leftIndex <= left.length; leftIndex += 1) {
|
| 303 |
+
const current = [leftIndex];
|
| 304 |
+
for (let rightIndex = 1; rightIndex <= right.length; rightIndex += 1) {
|
| 305 |
+
current.push(Math.min(
|
| 306 |
+
current[rightIndex - 1] + 1,
|
| 307 |
+
previous[rightIndex] + 1,
|
| 308 |
+
previous[rightIndex - 1] + (left[leftIndex - 1] === right[rightIndex - 1] ? 0 : 1),
|
| 309 |
+
));
|
| 310 |
+
}
|
| 311 |
+
previous = current;
|
| 312 |
+
}
|
| 313 |
+
return previous.at(-1);
|
| 314 |
+
};
|
| 315 |
+
const stagedInputs = [];
|
| 316 |
+
if (scheduleMode === "staged" && !decoderOnly) {
|
| 317 |
+
await markPhase("vision-load");
|
| 318 |
+
const visionWorker = new Worker("./vision-worker.mjs", { type: "module" });
|
| 319 |
+
let workerRequestId = 0;
|
| 320 |
+
const workerCalls = new Map();
|
| 321 |
+
visionWorker.onmessage = ({ data }) => {
|
| 322 |
+
const call = workerCalls.get(data.requestId);
|
| 323 |
+
if (!call) return;
|
| 324 |
+
workerCalls.delete(data.requestId);
|
| 325 |
+
if (data.error) call.reject(new Error(data.error));
|
| 326 |
+
else call.resolve(data.payload);
|
| 327 |
+
};
|
| 328 |
+
visionWorker.onerror = (event) => {
|
| 329 |
+
for (const call of workerCalls.values()) call.reject(event.error ?? new Error(event.message));
|
| 330 |
+
workerCalls.clear();
|
| 331 |
+
};
|
| 332 |
+
const callWorker = (type, payload, transfer = []) => new Promise((resolve, reject) => {
|
| 333 |
+
const requestId = ++workerRequestId;
|
| 334 |
+
workerCalls.set(requestId, { resolve, reject });
|
| 335 |
+
visionWorker.postMessage({ requestId, type, payload }, transfer);
|
| 336 |
+
});
|
| 337 |
+
await callWorker("init", {
|
| 338 |
+
modelPath: new URL(visionPath, location.href).href,
|
| 339 |
+
wasmPaths: new URL("./.work/ort/", location.href).href,
|
| 340 |
+
lean: memoryMode === "lean",
|
| 341 |
+
});
|
| 342 |
+
await sampleMemory("vision-worker-session-ready");
|
| 343 |
+
await markPhase("vision-encode");
|
| 344 |
+
for (const [id, expectedText] of expected) {
|
| 345 |
+
log(`ENCODE ${id}/13`);
|
| 346 |
+
const image = await loadImage(`./.work/hayai13/${id}.png`);
|
| 347 |
+
const pixels = preprocess(image);
|
| 348 |
+
image.close();
|
| 349 |
+
const encoded = await callWorker("encode", { pixels }, [pixels.buffer]);
|
| 350 |
+
stagedInputs.push({ id, expectedText, ...encoded });
|
| 351 |
+
}
|
| 352 |
+
await sampleMemory("all-vision-embeddings-ready");
|
| 353 |
+
await markPhase("vision-release");
|
| 354 |
+
await callWorker("close", {});
|
| 355 |
+
visionWorker.terminate();
|
| 356 |
+
await new Promise((resolve) => setTimeout(resolve, 100));
|
| 357 |
+
await sampleMemory("vision-worker-terminated");
|
| 358 |
+
await markPhase("decoder-load");
|
| 359 |
+
await createDecoder();
|
| 360 |
+
}
|
| 361 |
+
const benchmarkCases = decoderOnly ? [["synthetic", ""]] : expected;
|
| 362 |
+
const results = [];
|
| 363 |
+
await markPhase("decoder-run");
|
| 364 |
+
for (let caseIndex = 0; caseIndex < benchmarkCases.length; caseIndex += 1) {
|
| 365 |
+
const [id, expectedText] = benchmarkCases[caseIndex];
|
| 366 |
+
log(`CASE ${id}/${benchmarkCases.length}`);
|
| 367 |
+
let result;
|
| 368 |
+
if (decoderOnly) {
|
| 369 |
+
const visionEmbeds = new ort.Tensor("float32", new Float32Array(256 * 512), [1, 256, 512]);
|
| 370 |
+
result = await decode(visionEmbeds, 0);
|
| 371 |
+
} else if (scheduleMode === "staged") {
|
| 372 |
+
const staged = stagedInputs[caseIndex];
|
| 373 |
+
const visionEmbeds = new ort.Tensor("float32", staged.values, staged.dims);
|
| 374 |
+
result = await decode(visionEmbeds, staged.visionMs);
|
| 375 |
+
} else {
|
| 376 |
+
const image = await loadImage(`./.work/hayai13/${id}.png`);
|
| 377 |
+
result = await recognize(image);
|
| 378 |
+
image.close();
|
| 379 |
+
}
|
| 380 |
+
const normalizedExpected = normalize(expectedText);
|
| 381 |
+
const normalizedActual = normalize(result.text);
|
| 382 |
+
results.push({
|
| 383 |
+
id,
|
| 384 |
+
expected: expectedText,
|
| 385 |
+
actual: result.text,
|
| 386 |
+
distance: editDistance(normalizedExpected, normalizedActual),
|
| 387 |
+
characters: normalizedExpected.length,
|
| 388 |
+
exact: normalizedExpected === normalizedActual,
|
| 389 |
+
...result,
|
| 390 |
+
});
|
| 391 |
+
await sampleMemory(`after-case-${id}`);
|
| 392 |
+
}
|
| 393 |
+
const distance = results.reduce((sum, result) => sum + result.distance, 0);
|
| 394 |
+
const characters = results.reduce((sum, result) => sum + result.characters, 0);
|
| 395 |
+
const sortedTimes = results.map((result) => result.totalMs).sort((left, right) => left - right);
|
| 396 |
+
await markPhase("release");
|
| 397 |
+
if (vision) await vision.release();
|
| 398 |
+
if (decoder) await decoder.release();
|
| 399 |
+
if (gpuDevice) gpuDevice.destroy();
|
| 400 |
+
await new Promise((resolve) => setTimeout(resolve, 250));
|
| 401 |
+
await sampleMemory("released");
|
| 402 |
+
const summary = {
|
| 403 |
+
runId,
|
| 404 |
+
tier,
|
| 405 |
+
memoryMode,
|
| 406 |
+
loaderMode,
|
| 407 |
+
scheduleMode,
|
| 408 |
+
decoderMode,
|
| 409 |
+
benchmarkScope,
|
| 410 |
+
capability: {
|
| 411 |
+
layers: 6,
|
| 412 |
+
hiddenSize: 512,
|
| 413 |
+
intermediateSize: 1536,
|
| 414 |
+
attentionHeads: 8,
|
| 415 |
+
kvHeads: 2,
|
| 416 |
+
vocabulary: 14630,
|
| 417 |
+
maxTokens,
|
| 418 |
+
vision: tier === "121" ? "upstream INT4" : "upstream FP16",
|
| 419 |
+
decoder: ({
|
| 420 |
+
baseline: "complete INT8-QDQ unified Gather",
|
| 421 |
+
"gather-opt": "complete INT8-QDQ Gather-before-DQ",
|
| 422 |
+
"fp16-matmul": "layers 1-3 static FP16 MatMul with Gather-before-DQ embedding",
|
| 423 |
+
"fixed-kv": "complete INT8-QDQ Gather-before-DQ with fixed 384-slot KV I/O",
|
| 424 |
+
})[decoderMode],
|
| 425 |
+
},
|
| 426 |
+
cases: results.length,
|
| 427 |
+
nCER: characters ? distance / characters : null,
|
| 428 |
+
distance,
|
| 429 |
+
characters,
|
| 430 |
+
exact: results.filter((result) => result.exact).length,
|
| 431 |
+
medianMs: sortedTimes[Math.floor(sortedTimes.length / 2)],
|
| 432 |
+
peakCaseMs: sortedTimes.at(-1),
|
| 433 |
+
memorySamples,
|
| 434 |
+
results,
|
| 435 |
+
};
|
| 436 |
+
const response = await fetch("/benchmark-result", {
|
| 437 |
+
method: "POST",
|
| 438 |
+
headers: { "Content-Type": "application/json" },
|
| 439 |
+
body: JSON.stringify(summary),
|
| 440 |
+
});
|
| 441 |
+
if (!response.ok) throw new Error(`Failed persisting result: HTTP ${response.status}`);
|
| 442 |
+
summary.processMemory = await response.json();
|
| 443 |
+
window.__benchmarkResult = summary;
|
| 444 |
+
window.__benchmarkDone = true;
|
| 445 |
+
log(`SUITE_RESULT ${JSON.stringify(summary)}`);
|
| 446 |
+
</script>
|