ameraino11 commited on
Commit
cdf34e1
·
verified ·
1 Parent(s): 8020c75

Optimize Baberu WebGPU decoder memory execution

Browse files

Publish 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 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 and uses a quantized embedding
34
- `Gather` instead of a 14,630-element one-hot embedding MatMul.
 
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. It contains neither `DynamicQuantizeLinear` nor
 
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:3612d5787ffb1026ae09bfd4a16b3379ea75fe0dd803484f7bb41645be6b6239
3
- size 34647907
 
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:3612d5787ffb1026ae09bfd4a16b3379ea75fe0dd803484f7bb41645be6b6239
3
- size 34647907
 
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>