naklitechie commited on
Commit
cc08b9f
·
verified ·
1 Parent(s): 0bda4e5

diskformer.js c8d19ac: chat example, engines, store

Browse files
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Chirag Patnaik
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
README.md CHANGED
@@ -1,10 +1,23 @@
1
  ---
2
- title: Diskformer Chat
3
- emoji: 👁
4
- colorFrom: blue
5
- colorTo: purple
6
  sdk: static
 
7
  pinned: false
 
 
8
  ---
9
 
10
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: diskformer chat
3
+ colorFrom: gray
4
+ colorTo: blue
 
5
  sdk: static
6
+ app_file: examples/chat/index.html
7
  pinned: false
8
+ license: mit
9
+ short_description: A 14.4 GB model in 2.5 GB of GPU memory, in your tab
10
  ---
11
 
12
+ # diskformer chat
13
+
14
+ Gemma 4 26B-A4B (14.4 GB) or Qwen3.6 35B-A3B (36.9 GB) running in your browser tab with a GPU memory budget you
15
+ choose. The model downloads once into your browser's private storage (OPFS); the GPU holds the dense layers and a
16
+ small cache of experts, and the rest are read from disk as each token needs them.
17
+
18
+ Needs Chrome or Edge with WebGPU and shader-f16, and free disk space for the model. Open the app in its own tab,
19
+ so its storage is first-party: https://naklitechie-diskformer-chat.static.hf.space/examples/chat/
20
+
21
+ Source, measurements and tests: [github.com/NakliTechie/diskformer.js](https://github.com/NakliTechie/diskformer.js).
22
+ Model terms: Gemma ([ai.google.dev/gemma/terms](https://ai.google.dev/gemma/terms)); Qwen3.6, Apache-2.0. Engine
23
+ kernels are ported from gemma4-webgpu (Apache-2.0, see `engines/NOTICE`).
engines/LICENSE-APACHE ADDED
@@ -0,0 +1,201 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for describing the origin of the Work and
141
+ reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright 2026 Tyler Straub
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
engines/NOTICE ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ diskformer.js — engines/
2
+
3
+ The WGSL kernels in these engines are ports of browser-big-fast-lab's custom-kernels engine, an Apache-2.0
4
+ fork of gemma4-webgpu:
5
+
6
+ gemma4-webgpu
7
+ Copyright 2026 Tyler Straub
8
+ Licensed under the Apache License, Version 2.0 (engines/LICENSE-APACHE)
9
+ https://github.com/tylerstraub/gemma4-webgpu
10
+
11
+ Changes from the original: rewritten for mixture-of-experts models whose experts stream from OPFS (Q8_0 and Q4_0
12
+ split planes, slot-indexed expert GEMVs, GPU router top-k, chunked prefill, Q6_K head planes). Everything else in
13
+ this folder is MIT (../LICENSE), Copyright (c) 2026 Chirag Patnaik.
14
+
15
+ The engine files are copies of LocalMind's (https://github.com/NakliTechie/LocalMind), kept in step by
16
+ scripts/sync-localmind.mjs. Model weights are not distributed: each model is downloaded by the user under its own
17
+ terms (Gemma: https://ai.google.dev/gemma/terms; Qwen3.6: Apache-2.0).
engines/gemma4_moe_ssd.js ADDED
@@ -0,0 +1,1291 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* gemma4_moe_ssd.js — Gemma 4 26B-A4B (GGUF arch `gemma4`, MoE) with its routed experts on disk:
2
+ * rung 2c of the SSD-streaming ladder. Reuses rung 2a/2b's OPFS ingest, reader pool and expert streamer
3
+ * (qwen3_moe_ssd.js, moe-expert-stream.js); what differs is the file format and the trunk.
4
+ *
5
+ * The math follows llama.cpp b9830's src/models/gemma4.cpp (notes: LocalMind plan/gemma4-26b-a4b-port-notes.md;
6
+ * a numpy copy of it, scripts/gemma4-ref.py, gives per-layer states to compare against). Per layer:
7
+ * attention: RMS norm → Q/K/V (Q4_0) → per-head RMS norm on Q and K (weighted) and V (unweighted; full-attention
8
+ * layers have no V projection and reuse K's) → NeoX RoPE (sliding layers: 256-dim heads, base 1e4; full layers:
9
+ * 512-dim heads, base 1e6, `rope_freqs` leaves all but the first 64 pairs unrotated) → attention with scale 1.0,
10
+ * sliding layers over the last 1024 positions → output projection → post-norm → residual;
11
+ * FFN: a dense GELU MLP (2112) beside the MoE block (128 experts, top 8, GELU, a per-expert down scale), the router
12
+ * on rms_norm(x)·(1/√H)·scale; each branch post-normed, summed, normed again → residual → × layer_output_scale.
13
+ * Embedding: Q6_K rows × √H; head: the same Q6_K table, then a soft cap tanh(l/30)·30.
14
+ *
15
+ * File format (google/gemma-4-26B-A4B-it-qat-q4_0-gguf): matmul weights Q4_0, norms / router / scales F32,
16
+ * the token embedding (tied output head) Q6_K. Experts: fused `ffn_gate_up_exps` [H, 2F, E] (gate rows first)
17
+ * and `ffn_down_exps` [F, H, E], both Q4_0, plus `ffn_down_exps.scale` [E].
18
+ */
19
+
20
+ import { Qwen3MoeSsd, GGML, Q4_BLOCK, Q6K_BLOCK, tensorBytes, TOK, STORAGE, COPY_SRC, COPY_DST, MAP_READ } from './qwen3_moe_ssd.js';
21
+ import { OpfsReaderPool } from '../src/opfs-reader.js';
22
+ import { QWEN35_BATCH_KERNELS } from './qwen35_moe_ssd.js';
23
+
24
+ export function configFromGgufGemma4(kv) {
25
+ const a = kv['general.architecture'];
26
+ if (a !== 'gemma4') throw new Error(`expected a gemma4 GGUF, got ${a}`);
27
+ const g = (k) => kv[`${a}.${k}`];
28
+ const layers = g('block_count');
29
+ const perLayer = (v) => (Array.isArray(v) ? v : Array(layers).fill(v));
30
+ const swa = perLayer(g('attention.sliding_window_pattern')).map(Boolean);
31
+ if (!g('expert_count')) throw new Error('gemma4 GGUF without experts: this engine is for the MoE (26B-A4B)');
32
+ return {
33
+ arch: a, layers,
34
+ hidden: g('embedding_length'), denseFf: g('feed_forward_length'),
35
+ heads: g('attention.head_count'), kvHeads: perLayer(g('attention.head_count_kv')),
36
+ swa, // true: sliding-window layer
37
+ headDim: swa.map((s) => (s ? g('attention.key_length_swa') : g('attention.key_length'))),
38
+ ropeDims: swa.map((s) => (s ? g('rope.dimension_count_swa') : g('rope.dimension_count'))),
39
+ ropeTheta: swa.map((s) => (s ? g('rope.freq_base_swa') : g('rope.freq_base'))),
40
+ window: g('attention.sliding_window'),
41
+ experts: g('expert_count'), topK: g('expert_used_count'), expertFf: g('expert_feed_forward_length'),
42
+ softcap: g('final_logit_softcapping') || 0,
43
+ eps: g('attention.layer_norm_rms_epsilon'),
44
+ contextLength: g('context_length'),
45
+ vocab: kv['tokenizer.ggml.tokens'].length,
46
+ bos: kv['tokenizer.ggml.bos_token_id'], eos: kv['tokenizer.ggml.eos_token_id'],
47
+ };
48
+ }
49
+
50
+ // The OPFS layout from the header alone. One expert record = the fused gate/up planes then the down
51
+ // planes, each split into a nibble plane (16 B per 32 values) and an f16 scale plane (2 B per 32).
52
+ export function planLayoutGemma4(gguf) {
53
+ const cfg = configFromGgufGemma4(gguf.kv);
54
+ const H = cfg.hidden, F = cfg.expertFf;
55
+ const nib = (rows, cols) => rows * cols / 2, scl = (rows, cols) => rows * (cols / 32) * 2;
56
+ const parts = {
57
+ guQ: { off: 0, bytes: nib(2 * F, H) }, // gate rows 0..F-1, up rows F..2F-1
58
+ guS: { off: nib(2 * F, H), bytes: scl(2 * F, H) },
59
+ dQ: { off: nib(2 * F, H) + scl(2 * F, H), bytes: nib(H, F) },
60
+ dS: { off: nib(2 * F, H) + scl(2 * F, H) + nib(H, F), bytes: scl(H, F) },
61
+ };
62
+ const record = parts.dS.off + parts.dS.bytes;
63
+ const dense = {};
64
+ let off = 0;
65
+ const take = (n) => { const o = off; off += Math.ceil(n / 256) * 256; return o; };
66
+ for (const t of gguf.tensors) {
67
+ if (/_exps\.weight$/.test(t.name)) continue;
68
+ const n = t.dims.reduce((x, y) => x * y, 1);
69
+ if (t.type === GGML.Q4_0) {
70
+ dense[t.name] = { type: 'q4', dims: t.dims, q: { off: take(n / 2), bytes: n / 2 }, s: { off: take(n / 16), bytes: n / 16 } };
71
+ } else if (t.type === GGML.Q6_K) {
72
+ dense[t.name] = { type: 'q6k', dims: t.dims, raw: { off: take(tensorBytes(t)), bytes: tensorBytes(t) } };
73
+ } else if (t.type === GGML.F32) {
74
+ dense[t.name] = { type: 'f32', dims: t.dims, raw: { off: take(n * 4), bytes: n * 4 } };
75
+ } else throw new Error(`dense tensor ${t.name}: unsupported type ${t.type}`);
76
+ }
77
+ return {
78
+ config: cfg,
79
+ experts: { file: 'experts.bin', record, parts, layers: cfg.layers, perLayer: cfg.experts, bytes: record * cfg.layers * cfg.experts },
80
+ dense: { file: 'dense.bin', bytes: off, tensors: dense },
81
+ };
82
+ }
83
+
84
+ // Byte ranges of the GGUF to keep, each small enough to transform in memory: one unit per expert slice
85
+ // (kind q4), row bands of dense Q4_0 tensors (kind q4), whole Q6_K / F32 tensors (raw).
86
+ export function planUnitsGemma4(gguf, layout) {
87
+ const units = [];
88
+ const { record, parts, perLayer } = layout.experts;
89
+ for (const t of gguf.tensors) {
90
+ const abs = gguf.dataStart + t.offset;
91
+ const m = /^blk\.(\d+)\.ffn_(gate_up|down)_exps\.weight$/.exec(t.name);
92
+ if (m) {
93
+ if (t.type !== GGML.Q4_0) throw new Error(`${t.name}: experts must be Q4_0`);
94
+ const layer = Number(m[1]), gu = m[2] === 'gate_up';
95
+ const slice = tensorBytes(t) / perLayer;
96
+ for (let e = 0; e < perLayer; e++) {
97
+ const rec = (layer * perLayer + e) * record;
98
+ units.push({ kind: 'q4', src: abs + e * slice, len: slice, file: 'experts', q: rec + (gu ? parts.guQ.off : parts.dQ.off), s: rec + (gu ? parts.guS.off : parts.dS.off) });
99
+ }
100
+ continue;
101
+ }
102
+ const d = layout.dense.tensors[t.name];
103
+ if (d.type !== 'q4') { units.push({ src: abs, len: d.raw.bytes, file: 'dense', raw: d.raw.off }); continue; }
104
+ const cols = t.dims[0], rows = (d.q.bytes * 2) / cols, rowBytes = (cols / 32) * Q4_BLOCK;
105
+ const band = Math.max(1, Math.floor((8 << 20) / rowBytes));
106
+ for (let r = 0; r < rows; r += band) {
107
+ const nr = Math.min(band, rows - r);
108
+ units.push({ kind: 'q4', src: abs + r * rowBytes, len: nr * rowBytes, file: 'dense', q: d.q.off + r * cols / 2, s: d.s.off + r * (cols / 32) * 2 });
109
+ }
110
+ }
111
+ units.sort((a, b) => a.src - b.src);
112
+ return units;
113
+ }
114
+
115
+ export const GEMMA4_INGEST_PLAN = { layout: planLayoutGemma4, units: planUnitsGemma4 };
116
+
117
+ export const GEMMA4_26B_A4B = {
118
+ repo: 'google/gemma-4-26B-A4B-it-qat-q4_0-gguf',
119
+ file: 'gemma-4-26B_q4_0-it.gguf',
120
+ revision: 'd1c082be9cf3c8a514acf63b8761f4b41935842e',
121
+ sha256: '3eca3b8f6d7baf218a7dd6bba5fb59a56ee25fe2d567b6f5f589b4f697eca51d',
122
+ size: 14439363584,
123
+ };
124
+
125
+ // The GGUF's chat template (text turns): <bos>, a system turn when there is a system message or thinking is on
126
+ // (<|think|> first), then <|turn>user|model … <turn|> per message, and the model turn opened. With thinking off
127
+ // the model turn starts with an empty thought channel. Earlier model turns lose their thought channels.
128
+ const stripThinking = (t) => t.split('<channel|>').map((p) => (p.includes('<|channel>') ? p.split('<|channel>')[0] : p)).join('').trim();
129
+ export function gemmaChatPrompt(messages, { enableThinking = false } = {}) {
130
+ let s = '<bos>', msgs = messages;
131
+ const sys = msgs.length && msgs[0].role === 'system' ? String(msgs[0].content).trim() : null;
132
+ if (enableThinking || sys !== null) {
133
+ s += '<|turn>system\n' + (enableThinking ? '<|think|>\n' : '') + (sys ?? '') + '<turn|>\n';
134
+ if (sys !== null) msgs = msgs.slice(1);
135
+ }
136
+ for (const m of msgs) {
137
+ const role = m.role === 'assistant' ? 'model' : m.role;
138
+ s += `<|turn>${role}\n${role === 'model' ? stripThinking(String(m.content)) : String(m.content).trim()}<turn|>\n`;
139
+ }
140
+ return s + '<|turn>model\n' + (enableThinking ? '' : '<|channel>thought\n<channel|>');
141
+ }
142
+
143
+ // ── Tokenizer: SentencePiece-style BPE (GGUF tokenizer.ggml.model = gemma4) ──
144
+ // As llama.cpp's LLAMA_VOCAB_PRE_TYPE_GEMMA4: spaces become ▁, the text splits only into runs of newlines and
145
+ // runs of everything else, merges run on raw UTF-8 characters in rank order (leftmost first on equal ranks),
146
+ // and a run made only of newlines that is itself a token stays whole. A character with no token falls back
147
+ // to <0xXX> byte tokens (SentencePiece byte fallback; llama.cpp drops it instead).
148
+ export class GemmaTokenizer {
149
+ constructor(kv) {
150
+ this.tokens = kv['tokenizer.ggml.tokens'];
151
+ const types = kv['tokenizer.ggml.token_type'] || [];
152
+ const n = this.tokens.length;
153
+ this.ids = new Map(this.tokens.map((t, i) => [t, i]));
154
+ this.ranks = new Map((kv['tokenizer.ggml.merges'] || []).map((m, i) => [m, i]));
155
+ this.isSpecial = new Uint8Array(n); this.byteOf = new Int16Array(n).fill(-1); this.byteTok = new Int32Array(256).fill(-1);
156
+ const special = [];
157
+ for (let i = 0; i < n; i++) {
158
+ if (types[i] === 3 || types[i] === 4) { special.push(this.tokens[i]); this.isSpecial[i] = 1; }
159
+ const m = types[i] === 6 && /^<0x([0-9A-Fa-f]{2})>$/.exec(this.tokens[i]);
160
+ if (m) { const b = parseInt(m[1], 16); this.byteOf[i] = b; this.byteTok[b] = i; }
161
+ }
162
+ special.sort((a, b) => b.length - a.length);
163
+ const esc = (t) => t.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
164
+ this.specialRe = special.length ? new RegExp(`(${special.map(esc).join('|')})`) : null;
165
+ this.utf8 = new TextEncoder();
166
+ }
167
+ encode(text, { parseSpecial = true } = {}) {
168
+ const out = [];
169
+ const pieces = parseSpecial && this.specialRe ? text.split(this.specialRe) : [text];
170
+ for (let i = 0; i < pieces.length; i++) {
171
+ const piece = pieces[i];
172
+ if (!piece) continue;
173
+ if (parseSpecial && i % 2 === 1) { out.push(this.ids.get(piece)); continue; }
174
+ for (const [run] of piece.replaceAll(' ', '▁').matchAll(/[^\n]+|\n+/g)) {
175
+ const whole = run[0] === '\n' ? this.ids.get(run) : undefined;
176
+ if (whole !== undefined) out.push(whole); else this.bpe(run, out);
177
+ }
178
+ }
179
+ return out;
180
+ }
181
+ // Rank-ordered merges over a linked list of symbols, with a binary heap of candidate pairs.
182
+ bpe(word, out) {
183
+ const text = Array.from(word), n = text.length;
184
+ const prev = Int32Array.from({ length: n }, (_, i) => i - 1), next = Int32Array.from({ length: n }, (_, i) => (i + 1 < n ? i + 1 : -1));
185
+ const heap = [];
186
+ const less = (a, b) => a.rank < b.rank || (a.rank === b.rank && a.left < b.left);
187
+ const push = (e) => { heap.push(e); let i = heap.length - 1; while (i > 0) { const p = (i - 1) >> 1; if (!less(heap[i], heap[p])) break; [heap[i], heap[p]] = [heap[p], heap[i]]; i = p; } };
188
+ const pop = () => {
189
+ const top = heap[0], last = heap.pop();
190
+ if (heap.length) { heap[0] = last; let i = 0; for (;;) { const l = 2 * i + 1, r = l + 1; let m = i; if (l < heap.length && less(heap[l], heap[m])) m = l; if (r < heap.length && less(heap[r], heap[m])) m = r; if (m === i) break; [heap[i], heap[m]] = [heap[m], heap[i]]; i = m; } }
191
+ return top;
192
+ };
193
+ const pair = (left, right) => {
194
+ if (left < 0 || right < 0) return;
195
+ const t = text[left] + text[right], rank = this.ranks.get(text[left] + ' ' + text[right]);
196
+ if (rank !== undefined) push({ rank, left, right, t });
197
+ };
198
+ for (let i = 1; i < n; i++) pair(i - 1, i);
199
+ while (heap.length) {
200
+ const b = pop();
201
+ if (!text[b.left] || !text[b.right] || text[b.left] + text[b.right] !== b.t) continue;
202
+ text[b.left] = b.t; text[b.right] = '';
203
+ next[b.left] = next[b.right];
204
+ if (next[b.right] >= 0) prev[next[b.right]] = b.left;
205
+ pair(prev[b.left], b.left);
206
+ pair(b.left, next[b.left]);
207
+ }
208
+ for (let i = 0; i >= 0 && i < n; i = next[i]) {
209
+ if (!text[i]) continue;
210
+ const id = this.ids.get(text[i]);
211
+ if (id !== undefined) { out.push(id); continue; }
212
+ for (const byte of this.utf8.encode(text[i])) if (this.byteTok[byte] >= 0) out.push(this.byteTok[byte]);
213
+ }
214
+ return out;
215
+ }
216
+ tokenBytes(id) {
217
+ if (this.byteOf[id] >= 0) return Uint8Array.of(this.byteOf[id]);
218
+ const t = this.tokens[id];
219
+ return this.utf8.encode(this.isSpecial[id] ? t : t.replaceAll('▁', ' '));
220
+ }
221
+ decode(ids) {
222
+ const chunks = ids.map((id) => this.tokenBytes(id));
223
+ const all = new Uint8Array(chunks.reduce((a, c) => a + c.length, 0));
224
+ let o = 0; for (const c of chunks) { all.set(c, o); o += c.length; }
225
+ return new TextDecoder('utf-8', { ignoreBOM: true }).decode(all);
226
+ }
227
+ }
228
+
229
+ // ── WGSL ────────────────────────────────────────────────────────────────────
230
+ // Q4_0 as two planes (ingest's splitQ4): nibbles, 16 bytes per 32 values (byte j holds values j and j + 16),
231
+ // and one f16 scale per 32. Value = (nibble − 8) · scale, as ggml dequantizes it. Each thread reads one u32 of
232
+ // nibbles: values 4m..4m+3 (low nibbles) and 16+4m..16+4m+3 (high nibbles) of block b.
233
+ const Q4DOT = `
234
+ fn q4dot(wd: u32, d: f32, xb: u32) -> f32 {
235
+ let lo = (f32(wd & 15u) - 8.0) * xin[xb] + (f32((wd >> 8u) & 15u) - 8.0) * xin[xb + 1u]
236
+ + (f32((wd >> 16u) & 15u) - 8.0) * xin[xb + 2u] + (f32((wd >> 24u) & 15u) - 8.0) * xin[xb + 3u];
237
+ let hi = (f32((wd >> 4u) & 15u) - 8.0) * xin[xb + 16u] + (f32((wd >> 12u) & 15u) - 8.0) * xin[xb + 17u]
238
+ + (f32((wd >> 20u) & 15u) - 8.0) * xin[xb + 18u] + (f32((wd >> 28u) & 15u) - 8.0) * xin[xb + 19u];
239
+ return d * (lo + hi);
240
+ }`;
241
+ // Q6_K (ggml block_q6_K: ql[128], qh[64], int8 scales[16], f16 d per 256 values). The GGUF's 210-byte blocks
242
+ // sit at 2-byte alignment, so at load q6kPlanes splits the table on the GPU into aligned planes: ql (32 words
243
+ // per block), qh (16), scales (4) and d as f32. q6v reads value t of block bi exactly as dequantize_row_q6_K
244
+ // forms it: (d · scale) · (q − 32).
245
+ const Q6P = `
246
+ @group(0) @binding(0) var<storage, read> ql: array<u32>;
247
+ @group(0) @binding(1) var<storage, read> qh: array<u32>;
248
+ @group(0) @binding(2) var<storage, read> sc: array<u32>;
249
+ @group(0) @binding(3) var<storage, read> dd: array<f32>;
250
+ fn sc8(bi: u32, si: u32) -> f32 { let s8 = (sc[bi * 4u + (si >> 2u)] >> ((si & 3u) * 8u)) & 0xffu; return f32(select(i32(s8), i32(s8) - 256, s8 >= 128u)); }
251
+ fn q6v(bi: u32, t: u32) -> f32 {
252
+ let n = t >> 7u; let r = t & 127u; let qd = r >> 5u; let l = r & 31u;
253
+ let qb = n * 64u + l + (qd & 1u) * 32u;
254
+ let lo = (ql[bi * 32u + (qb >> 2u)] >> ((qb & 3u) * 8u)) & 0xffu;
255
+ let q4 = select(lo & 15u, lo >> 4u, qd >= 2u);
256
+ let hb = n * 32u + l;
257
+ let h = (qh[bi * 16u + (hb >> 2u)] >> ((hb & 3u) * 8u + qd * 2u)) & 3u;
258
+ return dd[bi] * sc8(bi, n * 8u + (l >> 4u) + qd * 2u) * f32(i32(q4 | (h << 4u)) - 32);
259
+ }`;
260
+ // tanh with its argument clamped: tanh(±10) is ±1 in f32, and a naive exp-based tanh would give NaN past ~44.
261
+ const TANH = 'fn tanhc(a: f32) -> f32 { return tanh(clamp(a, -10.0, 10.0)); }';
262
+ const GELU = `${TANH}
263
+ fn gelu(x: f32) -> f32 { return 0.5 * x * (1.0 + tanhc(0.7978845608 * x * (1.0 + 0.044715 * x * x))); }`;
264
+ // The window: sliding layers attend to the last `win` positions (llama.cpp masks p_q − p_k ≥ win).
265
+ const WIN = 'fn winStart(n: u32, win: u32) -> u32 { return select(0u, n - win, win > 0u && n > win); }';
266
+
267
+ const G4 = {
268
+ // The GGUF's Q6_K bytes (raw) → aligned planes; one thread per block.
269
+ q6kPlanes: `
270
+ struct P { blocks: u32 }
271
+ @group(0) @binding(0) var<storage, read> raw: array<u32>;
272
+ @group(0) @binding(1) var<storage, read_write> ql: array<u32>;
273
+ @group(0) @binding(2) var<storage, read_write> qh: array<u32>;
274
+ @group(0) @binding(3) var<storage, read_write> sc: array<u32>;
275
+ @group(0) @binding(4) var<storage, read_write> dd: array<f32>;
276
+ @group(0) @binding(5) var<uniform> p: P;
277
+ fn byteAt(i: u32) -> u32 { return (raw[i >> 2u] >> ((i & 3u) * 8u)) & 0xffu; }
278
+ fn wordAt(i: u32) -> u32 { return byteAt(i) | (byteAt(i + 1u) << 8u) | (byteAt(i + 2u) << 16u) | (byteAt(i + 3u) << 24u); }
279
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>) {
280
+ let b = g.y * ng.x * 256u + g.x; if (b >= p.blocks) { return; }
281
+ let o = b * 210u;
282
+ for (var i = 0u; i < 32u; i++) { ql[b * 32u + i] = wordAt(o + 4u * i); }
283
+ for (var i = 0u; i < 16u; i++) { qh[b * 16u + i] = wordAt(o + 128u + 4u * i); }
284
+ for (var i = 0u; i < 4u; i++) { sc[b * 4u + i] = wordAt(o + 192u + 4u * i); }
285
+ dd[b] = unpack2x16float(byteAt(o + 208u) | (byteAt(o + 209u) << 8u)).x;
286
+ }`,
287
+ embedQ6K: `
288
+ ${TOK}
289
+ struct P { n: u32, nb: u32, scale: f32 }
290
+ ${Q6P}
291
+ @group(0) @binding(4) var<storage, read_write> y: array<f32>;
292
+ @group(0) @binding(5) var<uniform> tok: Tok;
293
+ @group(0) @binding(6) var<uniform> p: P;
294
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
295
+ let i = g.x; if (i >= p.n) { return; }
296
+ y[i] = q6v(tok.token * p.nb + (i >> 8u), i & 255u) * p.scale;
297
+ }`,
298
+ // Logits over the Q6_K planes: 4 rows per 256-thread workgroup, 64 threads per row; thread j takes quad j
299
+ // (4 consecutive values: one ql word, one qh word, one scale) of every block of the row.
300
+ matmulQ6K: `
301
+ struct P { M: u32, N: u32 }
302
+ ${Q6P}
303
+ @group(0) @binding(4) var<storage, read> x: array<vec4<f32>>;
304
+ @group(0) @binding(5) var<storage, read_write> y: array<f32>;
305
+ @group(0) @binding(6) var<uniform> p: P;
306
+ var<workgroup> part: array<f32, 256>;
307
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
308
+ let j = l.x & 63u; let m = (wg.y * ng.x + wg.x) * 4u + (l.x >> 6u); let mr = min(m, p.M - 1u);
309
+ let nb = p.N / 256u; let n = j >> 5u; let w = j & 31u; let qd = w >> 3u; let k = w & 7u;
310
+ var acc = 0.0;
311
+ for (var bb = 0u; bb < nb; bb++) {
312
+ let bi = mr * nb + bb;
313
+ let lo = unpack4xU8(ql[bi * 32u + n * 16u + (qd & 1u) * 8u + k]);
314
+ let q4 = select(lo & vec4<u32>(15u), lo >> vec4<u32>(4u), qd >= 2u);
315
+ let h = (unpack4xU8(qh[bi * 16u + n * 8u + k]) >> vec4<u32>(qd * 2u)) & vec4<u32>(3u);
316
+ let q = vec4<f32>(vec4<i32>(q4 | (h << vec4<u32>(4u))) - vec4<i32>(32));
317
+ acc += (dd[bi] * sc8(bi, n * 8u + (k >> 2u) + qd * 2u)) * dot(q, x[bb * 64u + j]);
318
+ }
319
+ part[l.x] = acc; workgroupBarrier();
320
+ for (var s = 32u; s > 0u; s >>= 1u) { if (j < s) { part[l.x] += part[l.x + s]; } workgroupBarrier(); }
321
+ if (j == 0u && m < p.M) { y[m] = part[l.x]; }
322
+ }`,
323
+ // Dense Q4_0 GEMV, 4 output rows per workgroup (2-D dispatch ok).
324
+ matmulQ4: `enable f16;
325
+ struct P { M: u32, N: u32 }
326
+ @group(0) @binding(0) var<storage, read> xin: array<f32>;
327
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
328
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
329
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
330
+ @group(0) @binding(4) var<uniform> p: P;
331
+ ${Q4DOT}
332
+ var<workgroup> part: array<f32, 1024>;
333
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
334
+ let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let M = p.M; let nw = p.N / 8u; let nb = p.N / 32u;
335
+ if (m0 >= M) { return; }
336
+ let r0 = min(m0, M - 1u); let r1 = min(m0 + 1u, M - 1u); let r2 = min(m0 + 2u, M - 1u); let r3 = min(m0 + 3u, M - 1u);
337
+ var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
338
+ for (var w = t; w < nw; w += 256u) {
339
+ let b = w >> 2u; let xb = b * 32u + (w & 3u) * 4u;
340
+ a0 += q4dot(qw[r0 * nw + w], f32(qs[r0 * nb + b]), xb);
341
+ a1 += q4dot(qw[r1 * nw + w], f32(qs[r1 * nb + b]), xb);
342
+ a2 += q4dot(qw[r2 * nw + w], f32(qs[r2 * nb + b]), xb);
343
+ a3 += q4dot(qw[r3 * nw + w], f32(qs[r3 * nb + b]), xb);
344
+ }
345
+ part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
346
+ workgroupBarrier();
347
+ for (var s = 128u; s > 0u; s >>= 1u) {
348
+ if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
349
+ workgroupBarrier();
350
+ }
351
+ if (t < 4u && m0 + t < M) { y[m0 + t] = part[t * 256u]; }
352
+ }`,
353
+ // Routed experts, all k at once (workgroup z = top-k position); slots[z] picks the expert's rows in the pool.
354
+ expertQ4: `enable f16;
355
+ struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
356
+ @group(0) @binding(0) var<storage, read> xin: array<f32>;
357
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
358
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
359
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
360
+ @group(0) @binding(4) var<uniform> p: P;
361
+ @group(0) @binding(5) var<storage, read> slots: array<u32>;
362
+ ${Q4DOT}
363
+ var<workgroup> part: array<f32, 1024>;
364
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
365
+ let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let M = p.M; let nw = p.N / 8u; let nb = p.N / 32u;
366
+ if (m0 >= M) { return; }
367
+ let rb = slots[z] * p.rowsPerExpert;
368
+ let xo = select(0u, z * p.N, p.inputPerSlot == 1u);
369
+ let r0 = rb + min(m0, M - 1u); let r1 = rb + min(m0 + 1u, M - 1u); let r2 = rb + min(m0 + 2u, M - 1u); let r3 = rb + min(m0 + 3u, M - 1u);
370
+ var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
371
+ for (var w = t; w < nw; w += 256u) {
372
+ let b = w >> 2u; let xb = xo + b * 32u + (w & 3u) * 4u;
373
+ a0 += q4dot(qw[r0 * nw + w], f32(qs[r0 * nb + b]), xb);
374
+ a1 += q4dot(qw[r1 * nw + w], f32(qs[r1 * nb + b]), xb);
375
+ a2 += q4dot(qw[r2 * nw + w], f32(qs[r2 * nb + b]), xb);
376
+ a3 += q4dot(qw[r3 * nw + w], f32(qs[r3 * nb + b]), xb);
377
+ }
378
+ part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
379
+ workgroupBarrier();
380
+ for (var s = 128u; s > 0u; s >>= 1u) {
381
+ if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
382
+ workgroupBarrier();
383
+ }
384
+ if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
385
+ }`,
386
+ // gu = k × [gate F | up F] → act = k × F of GELU(gate)·up (ggml_geglu_split).
387
+ geglu: `
388
+ struct P { F: u32, k: u32 }
389
+ @group(0) @binding(0) var<storage, read> gu: array<f32>;
390
+ @group(0) @binding(1) var<storage, read_write> act: array<f32>;
391
+ @group(0) @binding(2) var<uniform> p: P;
392
+ ${GELU}
393
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
394
+ let i = g.x; if (i >= p.k * p.F) { return; }
395
+ let z = i / p.F; let j = i % p.F;
396
+ act[i] = gelu(gu[z * 2u * p.F + j]) * gu[z * 2u * p.F + p.F + j];
397
+ }`,
398
+ // Per-head RMS norm without a weight (V), from x into y. One workgroup per head.
399
+ vNorm: `
400
+ struct P { hd: u32, eps: f32 }
401
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
402
+ @group(0) @binding(1) var<storage, read_write> y: array<f32>;
403
+ @group(0) @binding(2) var<uniform> p: P;
404
+ var<workgroup> red: array<f32, 64>;
405
+ @compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
406
+ let base = wg.x * p.hd; let t = l.x;
407
+ var acc = 0.0;
408
+ for (var i = t; i < p.hd; i += 64u) { let v = x[base + i]; acc += v * v; }
409
+ red[t] = acc; workgroupBarrier();
410
+ for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
411
+ let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
412
+ for (var i = t; i < p.hd; i += 64u) { y[base + i] = x[base + i] * scale; }
413
+ }`,
414
+ // y = ((x · rms) · k) · w — the router's input (k = 1/√H, w = ffn_gate_inp.scale).
415
+ rmsnormK: `
416
+ struct P { n: u32, eps: f32, k: f32 }
417
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
418
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
419
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
420
+ @group(0) @binding(3) var<uniform> p: P;
421
+ var<workgroup> red: array<f32, 256>;
422
+ @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
423
+ let t = l.x; var acc = 0.0;
424
+ for (var i = t; i < p.n; i += 256u) { let v = x[i]; acc += v * v; }
425
+ red[t] = acc; workgroupBarrier();
426
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
427
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
428
+ for (var i = t; i < p.n; i += 256u) { y[i] = ((x[i] * scale) * p.k) * w[i]; }
429
+ }`,
430
+ // x += rms_norm(y) · w — the post-attention norm and the residual add.
431
+ rmsnormAdd: `
432
+ struct P { n: u32, eps: f32 }
433
+ @group(0) @binding(0) var<storage, read> y: array<f32>;
434
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
435
+ @group(0) @binding(2) var<storage, read_write> x: array<f32>;
436
+ @group(0) @binding(3) var<uniform> p: P;
437
+ var<workgroup> red: array<f32, 256>;
438
+ @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
439
+ let t = l.x; var acc = 0.0;
440
+ for (var i = t; i < p.n; i += 256u) { let v = y[i]; acc += v * v; }
441
+ red[t] = acc; workgroupBarrier();
442
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
443
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
444
+ for (var i = t; i < p.n; i += 256u) { x[i] = (y[i] * scale) * w[i] + x[i]; }
445
+ }`,
446
+ // x = (rms_norm(mlp + moe) · w + x) · out_scale — the FFN's last norm, the residual add, the layer scalar.
447
+ postFfn: `
448
+ struct P { n: u32, eps: f32 }
449
+ @group(0) @binding(0) var<storage, read> mlp: array<f32>;
450
+ @group(0) @binding(1) var<storage, read> moe: array<f32>;
451
+ @group(0) @binding(2) var<storage, read> w: array<f32>;
452
+ @group(0) @binding(3) var<storage, read> os: array<f32>;
453
+ @group(0) @binding(4) var<storage, read_write> x: array<f32>;
454
+ @group(0) @binding(5) var<uniform> p: P;
455
+ var<workgroup> red: array<f32, 256>;
456
+ @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
457
+ let t = l.x; var acc = 0.0;
458
+ for (var i = t; i < p.n; i += 256u) { let v = mlp[i] + moe[i]; acc += v * v; }
459
+ red[t] = acc; workgroupBarrier();
460
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
461
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
462
+ let k = os[0];
463
+ for (var i = t; i < p.n; i += 256u) { x[i] = (((mlp[i] + moe[i]) * scale) * w[i] + x[i]) * k; }
464
+ }`,
465
+ // moe[h] = Σ_j (dn[j·H + h] · ds[id_j]) · w_j, added in pick order (llama.cpp: down · scale · weight, then e0 + e1 + …).
466
+ moeSum: `
467
+ struct P { H: u32, k: u32 }
468
+ @group(0) @binding(0) var<storage, read> dn: array<f32>;
469
+ @group(0) @binding(1) var<storage, read> sel: array<u32>;
470
+ @group(0) @binding(2) var<storage, read> ds: array<f32>;
471
+ @group(0) @binding(3) var<storage, read_write> moe: array<f32>;
472
+ @group(0) @binding(4) var<uniform> p: P;
473
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
474
+ let h = g.x; if (h >= p.H) { return; }
475
+ var acc = (dn[h] * ds[sel[0]]) * bitcast<f32>(sel[p.k]);
476
+ for (var j = 1u; j < p.k; j++) { acc = acc + (dn[j * p.H + h] * ds[sel[j]]) * bitcast<f32>(sel[p.k + j]); }
477
+ moe[h] = acc;
478
+ }`,
479
+ // Attention with a window: scores, softmax and the weighted sum over positions [start, seqLen); scale 1.0.
480
+ attnScoreW: `enable f16;
481
+ ${TOK}
482
+ struct P { heads: u32, kvHeads: u32, hd: u32, win: u32 }
483
+ @group(0) @binding(0) var<storage, read> q: array<f32>;
484
+ @group(0) @binding(1) var<storage, read> kc: array<f16>;
485
+ @group(0) @binding(2) var<storage, read_write> sc: array<f32>;
486
+ @group(0) @binding(3) var<uniform> tok: Tok;
487
+ @group(0) @binding(4) var<uniform> p: P;
488
+ ${WIN}
489
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
490
+ let st = winStart(tok.seqLen, p.win); let n = tok.seqLen - st; let idx = g.x;
491
+ if (idx >= p.heads * n) { return; }
492
+ let h = idx / n; let j = idx % n; let kvh = h / (p.heads / p.kvHeads);
493
+ let qo = h * p.hd; let ko = (st + j) * p.kvHeads * p.hd + kvh * p.hd;
494
+ var acc = 0.0;
495
+ for (var d = 0u; d < p.hd; d++) { acc += q[qo + d] * f32(kc[ko + d]); }
496
+ sc[h * n + j] = acc;
497
+ }`,
498
+ softmaxW: `
499
+ ${TOK}
500
+ struct P { heads: u32, kvHeads: u32, hd: u32, win: u32 }
501
+ @group(0) @binding(0) var<storage, read_write> sc: array<f32>;
502
+ @group(0) @binding(1) var<uniform> tok: Tok;
503
+ @group(0) @binding(2) var<uniform> p: P;
504
+ ${WIN}
505
+ var<workgroup> red: array<f32, 256>;
506
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
507
+ let n = tok.seqLen - winStart(tok.seqLen, p.win); let base = wg.x * n; let t = l.x;
508
+ var mx = -3.4e38;
509
+ for (var i = t; i < n; i += 256u) { mx = max(mx, sc[base + i]); }
510
+ red[t] = mx; workgroupBarrier();
511
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
512
+ let m = red[0]; workgroupBarrier();
513
+ var sum = 0.0;
514
+ for (var i = t; i < n; i += 256u) { let e = exp(sc[base + i] - m); sc[base + i] = e; sum += e; }
515
+ red[t] = sum; workgroupBarrier();
516
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
517
+ let inv = 1.0 / red[0];
518
+ for (var i = t; i < n; i += 256u) { sc[base + i] = sc[base + i] * inv; }
519
+ }`,
520
+ attnOutW: `enable f16;
521
+ ${TOK}
522
+ struct P { heads: u32, kvHeads: u32, hd: u32, win: u32 }
523
+ @group(0) @binding(0) var<storage, read> pr: array<f32>;
524
+ @group(0) @binding(1) var<storage, read> vc: array<f16>;
525
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
526
+ @group(0) @binding(3) var<uniform> tok: Tok;
527
+ @group(0) @binding(4) var<uniform> p: P;
528
+ ${WIN}
529
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
530
+ let idx = g.x; if (idx >= p.heads * p.hd) { return; }
531
+ let st = winStart(tok.seqLen, p.win); let n = tok.seqLen - st;
532
+ let h = idx / p.hd; let d = idx % p.hd; let kvh = h / (p.heads / p.kvHeads);
533
+ let stride = p.kvHeads * p.hd;
534
+ var acc = 0.0;
535
+ for (var j = 0u; j < n; j++) { acc += pr[h * n + j] * f32(vc[(st + j) * stride + kvh * p.hd + d]); }
536
+ y[idx] = acc;
537
+ }`,
538
+ // Final logit soft cap, in place: l = tanh(l · (1/cap)) · cap.
539
+ softcap: `
540
+ struct P { n: u32, inv: f32, cap: f32 }
541
+ @group(0) @binding(0) var<storage, read_write> lg: array<f32>;
542
+ @group(0) @binding(1) var<uniform> p: P;
543
+ ${TANH}
544
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>) {
545
+ let i = g.y * ng.x * 256u + g.x; if (i >= p.n) { return; }
546
+ lg[i] = tanhc(lg[i] * p.inv) * p.cap;
547
+ }`,
548
+ // Debug capture: y[off + i] = x[i].
549
+ copyAt: `
550
+ struct P { n: u32, off: u32 }
551
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
552
+ @group(0) @binding(1) var<storage, read_write> y: array<f32>;
553
+ @group(0) @binding(2) var<uniform> p: P;
554
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
555
+ let i = g.x; if (i < p.n) { y[p.off + i] = x[i]; }
556
+ }`,
557
+ };
558
+
559
+ // ── Subgroup Q4_0 GEMVs (decode) ─────────────────────────────────────────────────
560
+ // One 32-lane subgroup per output row, 4 rows per 128-thread workgroup (the scalar kernels' dispatch
561
+ // geometry). A lane takes whole blocks: one vec4<u32> of nibbles, unpack4xU8 of the low and high nibble
562
+ // masks, vec4 dot products against the activations, the block scale once; subgroupAdd reduces the row.
563
+ // Used only when the adapter's subgroups are exactly 32 wide.
564
+ const Q4BLK = `
565
+ fn q4blk(w: vec4<u32>, xb: u32) -> f32 {
566
+ let m = 0x0f0f0f0fu;
567
+ return dot(vec4<f32>(unpack4xU8(w.x & m)) - 8.0, x[xb]) + dot(vec4<f32>(unpack4xU8((w.x >> 4u) & m)) - 8.0, x[xb + 4u])
568
+ + dot(vec4<f32>(unpack4xU8(w.y & m)) - 8.0, x[xb + 1u]) + dot(vec4<f32>(unpack4xU8((w.y >> 4u) & m)) - 8.0, x[xb + 5u])
569
+ + dot(vec4<f32>(unpack4xU8(w.z & m)) - 8.0, x[xb + 2u]) + dot(vec4<f32>(unpack4xU8((w.z >> 4u) & m)) - 8.0, x[xb + 6u])
570
+ + dot(vec4<f32>(unpack4xU8(w.w & m)) - 8.0, x[xb + 3u]) + dot(vec4<f32>(unpack4xU8((w.w >> 4u) & m)) - 8.0, x[xb + 7u]);
571
+ }`;
572
+ const G4SG = {
573
+ matmulQ4: `enable f16;
574
+ enable subgroups;
575
+ struct P { M: u32, N: u32 }
576
+ @group(0) @binding(0) var<storage, read> x: array<vec4<f32>>;
577
+ @group(0) @binding(1) var<storage, read> qw: array<vec4<u32>>;
578
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
579
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
580
+ @group(0) @binding(4) var<uniform> p: P;
581
+ ${Q4BLK}
582
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
583
+ let m = (wg.y * ng.x + wg.x) * 4u + l.x / 32u; let lane = l.x % 32u;
584
+ let nb = p.N / 32u; let rb = min(m, p.M - 1u) * nb;
585
+ var acc = 0.0;
586
+ for (var b = lane; b < nb; b += 32u) { acc += q4blk(qw[rb + b], b * 8u) * f32(qs[rb + b]); }
587
+ let total = subgroupAdd(acc);
588
+ if (lane == 0u && m < p.M) { y[m] = total; }
589
+ }`,
590
+ expertQ4: `enable f16;
591
+ enable subgroups;
592
+ struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
593
+ @group(0) @binding(0) var<storage, read> x: array<vec4<f32>>;
594
+ @group(0) @binding(1) var<storage, read> qw: array<vec4<u32>>;
595
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
596
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
597
+ @group(0) @binding(4) var<uniform> p: P;
598
+ @group(0) @binding(5) var<storage, read> slots: array<u32>;
599
+ ${Q4BLK}
600
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
601
+ let z = wg.z; let m = (wg.y * ng.x + wg.x) * 4u + l.x / 32u; let lane = l.x % 32u;
602
+ let nb = p.N / 32u; let rb = (slots[z] * p.rowsPerExpert + min(m, p.M - 1u)) * nb;
603
+ let xo = select(0u, z * p.N / 4u, p.inputPerSlot == 1u);
604
+ var acc = 0.0;
605
+ for (var b = lane; b < nb; b += 32u) { acc += q4blk(qw[rb + b], xo + b * 8u) * f32(qs[rb + b]); }
606
+ let total = subgroupAdd(acc);
607
+ if (lane == 0u && m < p.M) { y[z * p.M + m] = total; }
608
+ }`,
609
+ };
610
+
611
+ // ── Chunked prefill: T tokens per layer pass (rung 2b's pattern) ────────────────
612
+ // Rows are tokens, token-major. Generic kernels come from rung 2b (rmsnormB, matmulF32T, topkB, kvStoreB,
613
+ // softmaxB, attnOutB); these are Gemma's.
614
+ const QB = 'struct Q { T: u32, pos0: u32, S: u32, pad: u32 }';
615
+ const G4B = {
616
+ embedQ6KB: `
617
+ struct P { n: u32, nb: u32, scale: f32 }
618
+ ${QB}
619
+ ${Q6P}
620
+ @group(0) @binding(4) var<storage, read> ids: array<u32>;
621
+ @group(0) @binding(5) var<storage, read_write> y: array<f32>;
622
+ @group(0) @binding(6) var<uniform> p: P;
623
+ @group(0) @binding(7) var<uniform> qd: Q;
624
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
625
+ let i = g.x; let t = g.y; if (i >= p.n || t >= qd.T) { return; }
626
+ y[t * p.n + i] = q6v(ids[t] * p.nb + (i >> 8u), i & 255u) * p.scale;
627
+ }`,
628
+ // Tiled Q4_0 GEMM: y[t·M + m] = Σ_k x[t·N + k]·W[m][k]. 16 rows × 16 tokens per workgroup; each 32-wide
629
+ // block of W (decoded from its nibbles) and of x goes through shared memory, so a weight serves 16 tokens.
630
+ // Dispatch (ceil(M/16), ceil(T/16)).
631
+ matmulQ4T: `enable f16;
632
+ struct P { M: u32, N: u32 }
633
+ ${QB}
634
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
635
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
636
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
637
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
638
+ @group(0) @binding(4) var<uniform> p: P;
639
+ @group(0) @binding(5) var<uniform> qd: Q;
640
+ var<workgroup> ws: array<f32, 512>;
641
+ var<workgroup> xs: array<f32, 512>;
642
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
643
+ let m0 = wg.x * 16u; let t0 = wg.y * 16u; let r = l.x / 16u; let c = l.x % 16u;
644
+ let nb = p.N / 32u; let nw = p.N / 8u; var acc = 0.0;
645
+ for (var b = 0u; b < nb; b++) {
646
+ for (var i = l.x; i < 512u; i += 256u) {
647
+ let rr = i / 32u; let kk = i % 32u;
648
+ let m = m0 + rr; var wv = 0.0;
649
+ if (m < p.M) {
650
+ let wd = qw[m * nw + b * 4u + ((kk & 15u) >> 2u)];
651
+ wv = (f32((wd >> ((kk & 3u) * 8u + (kk >> 4u) * 4u)) & 15u) - 8.0) * f32(qs[m * nb + b]);
652
+ }
653
+ ws[i] = wv;
654
+ let t = t0 + rr; var xv = 0.0;
655
+ if (t < qd.T) { xv = x[t * p.N + b * 32u + kk]; }
656
+ xs[i] = xv;
657
+ }
658
+ workgroupBarrier();
659
+ for (var kk = 0u; kk < 32u; kk++) { acc += ws[r * 32u + kk] * xs[c * 32u + kk]; }
660
+ workgroupBarrier();
661
+ }
662
+ let m = m0 + r; let t = t0 + c;
663
+ if (m < p.M && t < qd.T) { y[t * p.M + m] = acc; }
664
+ }`,
665
+ // Routed experts for every (token, pick) pair: z = t·K + j, slots[z] its pool slot. Input: the token's row
666
+ // (gate/up, inputMode 0) or the pair's own row (down, inputMode 1).
667
+ expertQ4B: `enable f16;
668
+ struct P { M: u32, N: u32, rowsPerExpert: u32, inputMode: u32, K: u32 }
669
+ @group(0) @binding(0) var<storage, read> xin: array<f32>;
670
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
671
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
672
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
673
+ @group(0) @binding(4) var<uniform> p: P;
674
+ @group(0) @binding(5) var<storage, read> slots: array<u32>;
675
+ ${Q4DOT}
676
+ var<workgroup> part: array<f32, 1024>;
677
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
678
+ let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let M = p.M; let nw = p.N / 8u; let nb = p.N / 32u;
679
+ if (m0 >= M) { return; }
680
+ let rb = slots[z] * p.rowsPerExpert;
681
+ let xo = select((z / p.K) * p.N, z * p.N, p.inputMode == 1u);
682
+ let r0 = rb + min(m0, M - 1u); let r1 = rb + min(m0 + 1u, M - 1u); let r2 = rb + min(m0 + 2u, M - 1u); let r3 = rb + min(m0 + 3u, M - 1u);
683
+ var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
684
+ for (var w = t; w < nw; w += 256u) {
685
+ let b = w >> 2u; let xb = xo + b * 32u + (w & 3u) * 4u;
686
+ a0 += q4dot(qw[r0 * nw + w], f32(qs[r0 * nb + b]), xb);
687
+ a1 += q4dot(qw[r1 * nw + w], f32(qs[r1 * nb + b]), xb);
688
+ a2 += q4dot(qw[r2 * nw + w], f32(qs[r2 * nb + b]), xb);
689
+ a3 += q4dot(qw[r3 * nw + w], f32(qs[r3 * nb + b]), xb);
690
+ }
691
+ part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
692
+ workgroupBarrier();
693
+ for (var s = 128u; s > 0u; s >>= 1u) {
694
+ if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
695
+ workgroupBarrier();
696
+ }
697
+ if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
698
+ }`,
699
+ // Per-head RMS norm and NeoX RoPE in place for (head wg.x, token wg.y); rope holds hd values per token.
700
+ qkNormRopeB: `
701
+ struct P { hd: u32, row: u32, eps: f32 }
702
+ @group(0) @binding(0) var<storage, read_write> x: array<f32>;
703
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
704
+ @group(0) @binding(2) var<storage, read> rope: array<f32>;
705
+ @group(0) @binding(3) var<uniform> p: P;
706
+ var<workgroup> red: array<f32, 64>;
707
+ @compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
708
+ let base = wg.y * p.row + wg.x * p.hd; let rb = wg.y * p.hd; let t = l.x; let half = p.hd / 2u;
709
+ var acc = 0.0;
710
+ for (var i = t; i < p.hd; i += 64u) { let v = x[base + i]; acc += v * v; }
711
+ red[t] = acc; workgroupBarrier();
712
+ for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
713
+ let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
714
+ for (var i = t; i < half; i += 64u) {
715
+ let a = (x[base + i] * scale) * w[i];
716
+ let b = (x[base + i + half] * scale) * w[i + half];
717
+ let c = rope[rb + i]; let sn = rope[rb + half + i];
718
+ x[base + i] = a * c - b * sn;
719
+ x[base + i + half] = a * sn + b * c;
720
+ }
721
+ }`,
722
+ vNormB: `
723
+ struct P { hd: u32, row: u32, eps: f32 }
724
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
725
+ @group(0) @binding(1) var<storage, read_write> y: array<f32>;
726
+ @group(0) @binding(2) var<uniform> p: P;
727
+ var<workgroup> red: array<f32, 64>;
728
+ @compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
729
+ let base = wg.y * p.row + wg.x * p.hd; let t = l.x;
730
+ var acc = 0.0;
731
+ for (var i = t; i < p.hd; i += 64u) { let v = x[base + i]; acc += v * v; }
732
+ red[t] = acc; workgroupBarrier();
733
+ for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
734
+ let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
735
+ for (var i = t; i < p.hd; i += 64u) { y[base + i] = x[base + i] * scale; }
736
+ }`,
737
+ // Causal (and, on sliding layers, windowed) scores for T queries over S = pos0 + T keys; scale 1.0.
738
+ attnScoreWB: `enable f16;
739
+ struct P { heads: u32, kvHeads: u32, hd: u32, win: u32 }
740
+ ${QB}
741
+ @group(0) @binding(0) var<storage, read> q: array<f32>;
742
+ @group(0) @binding(1) var<storage, read> kc: array<f16>;
743
+ @group(0) @binding(2) var<storage, read_write> sc: array<f32>;
744
+ @group(0) @binding(3) var<uniform> p: P;
745
+ @group(0) @binding(4) var<uniform> qd: Q;
746
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>) {
747
+ let idx = g.y * ng.x * 256u + g.x; let S = qd.S; let T = qd.T;
748
+ if (idx >= p.heads * T * S) { return; }
749
+ let h = idx / (T * S); let rem = idx % (T * S); let t = rem / S; let pos = rem % S; let qp = qd.pos0 + t;
750
+ if (pos > qp || (p.win > 0u && qp - pos >= p.win)) { sc[idx] = -3.4e38; return; }
751
+ let kvh = h / (p.heads / p.kvHeads);
752
+ let qo = t * p.heads * p.hd + h * p.hd; let ko = pos * p.kvHeads * p.hd + kvh * p.hd;
753
+ var acc = 0.0;
754
+ for (var d = 0u; d < p.hd; d++) { acc += q[qo + d] * f32(kc[ko + d]); }
755
+ sc[idx] = acc;
756
+ }`,
757
+ // x += rms_norm(y) · w per token row (one workgroup per token).
758
+ rmsnormAddB: `
759
+ struct P { n: u32, eps: f32 }
760
+ @group(0) @binding(0) var<storage, read> y: array<f32>;
761
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
762
+ @group(0) @binding(2) var<storage, read_write> x: array<f32>;
763
+ @group(0) @binding(3) var<uniform> p: P;
764
+ var<workgroup> red: array<f32, 256>;
765
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
766
+ let b = wg.x * p.n; let t = l.x; var acc = 0.0;
767
+ for (var i = t; i < p.n; i += 256u) { let v = y[b + i]; acc += v * v; }
768
+ red[t] = acc; workgroupBarrier();
769
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
770
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
771
+ for (var i = t; i < p.n; i += 256u) { x[b + i] = (y[b + i] * scale) * w[i] + x[b + i]; }
772
+ }`,
773
+ rmsnormKB: `
774
+ struct P { n: u32, eps: f32, k: f32 }
775
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
776
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
777
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
778
+ @group(0) @binding(3) var<uniform> p: P;
779
+ var<workgroup> red: array<f32, 256>;
780
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
781
+ let b = wg.x * p.n; let t = l.x; var acc = 0.0;
782
+ for (var i = t; i < p.n; i += 256u) { let v = x[b + i]; acc += v * v; }
783
+ red[t] = acc; workgroupBarrier();
784
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
785
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
786
+ for (var i = t; i < p.n; i += 256u) { y[b + i] = ((x[b + i] * scale) * p.k) * w[i]; }
787
+ }`,
788
+ // act = GELU(g) · u over T × n (the dense MLP's separate gate and up).
789
+ gegluSplitB: `
790
+ struct P { n: u32 }
791
+ ${QB}
792
+ @group(0) @binding(0) var<storage, read> g: array<f32>;
793
+ @group(0) @binding(1) var<storage, read> u: array<f32>;
794
+ @group(0) @binding(2) var<storage, read_write> act: array<f32>;
795
+ @group(0) @binding(3) var<uniform> p: P;
796
+ @group(0) @binding(4) var<uniform> qd: Q;
797
+ ${GELU}
798
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
799
+ let i = gi.x; let t = gi.y; if (i >= p.n || t >= qd.T) { return; }
800
+ let k = t * p.n + i; act[k] = gelu(g[k]) * u[k];
801
+ }`,
802
+ moeSumB: `
803
+ struct P { H: u32, k: u32 }
804
+ ${QB}
805
+ @group(0) @binding(0) var<storage, read> dn: array<f32>;
806
+ @group(0) @binding(1) var<storage, read> sel: array<u32>;
807
+ @group(0) @binding(2) var<storage, read> ds: array<f32>;
808
+ @group(0) @binding(3) var<storage, read_write> moe: array<f32>;
809
+ @group(0) @binding(4) var<uniform> p: P;
810
+ @group(0) @binding(5) var<uniform> qd: Q;
811
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
812
+ let h = gi.x; let t = gi.y; if (h >= p.H || t >= qd.T) { return; }
813
+ let sb = t * 2u * p.k; let ob = t * p.k * p.H + h;
814
+ var acc = (dn[ob] * ds[sel[sb]]) * bitcast<f32>(sel[sb + p.k]);
815
+ for (var j = 1u; j < p.k; j++) { acc = acc + (dn[ob + j * p.H] * ds[sel[sb + j]]) * bitcast<f32>(sel[sb + p.k + j]); }
816
+ moe[t * p.H + h] = acc;
817
+ }`,
818
+ postFfnB: `
819
+ struct P { n: u32, eps: f32 }
820
+ @group(0) @binding(0) var<storage, read> mlp: array<f32>;
821
+ @group(0) @binding(1) var<storage, read> moe: array<f32>;
822
+ @group(0) @binding(2) var<storage, read> w: array<f32>;
823
+ @group(0) @binding(3) var<storage, read> os: array<f32>;
824
+ @group(0) @binding(4) var<storage, read_write> x: array<f32>;
825
+ @group(0) @binding(5) var<uniform> p: P;
826
+ var<workgroup> red: array<f32, 256>;
827
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
828
+ let b = wg.x * p.n; let t = l.x; var acc = 0.0;
829
+ for (var i = t; i < p.n; i += 256u) { let v = mlp[b + i] + moe[b + i]; acc += v * v; }
830
+ red[t] = acc; workgroupBarrier();
831
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
832
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
833
+ let k = os[0];
834
+ for (var i = t; i < p.n; i += 256u) { x[b + i] = (((mlp[b + i] + moe[b + i]) * scale) * w[i] + x[b + i]) * k; }
835
+ }`,
836
+ };
837
+ const REUSED_B = ['rmsnormB', 'matmulF32T', 'topkB', 'kvStoreB', 'softmaxB', 'attnOutB'];
838
+
839
+ // ── Engine ──────────────────────────────────────────────────────────────────
840
+ export class Gemma4MoeSsd extends Qwen3MoeSsd {
841
+ static configFrom(kv) { return configFromGgufGemma4(kv); }
842
+ static get ingestPlan() { return GEMMA4_INGEST_PLAN; }
843
+ static tokenizerFrom(kv) { return new GemmaTokenizer(kv); }
844
+ static load(modelId, opts = {}) { return super.load(modelId, { ...opts, source: opts.source || GEMMA4_26B_A4B }); }
845
+
846
+ chatPrompt(messages, opts) { return gemmaChatPrompt(messages, opts); }
847
+ get stopTokenIds() { return [this.cfg.eos, this.tokenizer.ids.get('<turn|>')].filter((t) => t !== undefined); }
848
+ // Thinking: the model writes <|channel>thought\n … <channel|> itself (nothing opens it in the prompt).
849
+ get thinkOpenTokenId() { return this.tokenizer.ids.get('<|channel>') ?? null; }
850
+ get thinkCloseTokenId() { return this.tokenizer.ids.get('<channel|>') ?? null; }
851
+ get thinkInPrompt() { return false; }
852
+ get thinkPreamble() { return 'thought\n'; }
853
+
854
+ get kernels() {
855
+ return (this._kernels ||= {
856
+ ...super.kernels, ...G4, ...(this.subgroupKernels ? G4SG : {}),
857
+ ...(this.batchPrefill ? { ...Object.fromEntries(REUSED_B.map((k) => [k, QWEN35_BATCH_KERNELS[k]])), ...G4B } : {}),
858
+ });
859
+ }
860
+ // Subgroup GEMVs replace the scalar Q4_0 ones when subgroups are exactly 32 wide (Apple, most others).
861
+ get subgroupKernels() {
862
+ const i = this.device.adapterInfo || this.opts.adapterInfo || {};
863
+ return this.opts.subgroups !== false && this.device.features.has('subgroups') && i.subgroupMinSize === 32 && i.subgroupMaxSize === 32;
864
+ }
865
+
866
+ // batchPrefill (default on): prompts go through in chunks of up to prefillChunk tokens, one routing
867
+ // readback per layer per chunk. capture (debug) records one-token steps, so it turns chunking off.
868
+ constructor(device, manifest, gguf, opts) {
869
+ super(device, manifest, gguf, opts);
870
+ this.batchPrefill = (opts.batchPrefill ?? true) && !opts.capture;
871
+ this.chunkTokens = Math.max(2, Math.floor(opts.prefillChunk || 256));
872
+ }
873
+
874
+ // A bind group whose entries are buffers or { buffer, offset, size } ranges.
875
+ bindR(name, entries) {
876
+ return this.device.createBindGroup({
877
+ layout: this.pipeline(name).getBindGroupLayout(0),
878
+ entries: entries.map((e, i) => ({ binding: i, resource: e.buffer ? e : { buffer: e } })),
879
+ });
880
+ }
881
+
882
+ // Dense upload; the Q6_K token table split into aligned planes on the GPU (the raw copy is then freed);
883
+ // the full-attention layers' RoPE frequency factors on the CPU side.
884
+ async uploadDense(onProgress) {
885
+ await super.uploadDense(onProgress);
886
+ const dev = this.device, raw = this.w['token_embd.weight'].raw, blocks = raw.size / Q6K_BLOCK | 0;
887
+ const P = (n) => this.buffer(n, STORAGE, 'dense');
888
+ this.q6 = { ql: P(blocks * 128), qh: P(blocks * 64), sc: P(blocks * 16), dd: P(blocks * 4) };
889
+ const enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
890
+ this.dispatch(pass, 'q6kPlanes', this.bind('q6kPlanes', [raw, this.q6.ql, this.q6.qh, this.q6.sc, this.q6.dd, this.uniform([blocks])]), Math.ceil(blocks / 256));
891
+ pass.end(); dev.queue.submit([enc.finish()]);
892
+ await dev.queue.onSubmittedWorkDone();
893
+ this.gpuBytes.dense -= raw.size; raw.destroy(); this.w['token_embd.weight'].raw = null;
894
+ const d = this.manifest.dense.tensors['rope_freqs.weight'];
895
+ this.ropeFreqs = null;
896
+ if (d) {
897
+ const r = await OpfsReaderPool.open(`${this.opts.dir}/${this.manifest.dense.file}`, { workers: 1 });
898
+ const got = await r.read(d.raw.off, d.raw.bytes);
899
+ this.ropeFreqs = new Float32Array(got.buf.slice(0, d.raw.bytes));
900
+ await r.close();
901
+ }
902
+ }
903
+
904
+ initBuffers() {
905
+ const c = this.cfg, H = c.hidden, K = c.topK, E = c.experts, F = c.expertFf, DF = c.denseFf, Lc = c.layers;
906
+ if (2 * K * 4 > 256) throw new Error(`top-k ${K} > 32`);
907
+ if (E > 256) throw new Error(`${E} experts: the router top-k kernel handles at most 256`);
908
+ const kvn = (l) => c.kvHeads[l] * c.headDim[l], qn = (l) => c.heads * c.headDim[l];
909
+ const QN = Math.max(...c.headDim.map((d) => c.heads * d)), KVN = Math.max(...c.headDim.map((_, l) => kvn(l)));
910
+ const HD = Math.max(...c.headDim);
911
+ this.kc = []; this.vc = [];
912
+ for (let l = 0; l < Lc; l++) {
913
+ this.kc.push(this.buffer(this.maxCtx * kvn(l) * 2, STORAGE, 'kv'));
914
+ this.vc.push(this.buffer(this.maxCtx * kvn(l) * 2, STORAGE, 'kv'));
915
+ }
916
+ const A = (n) => this.buffer(n * 4, STORAGE | COPY_SRC | COPY_DST);
917
+ const gateBytes = DF * 4;
918
+ if (gateBytes % 256) throw new Error(`dense FFN ${DF}: gate/up halves must sit on 256-byte offsets`);
919
+ this.a = {
920
+ x: A(H), xn: A(H), q: A(QN), k: A(KVN), v: A(KVN), vr: A(KVN), att: A(QN), o: A(H), sc: A(c.heads * this.maxCtx),
921
+ mgu: A(2 * DF), ma: A(DF), md: A(H), mlp: A(H), xm: A(H), rt: A(H), moe: A(H), moe2: A(H),
922
+ rl: A(E), sel: A(2 * K), prt: A(H), prl: A(E), psel: A(64 * this.lookahead), slots: A(K),
923
+ gu: A(K * 2 * F), act: A(K * F), dn: A(K * H), logits: A(c.vocab), am: A(4),
924
+ ropeS: A(HD), ropeF: A(HD),
925
+ };
926
+ if (this.opts.capture) this.a.cap = A(2 * Lc * H);
927
+ this.tok = this.buffer(16, 0x40 | COPY_DST);
928
+ this.rbSel = this.buffer(256 * (1 + this.lookahead), MAP_READ | COPY_DST);
929
+ this.rbLogits = this.buffer(c.vocab * 4, MAP_READ | COPY_DST);
930
+ this.rbArg = this.buffer(16, MAP_READ | COPY_DST);
931
+
932
+ const f = (v) => ({ f: v });
933
+ const eps = f(c.eps);
934
+ const u = {
935
+ embed: this.uniform([H, H / 256, f(Math.fround(Math.sqrt(H)))]),
936
+ rmsH: this.uniform([H, eps]), routerK: this.uniform([H, eps, f(Math.fround(1 / Math.fround(Math.sqrt(H))))]),
937
+ mmG: this.uniform([DF, H]), mmD: this.uniform([H, DF]), gegluD: this.uniform([DF, 1]),
938
+ router: this.uniform([E, H]), topk: this.uniform([E, K]),
939
+ gu: this.uniform([2 * F, H, 2 * F, 0]), dn: this.uniform([H, F, H, 1]), gegluE: this.uniform([F, K]), sum: this.uniform([H, K]),
940
+ lm: this.uniform([c.vocab, H]), cap: this.uniform([c.vocab, f(Math.fround(1 / c.softcap)), f(c.softcap)]), am: this.uniform([c.vocab]),
941
+ };
942
+ // Attention uniforms per layer shape (sliding / full).
943
+ const shape = new Map();
944
+ const shapeOf = (l) => {
945
+ const key = `${c.headDim[l]}/${c.kvHeads[l]}/${c.swa[l]}`;
946
+ if (!shape.has(key)) shape.set(key, {
947
+ mmQ: this.uniform([qn(l), H]), mmKV: this.uniform([kvn(l), H]), mmO: this.uniform([H, qn(l)]),
948
+ hd: this.uniform([c.headDim[l], eps]), kvn: this.uniform([kvn(l)]),
949
+ att: this.uniform([c.heads, c.kvHeads[l], c.headDim[l], c.swa[l] ? c.window : 0]),
950
+ });
951
+ return shape.get(key);
952
+ };
953
+ const a = this.a, W = this.w;
954
+ const q6 = [this.q6.ql, this.q6.qh, this.q6.sc, this.q6.dd];
955
+ this.g = { embed: this.bind('embedQ6K', [...q6, a.x, this.tok, u.embed]) };
956
+ this.layers = [];
957
+ for (let l = 0; l < Lc; l++) {
958
+ const p = (n) => W[`blk.${l}.${n}`], s = shapeOf(l), rope = c.swa[l] ? a.ropeS : a.ropeF;
959
+ const hasV = !!p('attn_v.weight');
960
+ const ly = {
961
+ swa: c.swa[l], hasV,
962
+ rmsA: this.bind('rmsnorm', [a.x, p('attn_norm.weight').raw, a.xn, u.rmsH]),
963
+ q: this.bind('matmulQ4', [a.xn, p('attn_q.weight').q, p('attn_q.weight').s, a.q, s.mmQ]),
964
+ k: this.bind('matmulQ4', [a.xn, p('attn_k.weight').q, p('attn_k.weight').s, a.k, s.mmKV]),
965
+ v: hasV ? this.bind('matmulQ4', [a.xn, p('attn_v.weight').q, p('attn_v.weight').s, a.vr, s.mmKV]) : null,
966
+ vn: this.bind('vNorm', [hasV ? a.vr : a.k, a.v, s.hd]), // no V projection: V is K's projection
967
+ ropeQ: this.bind('qkNormRope', [a.q, p('attn_q_norm.weight').raw, rope, s.hd]),
968
+ ropeK: this.bind('qkNormRope', [a.k, p('attn_k_norm.weight').raw, rope, s.hd]),
969
+ kv: this.bind('kvStore', [a.k, a.v, this.kc[l], this.vc[l], this.tok, s.kvn]),
970
+ score: this.bind('attnScoreW', [a.q, this.kc[l], a.sc, this.tok, s.att]),
971
+ soft: this.bind('softmaxW', [a.sc, this.tok, s.att]),
972
+ attOut: this.bind('attnOutW', [a.sc, this.vc[l], a.att, this.tok, s.att]),
973
+ o: this.bind('matmulQ4', [a.att, p('attn_output.weight').q, p('attn_output.weight').s, a.o, s.mmO]),
974
+ postA: this.bind('rmsnormAdd', [a.o, p('post_attention_norm.weight').raw, a.x, u.rmsH]),
975
+ rmsF: this.bind('rmsnorm', [a.x, p('ffn_norm.weight').raw, a.xn, u.rmsH]),
976
+ mg: this.bindR('matmulQ4', [a.xn, p('ffn_gate.weight').q, p('ffn_gate.weight').s, { buffer: a.mgu, offset: 0, size: gateBytes }, u.mmG]),
977
+ mu: this.bindR('matmulQ4', [a.xn, p('ffn_up.weight').q, p('ffn_up.weight').s, { buffer: a.mgu, offset: gateBytes, size: gateBytes }, u.mmG]),
978
+ mact: this.bind('geglu', [a.mgu, a.ma, u.gegluD]),
979
+ md: this.bind('matmulQ4', [a.ma, p('ffn_down.weight').q, p('ffn_down.weight').s, a.md, u.mmD]),
980
+ mpost: this.bind('rmsnorm', [a.md, p('post_ffw_norm_1.weight').raw, a.mlp, u.rmsH]),
981
+ xm: this.bind('rmsnorm', [a.x, p('pre_ffw_norm_2.weight').raw, a.xm, u.rmsH]),
982
+ rp: this.bind('rmsnormK', [a.x, p('ffn_gate_inp.scale').raw, a.rt, u.routerK]),
983
+ router: this.bind('matmulF32', [a.rt, p('ffn_gate_inp.weight').raw, a.rl, u.router]),
984
+ topk: this.bind('topk', [a.rl, a.sel, u.topk]),
985
+ // Routers of layers l+1..l+lookahead on this layer's post-attention state: the prefetch guesses.
986
+ pf: Array.from({ length: Math.min(this.lookahead, Lc - 1 - l) }, (_, i) => {
987
+ const at = (n) => W[`blk.${l + 1 + i}.${n}`];
988
+ return {
989
+ rp: this.bind('rmsnormK', [a.x, at('ffn_gate_inp.scale').raw, a.prt, u.routerK]),
990
+ router: this.bind('matmulF32', [a.prt, at('ffn_gate_inp.weight').raw, a.prl, u.router]),
991
+ topk: this.bindR('topk', [a.prl, { buffer: a.psel, offset: 256 * i, size: 8 * K }, u.topk]),
992
+ };
993
+ }),
994
+ gu: this.bind('expertQ4', [a.xm, this.pool.guQ, this.pool.guS, a.gu, u.gu, a.slots]),
995
+ act: this.bind('geglu', [a.gu, a.act, u.gegluE]),
996
+ dn: this.bind('expertQ4', [a.act, this.pool.dQ, this.pool.dS, a.dn, u.dn, a.slots]),
997
+ sum: this.bind('moeSum', [a.dn, a.sel, p('ffn_down_exps.scale').raw, a.moe, u.sum]),
998
+ mpost2: this.bind('rmsnorm', [a.moe, p('post_ffw_norm_2.weight').raw, a.moe2, u.rmsH]),
999
+ post: this.bind('postFfn', [a.mlp, a.moe2, p('post_ffw_norm.weight').raw, p('layer_output_scale.weight').raw, a.x, u.rmsH]),
1000
+ };
1001
+ if (a.cap) {
1002
+ ly.capA = this.bind('copyAt', [a.x, a.cap, this.uniform([H, l * H])]);
1003
+ ly.capX = this.bind('copyAt', [a.x, a.cap, this.uniform([H, (Lc + l) * H])]);
1004
+ }
1005
+ this.layers.push(ly);
1006
+ }
1007
+ this.g.rmsOut = this.bind('rmsnorm', [a.x, W['output_norm.weight'].raw, a.xn, u.rmsH]);
1008
+ this.g.lm = this.bind('matmulQ6K', [...q6, a.xn, a.logits, u.lm]);
1009
+ this.g.cap = this.bind('softcap', [a.logits, u.cap]);
1010
+ this.g.am = this.bind('argmax', [a.logits, a.am, u.am]);
1011
+ // A chunk pins min(experts, T·K) experts per layer at once; keep that under half the pool.
1012
+ if (Math.min(c.experts, this.chunkTokens * c.topK) > this.poolSlots / 2) this.chunkTokens = Math.max(2, Math.floor(this.poolSlots / 2 / c.topK));
1013
+ if (this.batchPrefill) this.initBatch();
1014
+ }
1015
+
1016
+ // Buffers and bind groups for chunked prefill (rows = the chunk's tokens).
1017
+ initBatch() {
1018
+ const c = this.cfg, B = this.chunkTokens, W = this.w, H = c.hidden, K = c.topK, E = c.experts, F = c.expertFf, DF = c.denseFf;
1019
+ const QN = Math.max(...c.headDim.map((d) => c.heads * d)), KVN = Math.max(...c.headDim.map((d, l) => c.kvHeads[l] * d)), HD = Math.max(...c.headDim);
1020
+ const A = (n) => this.buffer(n * 4, STORAGE | COPY_SRC | COPY_DST, 'batch');
1021
+ const b = this.b = {
1022
+ ids: A(B), x: A(B * H), xn: A(B * H), q: A(B * QN), k: A(B * KVN), v: A(B * KVN), vr: A(B * KVN), att: A(B * QN), o: A(B * H),
1023
+ sc: A(c.heads * B * this.maxCtx), mg: A(B * DF), mu: A(B * DF), ma: A(B * DF), md: A(B * H), mlp: A(B * H), xm: A(B * H), rt: A(B * H),
1024
+ rl: A(B * E), sel: A(B * 2 * K), slots: A(B * K), gu: A(B * K * 2 * F), act: A(B * K * F), dn: A(B * K * H), moe: A(B * H), moe2: A(B * H),
1025
+ ropeS: A(B * HD), ropeF: A(B * HD),
1026
+ };
1027
+ this.qb = this.buffer(16, 0x40 | COPY_DST, 'batch');
1028
+ this.rbSelB = this.buffer(B * 2 * K * 4, MAP_READ | COPY_DST, 'batch');
1029
+ const U = (w) => this.uniform(w), eps = { f: c.eps }, qb = this.qb;
1030
+ const u = {
1031
+ embed: U([H, H / 256, { f: Math.fround(Math.sqrt(H)) }]), rmsH: U([H, eps]),
1032
+ routerK: U([H, eps, { f: Math.fround(1 / Math.fround(Math.sqrt(H))) }]), router: U([E, H]), topk: U([E, K]),
1033
+ mg: U([DF, H]), md: U([H, DF]), geglu: U([DF]), gu: U([2 * F, H, 2 * F, 0, K]), dn: U([H, F, H, 1, K]), gegluE: U([F, B * K]), sum: U([H, K]),
1034
+ };
1035
+ const shape = new Map();
1036
+ const shapeOf = (l) => {
1037
+ const hd = c.headDim[l], qn = c.heads * hd, kvn = c.kvHeads[l] * hd, key = `${hd}/${c.kvHeads[l]}/${c.swa[l]}`;
1038
+ if (!shape.has(key)) shape.set(key, {
1039
+ q: U([qn, H]), kv: U([kvn, H]), o: U([H, qn]), ropeQ: U([hd, qn, eps]), ropeK: U([hd, kvn, eps]), kvn: U([kvn]),
1040
+ att: U([c.heads, c.kvHeads[l], hd, c.swa[l] ? c.window : 0]),
1041
+ });
1042
+ return shape.get(key);
1043
+ };
1044
+ const q4 = (name, x, y, uni) => this.bind('matmulQ4T', [x, W[name].q, W[name].s, y, uni, qb]);
1045
+ this.gB = { embed: this.bind('embedQ6KB', [this.q6.ql, this.q6.qh, this.q6.sc, this.q6.dd, b.ids, b.x, u.embed, qb]) };
1046
+ for (let l = 0; l < c.layers; l++) {
1047
+ const n = (t) => `blk.${l}.${t}`, raw = (t) => W[n(t)].raw, sh = shapeOf(l), hasV = !!W[n('attn_v.weight')];
1048
+ const rope = c.swa[l] ? b.ropeS : b.ropeF;
1049
+ this.layers[l].B = {
1050
+ rmsA: this.bind('rmsnormB', [b.x, raw('attn_norm.weight'), b.xn, u.rmsH]),
1051
+ q: q4(n('attn_q.weight'), b.xn, b.q, sh.q), k: q4(n('attn_k.weight'), b.xn, b.k, sh.kv),
1052
+ v: hasV ? q4(n('attn_v.weight'), b.xn, b.vr, sh.kv) : null,
1053
+ vn: this.bind('vNormB', [hasV ? b.vr : b.k, b.v, sh.ropeK]),
1054
+ ropeQ: this.bind('qkNormRopeB', [b.q, raw('attn_q_norm.weight'), rope, sh.ropeQ]),
1055
+ ropeK: this.bind('qkNormRopeB', [b.k, raw('attn_k_norm.weight'), rope, sh.ropeK]),
1056
+ kv: this.bind('kvStoreB', [b.k, b.v, this.kc[l], this.vc[l], sh.kvn, qb]),
1057
+ score: this.bind('attnScoreWB', [b.q, this.kc[l], b.sc, sh.att, qb]),
1058
+ soft: this.bind('softmaxB', [b.sc, qb]),
1059
+ attOut: this.bind('attnOutB', [b.sc, this.vc[l], b.att, sh.att, qb]),
1060
+ o: q4(n('attn_output.weight'), b.att, b.o, sh.o),
1061
+ postA: this.bind('rmsnormAddB', [b.o, raw('post_attention_norm.weight'), b.x, u.rmsH]),
1062
+ rmsF: this.bind('rmsnormB', [b.x, raw('ffn_norm.weight'), b.xn, u.rmsH]),
1063
+ mg: q4(n('ffn_gate.weight'), b.xn, b.mg, u.mg), mu: q4(n('ffn_up.weight'), b.xn, b.mu, u.mg),
1064
+ mact: this.bind('gegluSplitB', [b.mg, b.mu, b.ma, u.geglu, qb]),
1065
+ md: q4(n('ffn_down.weight'), b.ma, b.md, u.md),
1066
+ mpost: this.bind('rmsnormB', [b.md, raw('post_ffw_norm_1.weight'), b.mlp, u.rmsH]),
1067
+ xm: this.bind('rmsnormB', [b.x, raw('pre_ffw_norm_2.weight'), b.xm, u.rmsH]),
1068
+ rp: this.bind('rmsnormKB', [b.x, raw('ffn_gate_inp.scale'), b.rt, u.routerK]),
1069
+ router: this.bind('matmulF32T', [b.rt, raw('ffn_gate_inp.weight'), b.rl, u.router]),
1070
+ topk: this.bind('topkB', [b.rl, b.sel, u.topk]),
1071
+ gu: this.bind('expertQ4B', [b.xm, this.pool.guQ, this.pool.guS, b.gu, u.gu, b.slots]),
1072
+ act: this.bind('geglu', [b.gu, b.act, u.gegluE]),
1073
+ dn: this.bind('expertQ4B', [b.act, this.pool.dQ, this.pool.dS, b.dn, u.dn, b.slots]),
1074
+ sum: this.bind('moeSumB', [b.dn, b.sel, raw('ffn_down_exps.scale'), b.moe, u.sum, qb]),
1075
+ mpost2: this.bind('rmsnormB', [b.moe, raw('post_ffw_norm_2.weight'), b.moe2, u.rmsH]),
1076
+ post: this.bind('postFfnB', [b.mlp, b.moe2, raw('post_ffw_norm.weight'), raw('layer_output_scale.weight'), b.x, u.rmsH]),
1077
+ };
1078
+ }
1079
+ }
1080
+
1081
+ // Layer l over the chunk's T tokens up to its routing (the batched twin of encodeAttention).
1082
+ encodeTrunkB(pass, l, T, S) {
1083
+ const c = this.cfg, g = this.layers[l].B, d = (nm, gr, x, y, z) => this.dispatch(pass, nm, gr, x, y, z);
1084
+ const H = c.hidden, hd = c.headDim[l], QN = c.heads * hd, kvn = c.kvHeads[l] * hd, DF = c.denseFf;
1085
+ const tt = Math.ceil(T / 16), mm = (n) => Math.ceil(n / 16);
1086
+ d('rmsnormB', g.rmsA, T);
1087
+ d('matmulQ4T', g.q, mm(QN), tt); d('matmulQ4T', g.k, mm(kvn), tt);
1088
+ if (g.v) d('matmulQ4T', g.v, mm(kvn), tt);
1089
+ d('vNormB', g.vn, c.kvHeads[l], T);
1090
+ d('qkNormRopeB', g.ropeQ, c.heads, T); d('qkNormRopeB', g.ropeK, c.kvHeads[l], T);
1091
+ d('kvStoreB', g.kv, Math.ceil(kvn / 256), T);
1092
+ d('attnScoreWB', g.score, Math.ceil(c.heads * T * S / 256));
1093
+ d('softmaxB', g.soft, c.heads * T);
1094
+ d('attnOutB', g.attOut, Math.ceil(QN / 256), T);
1095
+ d('matmulQ4T', g.o, mm(H), tt);
1096
+ d('rmsnormAddB', g.postA, T);
1097
+ d('rmsnormB', g.rmsF, T);
1098
+ d('matmulQ4T', g.mg, mm(DF), tt); d('matmulQ4T', g.mu, mm(DF), tt);
1099
+ d('gegluSplitB', g.mact, Math.ceil(DF / 256), T);
1100
+ d('matmulQ4T', g.md, mm(H), tt);
1101
+ d('rmsnormB', g.mpost, T);
1102
+ d('rmsnormB', g.xm, T);
1103
+ d('rmsnormKB', g.rp, T);
1104
+ d('matmulF32T', g.router, c.experts, T);
1105
+ d('topkB', g.topk, T);
1106
+ }
1107
+ encodeExpertsB(pass, l, T) {
1108
+ const c = this.cfg, g = this.layers[l].B, d = (nm, gr, x, y, z) => this.dispatch(pass, nm, gr, x, y, z);
1109
+ d('expertQ4B', g.gu, Math.ceil(2 * c.expertFf / 4), 1, T * c.topK);
1110
+ d('geglu', g.act, Math.ceil(T * c.topK * c.expertFf / 256));
1111
+ d('expertQ4B', g.dn, Math.ceil(c.hidden / 4), 1, T * c.topK);
1112
+ d('moeSumB', g.sum, Math.ceil(c.hidden / 256), T);
1113
+ d('rmsnormB', g.mpost2, T);
1114
+ d('postFfnB', g.post, T);
1115
+ }
1116
+
1117
+ // Feeds T = ids.length tokens at this.position through every layer: one routing readback per layer,
1118
+ // the union of the chunk's experts loaded once. want as in step(), for the last token.
1119
+ async prefillChunk(ids, want = 'none') {
1120
+ const c = this.cfg, dev = this.device, K = c.topK, H = c.hidden, T = ids.length, pos0 = this.position, S = pos0 + T;
1121
+ if (T > this.chunkTokens) throw new Error(`chunk of ${T} > ${this.chunkTokens}`);
1122
+ if (S > this.maxCtx) throw new Error(`context full (${this.maxCtx} tokens)`);
1123
+ const tStart = performance.now();
1124
+ dev.queue.writeBuffer(this.b.ids, 0, Uint32Array.from(ids));
1125
+ dev.queue.writeBuffer(this.qb, 0, new Uint32Array([T, pos0, S, 0]));
1126
+ const ls = c.swa.indexOf(true), lf = c.swa.indexOf(false);
1127
+ const tables = (l, ff, buf) => {
1128
+ if (l < 0) return;
1129
+ const d = c.headDim[l], r = new Float32Array(T * d);
1130
+ for (let t = 0; t < T; t++) r.set(this.ropeTable(d, c.ropeTheta[l], ff, pos0 + t), t * d);
1131
+ dev.queue.writeBuffer(buf, 0, r);
1132
+ };
1133
+ tables(ls, null, this.b.ropeS); tables(lf, this.ropeFreqs, this.b.ropeF);
1134
+ const selBytes = T * 2 * K * 4;
1135
+ let enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
1136
+ this.dispatch(pass, 'embedQ6KB', this.gB.embed, Math.ceil(H / 256), T);
1137
+ this.encodeTrunkB(pass, 0, T, S);
1138
+ pass.end();
1139
+ enc.copyBufferToBuffer(this.b.sel, 0, this.rbSelB, 0, selBytes);
1140
+ dev.queue.submit([enc.finish()]);
1141
+ for (let l = 0; l < c.layers; l++) {
1142
+ const t0 = performance.now();
1143
+ await this.rbSelB.mapAsync(1, 0, selBytes);
1144
+ this.counters.gpuWaitMs += performance.now() - t0;
1145
+ const sel = new Uint32Array(this.rbSelB.getMappedRange(0, selBytes).slice(0));
1146
+ this.rbSelB.unmap();
1147
+ const pairs = new Uint32Array(T * K);
1148
+ for (let t = 0; t < T; t++) for (let j = 0; j < K; j++) pairs[t * K + j] = sel[t * 2 * K + j];
1149
+ const unique = [...new Set(pairs)];
1150
+ const tq = performance.now();
1151
+ const slotsU = await this.xs.ensure(l, unique);
1152
+ this.counters.ensureMs += performance.now() - tq;
1153
+ const slotOf = new Map(unique.map((e, i) => [e, slotsU[i]]));
1154
+ dev.queue.writeBuffer(this.b.slots, 0, pairs.map((e) => slotOf.get(e)));
1155
+ enc = dev.createCommandEncoder(); pass = enc.beginComputePass();
1156
+ this.encodeExpertsB(pass, l, T);
1157
+ if (l + 1 < c.layers) {
1158
+ this.encodeTrunkB(pass, l + 1, T, S);
1159
+ pass.end();
1160
+ enc.copyBufferToBuffer(this.b.sel, 0, this.rbSelB, 0, selBytes);
1161
+ } else {
1162
+ pass.end();
1163
+ enc.copyBufferToBuffer(this.b.x, (T - 1) * H * 4, this.a.x, 0, H * 4);
1164
+ if (want !== 'none') {
1165
+ pass = enc.beginComputePass();
1166
+ this.encodeHead(pass, want);
1167
+ pass.end();
1168
+ if (want === 'argmax') enc.copyBufferToBuffer(this.a.am, 0, this.rbArg, 0, 4);
1169
+ if (want === 'logits') enc.copyBufferToBuffer(this.a.logits, 0, this.rbLogits, 0, c.vocab * 4);
1170
+ }
1171
+ }
1172
+ dev.queue.submit([enc.finish()]);
1173
+ this.xs.release(slotsU);
1174
+ }
1175
+ this.position += T;
1176
+ this.cached.push(...ids);
1177
+ let result = null;
1178
+ if (want === 'argmax') {
1179
+ await this.rbArg.mapAsync(1, 0, 4);
1180
+ result = new Uint32Array(this.rbArg.getMappedRange(0, 4))[0];
1181
+ this.rbArg.unmap();
1182
+ } else if (want === 'logits') {
1183
+ await this.rbLogits.mapAsync(1);
1184
+ result = new Float32Array(this.rbLogits.getMappedRange().slice(0));
1185
+ this.rbLogits.unmap();
1186
+ } else {
1187
+ await dev.queue.onSubmittedWorkDone();
1188
+ }
1189
+ this.counters.tokens += T;
1190
+ this.counters.wallMs += performance.now() - tStart;
1191
+ return result;
1192
+ }
1193
+
1194
+ // Feeds `ids` after the cached prefix they share; KV past the prefix is simply overwritten (no recurrent
1195
+ // state to rewind). Chunks of up to chunkTokens when there are at least two new tokens.
1196
+ async prefill(ids, want = 'argmax') {
1197
+ let common = 0;
1198
+ while (common < ids.length - 1 && common < this.cached.length && this.cached[common] === ids[common]) common++;
1199
+ if (common < this.cached.length) { this.position = common; this.cached.length = common; }
1200
+ const rest = ids.slice(common);
1201
+ if (!this.batchPrefill || rest.length < 2) return super.prefill(ids, want);
1202
+ let r = null;
1203
+ for (let i = 0; i < rest.length; i += this.chunkTokens) {
1204
+ r = await this.prefillChunk(rest.slice(i, i + this.chunkTokens), i + this.chunkTokens >= rest.length ? want : 'none');
1205
+ }
1206
+ return r;
1207
+ }
1208
+
1209
+ encodeEmbed(pass) { this.dispatch(pass, 'embedQ6K', this.g.embed, Math.ceil(this.cfg.hidden / 256)); }
1210
+ encodeHead(pass, want) {
1211
+ const c = this.cfg;
1212
+ this.dispatch(pass, 'rmsnorm', this.g.rmsOut, 1);
1213
+ this.dispatch(pass, 'matmulQ6K', this.g.lm, Math.ceil(c.vocab / 4));
1214
+ if (c.softcap) this.dispatch(pass, 'softcap', this.g.cap, Math.ceil(c.vocab / 256));
1215
+ if (want === 'argmax') this.dispatch(pass, 'argmax', this.g.am, 1);
1216
+ }
1217
+
1218
+ // Attention, the dense MLP and the router of layer l (all before the routing readback).
1219
+ encodeAttention(pass, l, seqLen) {
1220
+ const c = this.cfg, g = this.layers[l], d = (n, gr, x, y, z) => this.dispatch(pass, n, gr, x, y, z);
1221
+ const hd = c.headDim[l], QN = c.heads * hd, kvn = c.kvHeads[l] * hd;
1222
+ const n = c.swa[l] ? Math.min(seqLen, c.window) : seqLen;
1223
+ d('rmsnorm', g.rmsA, 1);
1224
+ d('matmulQ4', g.q, Math.ceil(QN / 4)); d('matmulQ4', g.k, Math.ceil(kvn / 4));
1225
+ if (g.hasV) d('matmulQ4', g.v, Math.ceil(kvn / 4));
1226
+ d('vNorm', g.vn, c.kvHeads[l]);
1227
+ d('qkNormRope', g.ropeQ, c.heads); d('qkNormRope', g.ropeK, c.kvHeads[l]);
1228
+ d('kvStore', g.kv, Math.ceil(kvn / 256));
1229
+ d('attnScoreW', g.score, Math.ceil(c.heads * n / 256));
1230
+ d('softmaxW', g.soft, c.heads);
1231
+ d('attnOutW', g.attOut, Math.ceil(QN / 256));
1232
+ d('matmulQ4', g.o, Math.ceil(c.hidden / 4));
1233
+ d('rmsnormAdd', g.postA, 1);
1234
+ if (g.capA) d('copyAt', g.capA, Math.ceil(c.hidden / 256));
1235
+ d('rmsnorm', g.rmsF, 1);
1236
+ d('matmulQ4', g.mg, Math.ceil(c.denseFf / 4)); d('matmulQ4', g.mu, Math.ceil(c.denseFf / 4));
1237
+ d('geglu', g.mact, Math.ceil(c.denseFf / 256));
1238
+ d('matmulQ4', g.md, Math.ceil(c.hidden / 4));
1239
+ d('rmsnorm', g.mpost, 1);
1240
+ d('rmsnorm', g.xm, 1);
1241
+ if (this.prefetch) for (const p of g.pf) { d('rmsnormK', p.rp, 1); d('matmulF32', p.router, c.experts); d('topk', p.topk, 1); }
1242
+ d('rmsnormK', g.rp, 1);
1243
+ d('matmulF32', g.router, c.experts);
1244
+ d('topk', g.topk, 1);
1245
+ }
1246
+ encodeExperts(pass, l) {
1247
+ const c = this.cfg, g = this.layers[l], d = (n, gr, x, y, z) => this.dispatch(pass, n, gr, x, y, z);
1248
+ d('expertQ4', g.gu, Math.ceil(2 * c.expertFf / 4), 1, c.topK);
1249
+ d('geglu', g.act, Math.ceil(c.topK * c.expertFf / 256));
1250
+ d('expertQ4', g.dn, Math.ceil(c.hidden / 4), 1, c.topK);
1251
+ d('moeSum', g.sum, Math.ceil(c.hidden / 256));
1252
+ d('rmsnorm', g.mpost2, 1);
1253
+ d('postFfn', g.post, 1);
1254
+ if (g.capX) d('copyAt', g.capX, Math.ceil(c.hidden / 256));
1255
+ }
1256
+
1257
+ // NeoX RoPE angles as llama.cpp's Metal kernel forms them: f32 pos × base^(−2i/d), ÷ the frequency factor
1258
+ // on full-attention layers. One table per layer kind; [0, d/2) = cos, [d/2, d) = sin.
1259
+ // KV cache bytes: per-layer head counts and sizes (sliding and full-attention layers differ).
1260
+ kvBytes() { const c = this.cfg; return c.headDim.reduce((sum, d, l) => sum + 2 * 2 * this.maxCtx * c.kvHeads[l] * d, 0); }
1261
+ ropeTable(d, base, ff, pos) {
1262
+ const half = d / 2, r = new Float32Array(d), inv = Math.fround(-1 / d);
1263
+ for (let i = 0; i < half; i++) {
1264
+ let th = Math.fround(pos * Math.fround(Math.pow(base, Math.fround(inv * 2 * i))));
1265
+ if (ff) th = Math.fround(th / ff[i]);
1266
+ r[i] = Math.cos(th); r[half + i] = Math.sin(th);
1267
+ }
1268
+ return r;
1269
+ }
1270
+ writeTokenUniforms(token, pos) {
1271
+ const c = this.cfg;
1272
+ this.device.queue.writeBuffer(this.tok, 0, new Uint32Array([token, pos, pos + 1, 0]));
1273
+ const ls = c.swa.indexOf(true), lf = c.swa.indexOf(false);
1274
+ if (ls >= 0) this.device.queue.writeBuffer(this.a.ropeS, 0, this.ropeTable(c.headDim[ls], c.ropeTheta[ls], null, pos));
1275
+ if (lf >= 0) this.device.queue.writeBuffer(this.a.ropeF, 0, this.ropeTable(c.headDim[lf], c.ropeTheta[lf], this.ropeFreqs, pos));
1276
+ }
1277
+
1278
+ // Debug: the last step's per-layer states — attn_out (after the attention residual) and each layer's output.
1279
+ async readCapture() {
1280
+ const c = this.cfg, n = 2 * c.layers * c.hidden * 4;
1281
+ const rb = this.device.createBuffer({ size: n, usage: MAP_READ | COPY_DST });
1282
+ const enc = this.device.createCommandEncoder();
1283
+ enc.copyBufferToBuffer(this.a.cap, 0, rb, 0, n);
1284
+ this.device.queue.submit([enc.finish()]);
1285
+ await rb.mapAsync(1);
1286
+ const all = new Float32Array(rb.getMappedRange().slice(0));
1287
+ rb.unmap(); rb.destroy();
1288
+ const H = c.hidden, L = c.layers;
1289
+ return { attnOut: Array.from({ length: L }, (_, l) => all.subarray(l * H, (l + 1) * H)), layers: Array.from({ length: L }, (_, l) => all.subarray((L + l) * H, (L + l + 1) * H)) };
1290
+ }
1291
+ }
engines/qwen35_moe_ssd.js ADDED
@@ -0,0 +1,1231 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* qwen35_moe_ssd.js — Qwen3.5/3.6 MoE (Qwen3.6-35B-A3B, GGUF arch `qwen35moe`) in a browser tab,
2
+ * larger than RAM: rung 2b of the SSD-streaming ladder. The routed experts stream from OPFS into
3
+ * a GPU slot pool exactly as in rung 2a (qwen3_moe_ssd.js, which this module extends); what is new
4
+ * is the hybrid trunk:
5
+ * - 3 of every 4 layers are Gated DeltaNet (linear attention): a 4-tap causal conv over the
6
+ * q|k|v projection, L2-normalised q and k, and the gated delta rule on a 128×128 state per
7
+ * value head; output = RMSNorm(o)·w · SiLU(z), then a projection.
8
+ * - every 4th layer is full attention with a per-head output gate (the q projection carries
9
+ * q and gate per head), q/k RMSNorm and NeoX RoPE on the first 64 of 256 dims.
10
+ * - each MoE block adds a shared expert, scaled by sigmoid(x·w_gate).
11
+ * The math follows llama.cpp's qwen35moe graph and its CPU ops (ggml_gated_delta_net, ssm_conv,
12
+ * rope_multi with text positions), so greedy output can be gated against llama-server on the
13
+ * same GGUF. Decode is one token at a time (prefill feeds tokens through the same step).
14
+ *
15
+ * const m = await Qwen35MoeSsd.load(null, { url: '…/Qwen3.6-35B-A3B-Q8_0.gguf', poolBytes });
16
+ */
17
+
18
+ import { Qwen3MoeSsd, STORAGE, COPY_SRC, COPY_DST, MAP_READ, UNIFORM, I8 } from './qwen3_moe_ssd.js';
19
+
20
+ export const QWEN36_35B_A3B = {
21
+ repo: 'unsloth/Qwen3.6-35B-A3B-GGUF',
22
+ file: 'Qwen3.6-35B-A3B-Q8_0.gguf',
23
+ revision: 'a483e9e6cbd595906af30beda3187c2663a1118c',
24
+ sha256: 'd1a395809f65a43a13ad119eb4e7acdef1ac6d68120f39902c8ab96e72794a59',
25
+ size: 36903140320,
26
+ };
27
+
28
+ const SIGMOID = 'fn sigm(x: f32) -> f32 { return 1.0 / (1.0 + exp(-x)); }';
29
+
30
+ const KERNELS = {
31
+ // Per-head RMSNorm (QK-norm), then NeoX RoPE on the first nRot dims (pairs i, i + nRot/2);
32
+ // dims ≥ nRot keep the normed value. src rows are srcStride apart (the q projection holds
33
+ // [q hd | gate hd] per head); with gate = 1 the gate half is copied to gout. One workgroup per head.
34
+ qkNormRopeP: `
35
+ struct P { hd: u32, srcStride: u32, nRot: u32, gate: u32, eps: f32 }
36
+ @group(0) @binding(0) var<storage, read> src: array<f32>;
37
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
38
+ @group(0) @binding(2) var<storage, read> rope: array<f32>;
39
+ @group(0) @binding(3) var<storage, read_write> dst: array<f32>;
40
+ @group(0) @binding(4) var<storage, read_write> gout: array<f32>;
41
+ @group(0) @binding(5) var<uniform> p: P;
42
+ var<workgroup> red: array<f32, 64>;
43
+ @compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
44
+ let sb = wg.x * p.srcStride; let ob = wg.x * p.hd; let t = l.x; let half = p.nRot / 2u;
45
+ var acc = 0.0;
46
+ for (var i = t; i < p.hd; i += 64u) { let v = src[sb + i]; acc += v * v; }
47
+ red[t] = acc; workgroupBarrier();
48
+ for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
49
+ let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
50
+ for (var i = t; i < p.hd; i += 64u) {
51
+ if (i < half) {
52
+ let a = (src[sb + i] * scale) * w[i];
53
+ let b = (src[sb + i + half] * scale) * w[i + half];
54
+ let c = rope[i]; let sn = rope[half + i];
55
+ dst[ob + i] = a * c - b * sn;
56
+ dst[ob + i + half] = a * sn + b * c;
57
+ } else if (i >= p.nRot) {
58
+ dst[ob + i] = (src[sb + i] * scale) * w[i];
59
+ }
60
+ if (p.gate == 1u) { gout[ob + i] = src[sb + p.hd + i]; }
61
+ }
62
+ }`,
63
+ // x[i] *= sigmoid(g[i]) — the full-attention output gate.
64
+ sigmoidMul: `
65
+ struct P { n: u32 }
66
+ @group(0) @binding(0) var<storage, read_write> x: array<f32>;
67
+ @group(0) @binding(1) var<storage, read> g: array<f32>;
68
+ @group(0) @binding(2) var<uniform> p: P;
69
+ ${SIGMOID}
70
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
71
+ let i = gi.x; if (i < p.n) { x[i] = x[i] * sigm(g[i]); }
72
+ }`,
73
+ // Causal depthwise conv, kernel 4, over the q|k|v projection, then SiLU. state holds each
74
+ // channel's previous 3 inputs, oldest first (llama.cpp's conv state), and rolls forward.
75
+ gdnConv: `
76
+ struct P { n: u32 }
77
+ @group(0) @binding(0) var<storage, read> mixed: array<f32>;
78
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
79
+ @group(0) @binding(2) var<storage, read_write> st: array<f32>;
80
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
81
+ @group(0) @binding(4) var<uniform> p: P;
82
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
83
+ let c = gi.x; if (c >= p.n) { return; }
84
+ let s0 = st[c * 3u]; let s1 = st[c * 3u + 1u]; let s2 = st[c * 3u + 2u]; let x = mixed[c];
85
+ var sum = 0.0;
86
+ sum += s0 * w[c * 4u]; sum += s1 * w[c * 4u + 1u]; sum += s2 * w[c * 4u + 2u]; sum += x * w[c * 4u + 3u];
87
+ y[c] = sum / (1.0 + exp(-sum));
88
+ st[c * 3u] = s1; st[c * 3u + 1u] = s2; st[c * 3u + 2u] = x;
89
+ }`,
90
+ // The gated delta rule for one token. Workgroup h = value head, thread j = value row j of the
91
+ // head's state M (M[j][i] = S[i][j], llama.cpp's layout; i runs over the key dim). k-head = h % kH.
92
+ // q, k ← L2-normalised (x / sqrt(Σx² + eps)); beta = σ(b); g = softplus(a + dt) · A
93
+ // M ← M·e^g; d_j = β (v_j − Σ_i M[j][i] k_i); M[j][i] += k_i d_j; o_j = (Σ_i M[j][i] q_i) / √dS
94
+ gdnStep: `
95
+ struct P { kH: u32, dS: u32, keyDim: u32, pad: u32, eps: f32, scale: f32 }
96
+ @group(0) @binding(0) var<storage, read> cv: array<f32>;
97
+ @group(0) @binding(1) var<storage, read> bb: array<f32>;
98
+ @group(0) @binding(2) var<storage, read> al: array<f32>;
99
+ @group(0) @binding(3) var<storage, read> dt: array<f32>;
100
+ @group(0) @binding(4) var<storage, read> aa: array<f32>;
101
+ @group(0) @binding(5) var<storage, read_write> M: array<f32>;
102
+ @group(0) @binding(6) var<storage, read_write> o: array<f32>;
103
+ @group(0) @binding(7) var<uniform> p: P;
104
+ ${SIGMOID}
105
+ var<workgroup> kq: array<f32, 128>;
106
+ var<workgroup> qq: array<f32, 128>;
107
+ var<workgroup> rq: array<f32, 128>;
108
+ var<workgroup> rk: array<f32, 128>;
109
+ fn log1p_(u: f32) -> f32 { let y = 1.0 + u; return log(y) - ((y - 1.0) - u) / y; }
110
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
111
+ let h = wg.x; let j = l.x; let dS = p.dS; let kh = h % p.kH;
112
+ let qv = cv[kh * dS + j]; let kv = cv[p.keyDim + kh * dS + j]; let vv = cv[2u * p.keyDim + h * dS + j];
113
+ rq[j] = qv * qv; rk[j] = kv * kv; workgroupBarrier();
114
+ for (var s = 64u; s > 0u; s >>= 1u) { if (j < s) { rq[j] += rq[j + s]; rk[j] += rk[j + s]; } workgroupBarrier(); }
115
+ let n = f32(dS); let inv = 1.0 / sqrt(n);
116
+ qq[j] = (qv * (1.0 / sqrt(rq[0] / n + p.eps / n))) * inv;
117
+ kq[j] = (kv * (1.0 / sqrt(rk[0] / n + p.eps / n))) * inv;
118
+ workgroupBarrier();
119
+ let beta = sigm(bb[h]);
120
+ let x = al[h] + dt[h];
121
+ let sp = select(log1p_(exp(x)), x, x > 20.0);
122
+ let decay = exp(sp * aa[h]);
123
+ let row = (h * dS + j) * dS;
124
+ var sk = 0.0;
125
+ for (var i = 0u; i < dS; i++) { let m = M[row + i] * decay; M[row + i] = m; sk += m * kq[i]; }
126
+ let d = (vv - sk) * beta;
127
+ var acc = 0.0;
128
+ for (var i = 0u; i < dS; i++) { let m = M[row + i] + kq[i] * d; M[row + i] = m; acc += m * qq[i]; }
129
+ o[h * dS + j] = acc * p.scale;
130
+ }`,
131
+ // Per value head: RMSNorm(o)·w · SiLU(z).
132
+ gdnNormGate: `
133
+ struct P { dS: u32, eps: f32 }
134
+ @group(0) @binding(0) var<storage, read> o: array<f32>;
135
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
136
+ @group(0) @binding(2) var<storage, read> z: array<f32>;
137
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
138
+ @group(0) @binding(4) var<uniform> p: P;
139
+ var<workgroup> red: array<f32, 128>;
140
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
141
+ let h = wg.x; let j = l.x; let i = h * p.dS + j;
142
+ let v = o[i]; red[j] = v * v; workgroupBarrier();
143
+ for (var s = 64u; s > 0u; s >>= 1u) { if (j < s) { red[j] += red[j + s]; } workgroupBarrier(); }
144
+ let scale = 1.0 / sqrt(red[0] / f32(p.dS) + p.eps);
145
+ let zz = z[i];
146
+ y[i] = ((v * scale) * w[j]) * (zz / (1.0 + exp(-zz)));
147
+ }`,
148
+ // x[h] = (Σ_j w_j·out[j·H + h] + sh[h]·σ(sg)) + x[h]: routed experts, the gated shared expert,
149
+ // then the residual — llama.cpp's order.
150
+ moeAccumShared: `
151
+ struct P { H: u32, k: u32 }
152
+ @group(0) @binding(0) var<storage, read> out: array<f32>;
153
+ @group(0) @binding(1) var<storage, read> sel: array<u32>;
154
+ @group(0) @binding(2) var<storage, read> sh: array<f32>;
155
+ @group(0) @binding(3) var<storage, read> sg: array<f32>;
156
+ @group(0) @binding(4) var<storage, read_write> x: array<f32>;
157
+ @group(0) @binding(5) var<uniform> p: P;
158
+ ${SIGMOID}
159
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
160
+ let h = gi.x; if (h >= p.H) { return; }
161
+ var acc = 0.0;
162
+ for (var j = 0u; j < p.k; j++) { acc += bitcast<f32>(sel[p.k + j]) * out[j * p.H + h]; }
163
+ x[h] = (acc + sh[h] * sigm(sg[0])) + x[h];
164
+ }`,
165
+ };
166
+
167
+
168
+ // ── Batched prefill: a chunk of T prompt tokens per layer ───────────────────────────────────
169
+ // Every batched buffer is token-major (row t holds token pos0 + t). Q is the per-chunk uniform.
170
+ const QB = 'struct Q { T: u32, pos0: u32, S: u32, pad: u32 }';
171
+ const BATCH_KERNELS = {
172
+ embedQ8B: `enable f16;
173
+ struct P { n: u32 }
174
+ ${QB}
175
+ @group(0) @binding(0) var<storage, read> q: array<u32>;
176
+ @group(0) @binding(1) var<storage, read> s: array<f16>;
177
+ @group(0) @binding(2) var<storage, read> ids: array<u32>;
178
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
179
+ @group(0) @binding(4) var<uniform> p: P;
180
+ @group(0) @binding(5) var<uniform> qd: Q;
181
+ ${I8}
182
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
183
+ let i = g.x; let t = g.y; if (i >= p.n || t >= qd.T) { return; }
184
+ let e = ids[t] * p.n + i;
185
+ y[t * p.n + i] = i8at(q[e >> 2u], e & 3u) * f32(s[e >> 5u]);
186
+ }`,
187
+ // One workgroup per token row.
188
+ rmsnormB: `
189
+ struct P { n: u32, eps: f32 }
190
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
191
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
192
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
193
+ @group(0) @binding(3) var<uniform> p: P;
194
+ var<workgroup> red: array<f32, 256>;
195
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
196
+ let b = wg.x * p.n; let t = l.x; var acc = 0.0;
197
+ for (var i = t; i < p.n; i += 256u) { let v = x[b + i]; acc += v * v; }
198
+ red[t] = acc; workgroupBarrier();
199
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
200
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
201
+ for (var i = t; i < p.n; i += 256u) { y[b + i] = (x[b + i] * scale) * w[i]; }
202
+ }`,
203
+ // Tiled Q8_0 GEMM: y[t·M + m] = Σ_k x[t·N + k]·W[m][k]. A workgroup computes 16 rows × 16 tokens;
204
+ // each 32-wide k block (one Q8_0 block) of W and x goes through shared memory, so every weight
205
+ // byte read serves 16 tokens. Dispatch (ceil(M/16), ceil(T/16)).
206
+ matmulQ8T: `enable f16;
207
+ struct P { M: u32, N: u32 }
208
+ ${QB}
209
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
210
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
211
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
212
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
213
+ @group(0) @binding(4) var<uniform> p: P;
214
+ @group(0) @binding(5) var<uniform> qd: Q;
215
+ ${I8}
216
+ var<workgroup> ws: array<f32, 512>;
217
+ var<workgroup> xs: array<f32, 512>;
218
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
219
+ let m0 = wg.x * 16u; let t0 = wg.y * 16u; let r = l.x / 16u; let c = l.x % 16u;
220
+ let nb = p.N / 32u; var acc = 0.0;
221
+ for (var b = 0u; b < nb; b++) {
222
+ for (var i = l.x; i < 512u; i += 256u) {
223
+ let rr = i / 32u; let kk = i % 32u;
224
+ let m = m0 + rr; var wv = 0.0;
225
+ if (m < p.M) { let e = m * p.N + b * 32u + kk; wv = i8at(qw[e >> 2u], e & 3u) * f32(qs[m * nb + b]); }
226
+ ws[i] = wv;
227
+ let t = t0 + rr; var xv = 0.0;
228
+ if (t < qd.T) { xv = x[t * p.N + b * 32u + kk]; }
229
+ xs[i] = xv;
230
+ }
231
+ workgroupBarrier();
232
+ for (var kk = 0u; kk < 32u; kk++) { acc += ws[r * 32u + kk] * xs[c * 32u + kk]; }
233
+ workgroupBarrier();
234
+ }
235
+ let m = m0 + r; let t = t0 + c;
236
+ if (m < p.M && t < qd.T) { y[t * p.M + m] = acc; }
237
+ }`,
238
+ // F32 matrix × T vectors, one workgroup per (row, token). Dispatch (M, T).
239
+ matmulF32T: `
240
+ struct P { M: u32, N: u32 }
241
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
242
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
243
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
244
+ @group(0) @binding(3) var<uniform> p: P;
245
+ var<workgroup> red: array<f32, 256>;
246
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
247
+ let m = wg.x; let t = wg.y; let i = l.x; var acc = 0.0;
248
+ for (var k = i; k < p.N; k += 256u) { acc += w[m * p.N + k] * x[t * p.N + k]; }
249
+ red[i] = acc; workgroupBarrier();
250
+ for (var s = 128u; s > 0u; s >>= 1u) { if (i < s) { red[i] += red[i + s]; } workgroupBarrier(); }
251
+ if (i == 0u) { y[t * p.M + m] = red[0]; }
252
+ }`,
253
+ // The router top-k for token wg.x (same picks and weights as the one-token kernel).
254
+ topkB: `
255
+ struct P { n: u32, k: u32 }
256
+ @group(0) @binding(0) var<storage, read> lg: array<f32>;
257
+ @group(0) @binding(1) var<storage, read_write> sel: array<u32>;
258
+ @group(0) @binding(2) var<uniform> p: P;
259
+ var<workgroup> red: array<f32, 256>;
260
+ var<workgroup> bi: array<u32, 256>;
261
+ var<workgroup> gmax: f32;
262
+ var<workgroup> gsum: f32;
263
+ var<workgroup> picks: array<u32, 32>;
264
+ var<workgroup> probs: array<f32, 32>;
265
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
266
+ let t = l.x; let n = p.n; let lb = wg.x * n; let sb = wg.x * 2u * p.k;
267
+ red[t] = select(-3.4e38, lg[lb + min(t, n - 1u)], t < n); workgroupBarrier();
268
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
269
+ if (t == 0u) { gmax = red[0]; } workgroupBarrier();
270
+ red[t] = select(0.0, exp(lg[lb + min(t, n - 1u)] - gmax), t < n); workgroupBarrier();
271
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
272
+ if (t == 0u) { gsum = red[0]; } workgroupBarrier();
273
+ var mine = select(-1.0, exp(lg[lb + min(t, n - 1u)] - gmax) / gsum, t < n);
274
+ for (var j = 0u; j < p.k; j++) {
275
+ red[t] = mine; bi[t] = t; workgroupBarrier();
276
+ for (var s = 128u; s > 0u; s >>= 1u) {
277
+ if (t < s) { let o = red[t + s]; if (o > red[t] || (o == red[t] && bi[t + s] < bi[t])) { red[t] = o; bi[t] = bi[t + s]; } }
278
+ workgroupBarrier();
279
+ }
280
+ if (t == 0u) { picks[j] = bi[0]; probs[j] = red[0]; }
281
+ workgroupBarrier();
282
+ if (t == picks[j]) { mine = -1.0; }
283
+ workgroupBarrier();
284
+ }
285
+ if (t == 0u) {
286
+ var wsum = 0.0;
287
+ for (var j = 0u; j < p.k; j++) { wsum += probs[j]; }
288
+ for (var j = 0u; j < p.k; j++) { sel[sb + j] = picks[j]; sel[sb + p.k + j] = bitcast<u32>(probs[j] / wsum); }
289
+ }
290
+ }`,
291
+ // The 4-tap conv over T tokens in order, rolling the state.
292
+ gdnConvB: `
293
+ struct P { n: u32 }
294
+ ${QB}
295
+ @group(0) @binding(0) var<storage, read> mixed: array<f32>;
296
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
297
+ @group(0) @binding(2) var<storage, read_write> st: array<f32>;
298
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
299
+ @group(0) @binding(4) var<uniform> p: P;
300
+ @group(0) @binding(5) var<uniform> qd: Q;
301
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
302
+ let c = gi.x; if (c >= p.n) { return; }
303
+ var s0 = st[c * 3u]; var s1 = st[c * 3u + 1u]; var s2 = st[c * 3u + 2u];
304
+ let w0 = w[c * 4u]; let w1 = w[c * 4u + 1u]; let w2 = w[c * 4u + 2u]; let w3 = w[c * 4u + 3u];
305
+ for (var t = 0u; t < qd.T; t++) {
306
+ let x = mixed[t * p.n + c];
307
+ var sum = 0.0;
308
+ sum += s0 * w0; sum += s1 * w1; sum += s2 * w2; sum += x * w3;
309
+ y[t * p.n + c] = sum / (1.0 + exp(-sum));
310
+ s0 = s1; s1 = s2; s2 = x;
311
+ }
312
+ st[c * 3u] = s0; st[c * 3u + 1u] = s1; st[c * 3u + 2u] = s2;
313
+ }`,
314
+ // The gated delta rule over T tokens in order (one workgroup per value head, as gdnStep).
315
+ gdnStepB: `
316
+ struct P { kH: u32, dS: u32, keyDim: u32, convDim: u32, vH: u32, pad: u32, eps: f32, scale: f32 }
317
+ ${QB}
318
+ @group(0) @binding(0) var<storage, read> cv: array<f32>;
319
+ @group(0) @binding(1) var<storage, read> bb: array<f32>;
320
+ @group(0) @binding(2) var<storage, read> al: array<f32>;
321
+ @group(0) @binding(3) var<storage, read> dt: array<f32>;
322
+ @group(0) @binding(4) var<storage, read> aa: array<f32>;
323
+ @group(0) @binding(5) var<storage, read_write> M: array<f32>;
324
+ @group(0) @binding(6) var<storage, read_write> o: array<f32>;
325
+ @group(0) @binding(7) var<uniform> p: P;
326
+ @group(0) @binding(8) var<uniform> qd: Q;
327
+ ${SIGMOID}
328
+ var<workgroup> kq: array<f32, 128>;
329
+ var<workgroup> qq: array<f32, 128>;
330
+ var<workgroup> rq: array<f32, 128>;
331
+ var<workgroup> rk: array<f32, 128>;
332
+ fn log1p_(u: f32) -> f32 { let y = 1.0 + u; return log(y) - ((y - 1.0) - u) / y; }
333
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
334
+ let h = wg.x; let j = l.x; let dS = p.dS; let kh = h % p.kH;
335
+ let n = f32(dS); let inv = 1.0 / sqrt(n);
336
+ let row = (h * dS + j) * dS;
337
+ let dth = dt[h]; let aah = aa[h];
338
+ for (var t = 0u; t < qd.T; t++) {
339
+ let cb = t * p.convDim;
340
+ let qv = cv[cb + kh * dS + j]; let kv = cv[cb + p.keyDim + kh * dS + j]; let vv = cv[cb + 2u * p.keyDim + h * dS + j];
341
+ rq[j] = qv * qv; rk[j] = kv * kv; workgroupBarrier();
342
+ for (var s = 64u; s > 0u; s >>= 1u) { if (j < s) { rq[j] += rq[j + s]; rk[j] += rk[j + s]; } workgroupBarrier(); }
343
+ qq[j] = (qv * (1.0 / sqrt(rq[0] / n + p.eps / n))) * inv;
344
+ kq[j] = (kv * (1.0 / sqrt(rk[0] / n + p.eps / n))) * inv;
345
+ workgroupBarrier();
346
+ let beta = sigm(bb[t * p.vH + h]);
347
+ let x = al[t * p.vH + h] + dth;
348
+ let sp = select(log1p_(exp(x)), x, x > 20.0);
349
+ let decay = exp(sp * aah);
350
+ var sk = 0.0;
351
+ for (var i = 0u; i < dS; i++) { let m = M[row + i] * decay; M[row + i] = m; sk += m * kq[i]; }
352
+ let d = (vv - sk) * beta;
353
+ var acc = 0.0;
354
+ for (var i = 0u; i < dS; i++) { let m = M[row + i] + kq[i] * d; M[row + i] = m; acc += m * qq[i]; }
355
+ o[t * p.vH * dS + h * dS + j] = acc * p.scale;
356
+ workgroupBarrier();
357
+ }
358
+ }`,
359
+ // RMSNorm(o)·w·SiLU(z) for (value head wg.x, token wg.y).
360
+ gdnNormGateB: `
361
+ struct P { dS: u32, dInner: u32, eps: f32 }
362
+ @group(0) @binding(0) var<storage, read> o: array<f32>;
363
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
364
+ @group(0) @binding(2) var<storage, read> z: array<f32>;
365
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
366
+ @group(0) @binding(4) var<uniform> p: P;
367
+ var<workgroup> red: array<f32, 128>;
368
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
369
+ let j = l.x; let i = wg.y * p.dInner + wg.x * p.dS + j;
370
+ let v = o[i]; red[j] = v * v; workgroupBarrier();
371
+ for (var s = 64u; s > 0u; s >>= 1u) { if (j < s) { red[j] += red[j + s]; } workgroupBarrier(); }
372
+ let scale = 1.0 / sqrt(red[0] / f32(p.dS) + p.eps);
373
+ let zz = z[i];
374
+ y[i] = ((v * scale) * w[j]) * (zz / (1.0 + exp(-zz)));
375
+ }`,
376
+ // qkNormRopeP for (head wg.x, token wg.y); rope holds nRot values per token.
377
+ qkNormRopePB: `
378
+ struct P { hd: u32, srcStride: u32, srcRow: u32, dstRow: u32, nRot: u32, gate: u32, eps: f32 }
379
+ @group(0) @binding(0) var<storage, read> src: array<f32>;
380
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
381
+ @group(0) @binding(2) var<storage, read> rope: array<f32>;
382
+ @group(0) @binding(3) var<storage, read_write> dst: array<f32>;
383
+ @group(0) @binding(4) var<storage, read_write> gout: array<f32>;
384
+ @group(0) @binding(5) var<uniform> p: P;
385
+ var<workgroup> red: array<f32, 64>;
386
+ @compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
387
+ let tk = wg.y; let sb = tk * p.srcRow + wg.x * p.srcStride; let ob = tk * p.dstRow + wg.x * p.hd; let rb = tk * p.nRot;
388
+ let t = l.x; let half = p.nRot / 2u;
389
+ var acc = 0.0;
390
+ for (var i = t; i < p.hd; i += 64u) { let v = src[sb + i]; acc += v * v; }
391
+ red[t] = acc; workgroupBarrier();
392
+ for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
393
+ let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
394
+ for (var i = t; i < p.hd; i += 64u) {
395
+ if (i < half) {
396
+ let a = (src[sb + i] * scale) * w[i];
397
+ let b = (src[sb + i + half] * scale) * w[i + half];
398
+ let c = rope[rb + i]; let sn = rope[rb + half + i];
399
+ dst[ob + i] = a * c - b * sn;
400
+ dst[ob + i + half] = a * sn + b * c;
401
+ } else if (i >= p.nRot) {
402
+ dst[ob + i] = (src[sb + i] * scale) * w[i];
403
+ }
404
+ if (p.gate == 1u) { gout[ob + i] = src[sb + p.hd + i]; }
405
+ }
406
+ }`,
407
+ kvStoreB: `enable f16;
408
+ struct P { n: u32 }
409
+ ${QB}
410
+ @group(0) @binding(0) var<storage, read> k: array<f32>;
411
+ @group(0) @binding(1) var<storage, read> v: array<f32>;
412
+ @group(0) @binding(2) var<storage, read_write> kc: array<f16>;
413
+ @group(0) @binding(3) var<storage, read_write> vc: array<f16>;
414
+ @group(0) @binding(4) var<uniform> p: P;
415
+ @group(0) @binding(5) var<uniform> qd: Q;
416
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
417
+ let i = g.x; let t = g.y; if (i >= p.n || t >= qd.T) { return; }
418
+ let pos = qd.pos0 + t;
419
+ kc[pos * p.n + i] = f16(k[t * p.n + i]);
420
+ vc[pos * p.n + i] = f16(v[t * p.n + i]);
421
+ }`,
422
+ // Causal scores for T queries over S = pos0 + T keys; masked keys get -3.4e38.
423
+ attnScoreB: `enable f16;
424
+ struct P { heads: u32, kvHeads: u32, hd: u32, scale: f32 }
425
+ ${QB}
426
+ @group(0) @binding(0) var<storage, read> q: array<f32>;
427
+ @group(0) @binding(1) var<storage, read> kc: array<f16>;
428
+ @group(0) @binding(2) var<storage, read_write> sc: array<f32>;
429
+ @group(0) @binding(3) var<uniform> p: P;
430
+ @group(0) @binding(4) var<uniform> qd: Q;
431
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>) {
432
+ let idx = g.y * ng.x * 256u + g.x; let S = qd.S; let T = qd.T;
433
+ if (idx >= p.heads * T * S) { return; }
434
+ let h = idx / (T * S); let rem = idx % (T * S); let t = rem / S; let pos = rem % S;
435
+ if (pos > qd.pos0 + t) { sc[idx] = -3.4e38; return; }
436
+ let kvh = h / (p.heads / p.kvHeads);
437
+ let qo = t * p.heads * p.hd + h * p.hd; let ko = pos * p.kvHeads * p.hd + kvh * p.hd;
438
+ var acc = 0.0;
439
+ for (var d = 0u; d < p.hd; d++) { acc += q[qo + d] * f32(kc[ko + d]); }
440
+ sc[idx] = acc * p.scale;
441
+ }`,
442
+ softmaxB: `
443
+ ${QB}
444
+ @group(0) @binding(0) var<storage, read_write> sc: array<f32>;
445
+ @group(0) @binding(1) var<uniform> qd: Q;
446
+ var<workgroup> red: array<f32, 256>;
447
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
448
+ let n = qd.S; let base = wg.x * n; let t = l.x;
449
+ var mx = -3.4e38;
450
+ for (var i = t; i < n; i += 256u) { mx = max(mx, sc[base + i]); }
451
+ red[t] = mx; workgroupBarrier();
452
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
453
+ let m = red[0]; workgroupBarrier();
454
+ var sum = 0.0;
455
+ for (var i = t; i < n; i += 256u) { let e = exp(sc[base + i] - m); sc[base + i] = e; sum += e; }
456
+ red[t] = sum; workgroupBarrier();
457
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
458
+ let inv = 1.0 / red[0];
459
+ for (var i = t; i < n; i += 256u) { sc[base + i] = sc[base + i] * inv; }
460
+ }`,
461
+ attnOutB: `enable f16;
462
+ struct P { heads: u32, kvHeads: u32, hd: u32, pad: u32 }
463
+ ${QB}
464
+ @group(0) @binding(0) var<storage, read> pr: array<f32>;
465
+ @group(0) @binding(1) var<storage, read> vc: array<f16>;
466
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
467
+ @group(0) @binding(3) var<uniform> p: P;
468
+ @group(0) @binding(4) var<uniform> qd: Q;
469
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
470
+ let QN = p.heads * p.hd; let idx = g.x; let t = g.y;
471
+ if (idx >= QN || t >= qd.T) { return; }
472
+ let h = idx / p.hd; let d = idx % p.hd; let kvh = h / (p.heads / p.kvHeads);
473
+ let stride = p.kvHeads * p.hd; let rb = (h * qd.T + t) * qd.S; let last = qd.pos0 + t;
474
+ var acc = 0.0;
475
+ for (var pos = 0u; pos <= last; pos++) { acc += pr[rb + pos] * f32(vc[pos * stride + kvh * p.hd + d]); }
476
+ y[t * QN + idx] = acc;
477
+ }`,
478
+ sigmoidMulB: `
479
+ struct P { n: u32 }
480
+ ${QB}
481
+ @group(0) @binding(0) var<storage, read_write> x: array<f32>;
482
+ @group(0) @binding(1) var<storage, read> g: array<f32>;
483
+ @group(0) @binding(2) var<uniform> p: P;
484
+ @group(0) @binding(3) var<uniform> qd: Q;
485
+ ${SIGMOID}
486
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
487
+ let i = gi.x; let t = gi.y; if (i >= p.n || t >= qd.T) { return; }
488
+ let k = t * p.n + i; x[k] = x[k] * sigm(g[k]);
489
+ }`,
490
+ addInPlaceB: `
491
+ struct P { n: u32 }
492
+ ${QB}
493
+ @group(0) @binding(0) var<storage, read_write> x: array<f32>;
494
+ @group(0) @binding(1) var<storage, read> y: array<f32>;
495
+ @group(0) @binding(2) var<uniform> p: P;
496
+ @group(0) @binding(3) var<uniform> qd: Q;
497
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
498
+ let i = gi.x; let t = gi.y; if (i >= p.n || t >= qd.T) { return; }
499
+ let k = t * p.n + i; x[k] = x[k] + y[k];
500
+ }`,
501
+ // Routed experts for all (token, pick) pairs: z = t·K + j, slots[z] its pool slot. Input is the
502
+ // token's row (gate/up, inputMode 0) or the pair's own row (down, inputMode 1).
503
+ expertQ8B: `enable f16;
504
+ struct P { M: u32, N: u32, rowsPerExpert: u32, inputMode: u32, K: u32 }
505
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
506
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
507
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
508
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
509
+ @group(0) @binding(4) var<uniform> p: P;
510
+ @group(0) @binding(5) var<storage, read> slots: array<u32>;
511
+ ${I8}
512
+ var<workgroup> part: array<f32, 1024>;
513
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
514
+ let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let N = p.N; let M = p.M; let nb = N / 32u;
515
+ if (m0 >= M) { return; }
516
+ let rb = slots[z] * p.rowsPerExpert;
517
+ let xb = select((z / p.K) * N, z * N, p.inputMode == 1u);
518
+ let r0 = rb + min(m0, M - 1u); let r1 = rb + min(m0 + 1u, M - 1u); let r2 = rb + min(m0 + 2u, M - 1u); let r3 = rb + min(m0 + 3u, M - 1u);
519
+ let q0 = (r0 * N) >> 2u; let q1 = (r1 * N) >> 2u; let q2 = (r2 * N) >> 2u; let q3 = (r3 * N) >> 2u;
520
+ var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
521
+ for (var k = t; k < N; k += 256u) {
522
+ let xv = x[xb + k]; let wi = k >> 2u; let ln = k & 3u; let b = k >> 5u;
523
+ a0 += i8at(qw[q0 + wi], ln) * f32(qs[r0 * nb + b]) * xv;
524
+ a1 += i8at(qw[q1 + wi], ln) * f32(qs[r1 * nb + b]) * xv;
525
+ a2 += i8at(qw[q2 + wi], ln) * f32(qs[r2 * nb + b]) * xv;
526
+ a3 += i8at(qw[q3 + wi], ln) * f32(qs[r3 * nb + b]) * xv;
527
+ }
528
+ part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
529
+ workgroupBarrier();
530
+ for (var s = 128u; s > 0u; s >>= 1u) {
531
+ if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
532
+ workgroupBarrier();
533
+ }
534
+ if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
535
+ }`,
536
+ // act = SiLU(g)·u, elementwise over T × n (the shared expert's separate gate and up).
537
+ siluMul2B: `
538
+ struct P { n: u32 }
539
+ ${QB}
540
+ @group(0) @binding(0) var<storage, read> g: array<f32>;
541
+ @group(0) @binding(1) var<storage, read> u: array<f32>;
542
+ @group(0) @binding(2) var<storage, read_write> act: array<f32>;
543
+ @group(0) @binding(3) var<uniform> p: P;
544
+ @group(0) @binding(4) var<uniform> qd: Q;
545
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
546
+ let i = gi.x; let t = gi.y; if (i >= p.n || t >= qd.T) { return; }
547
+ let k = t * p.n + i; let gv = g[k];
548
+ act[k] = (gv / (1.0 + exp(-gv))) * u[k];
549
+ }`,
550
+ moeAccumSharedB: `
551
+ struct P { H: u32, k: u32 }
552
+ ${QB}
553
+ @group(0) @binding(0) var<storage, read> out: array<f32>;
554
+ @group(0) @binding(1) var<storage, read> sel: array<u32>;
555
+ @group(0) @binding(2) var<storage, read> sh: array<f32>;
556
+ @group(0) @binding(3) var<storage, read> sg: array<f32>;
557
+ @group(0) @binding(4) var<storage, read_write> x: array<f32>;
558
+ @group(0) @binding(5) var<uniform> p: P;
559
+ @group(0) @binding(6) var<uniform> qd: Q;
560
+ ${SIGMOID}
561
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gi: vec3<u32>) {
562
+ let h = gi.x; let t = gi.y; if (h >= p.H || t >= qd.T) { return; }
563
+ let sb = t * 2u * p.k; var acc = 0.0;
564
+ for (var j = 0u; j < p.k; j++) { acc += bitcast<f32>(sel[sb + p.k + j]) * out[(t * p.k + j) * p.H + h]; }
565
+ let r = t * p.H + h;
566
+ x[r] = (acc + sh[r] * sigm(sg[t])) + x[r];
567
+ }`,
568
+ };
569
+
570
+ // ── Subgroup Q8_0 GEMV (the LFM2.5 / Gemma 4 kernels' pattern) ──────────────────────────────
571
+ // One 32-lane subgroup per output row, 4 rows per 128-thread workgroup (the same dispatch
572
+ // geometry as matmulQ8/expertQ8). Each lane takes whole Q8_0 blocks: two vec4<u32> of int8
573
+ // codes, unpack4xI8, vec4 dot products against the activations, the block scale applied once;
574
+ // one subgroupAdd reduces the row. Used only when the adapter's subgroups are exactly 32 wide.
575
+ const SG_KERNELS = {
576
+ matmulQ8: `enable f16;
577
+ enable subgroups;
578
+ struct P { M: u32, N: u32 }
579
+ @group(0) @binding(0) var<storage, read> x: array<vec4<f32>>;
580
+ @group(0) @binding(1) var<storage, read> qw: array<vec4<u32>>;
581
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
582
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
583
+ @group(0) @binding(4) var<uniform> p: P;
584
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
585
+ let m = (wg.y * ng.x + wg.x) * 4u + l.x / 32u; let lane = l.x % 32u;
586
+ let nb = p.N / 32u; let mr = min(m, p.M - 1u);
587
+ let wb = mr * nb * 2u; let sb = mr * nb;
588
+ var acc = 0.0;
589
+ for (var b = lane; b < nb; b += 32u) {
590
+ let w0 = qw[wb + b * 2u]; let w1 = qw[wb + b * 2u + 1u]; let xb = b * 8u;
591
+ var s = dot(vec4<f32>(unpack4xI8(w0.x)), x[xb]) + dot(vec4<f32>(unpack4xI8(w0.y)), x[xb + 1u]);
592
+ s += dot(vec4<f32>(unpack4xI8(w0.z)), x[xb + 2u]) + dot(vec4<f32>(unpack4xI8(w0.w)), x[xb + 3u]);
593
+ s += dot(vec4<f32>(unpack4xI8(w1.x)), x[xb + 4u]) + dot(vec4<f32>(unpack4xI8(w1.y)), x[xb + 5u]);
594
+ s += dot(vec4<f32>(unpack4xI8(w1.z)), x[xb + 6u]) + dot(vec4<f32>(unpack4xI8(w1.w)), x[xb + 7u]);
595
+ acc += s * f32(qs[sb + b]);
596
+ }
597
+ let total = subgroupAdd(acc);
598
+ if (lane == 0u && m < p.M) { y[m] = total; }
599
+ }`,
600
+ expertQ8: `enable f16;
601
+ enable subgroups;
602
+ struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
603
+ @group(0) @binding(0) var<storage, read> x: array<vec4<f32>>;
604
+ @group(0) @binding(1) var<storage, read> qw: array<vec4<u32>>;
605
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
606
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
607
+ @group(0) @binding(4) var<uniform> p: P;
608
+ @group(0) @binding(5) var<storage, read> slots: array<u32>;
609
+ @compute @workgroup_size(128) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
610
+ let z = wg.z; let m = (wg.y * ng.x + wg.x) * 4u + l.x / 32u; let lane = l.x % 32u;
611
+ let nb = p.N / 32u; let r = slots[z] * p.rowsPerExpert + min(m, p.M - 1u);
612
+ let wb = r * nb * 2u; let sb = r * nb;
613
+ let xo = select(0u, z * p.N / 4u, p.inputPerSlot == 1u);
614
+ var acc = 0.0;
615
+ for (var b = lane; b < nb; b += 32u) {
616
+ let w0 = qw[wb + b * 2u]; let w1 = qw[wb + b * 2u + 1u]; let xb = xo + b * 8u;
617
+ var s = dot(vec4<f32>(unpack4xI8(w0.x)), x[xb]) + dot(vec4<f32>(unpack4xI8(w0.y)), x[xb + 1u]);
618
+ s += dot(vec4<f32>(unpack4xI8(w0.z)), x[xb + 2u]) + dot(vec4<f32>(unpack4xI8(w0.w)), x[xb + 3u]);
619
+ s += dot(vec4<f32>(unpack4xI8(w1.x)), x[xb + 4u]) + dot(vec4<f32>(unpack4xI8(w1.y)), x[xb + 5u]);
620
+ s += dot(vec4<f32>(unpack4xI8(w1.z)), x[xb + 6u]) + dot(vec4<f32>(unpack4xI8(w1.w)), x[xb + 7u]);
621
+ acc += s * f32(qs[sb + b]);
622
+ }
623
+ let total = subgroupAdd(acc);
624
+ if (lane == 0u && m < p.M) { y[z * p.M + m] = total; }
625
+ }`,
626
+ };
627
+
628
+ // Routing on the GPU (decode): after layer l's top-k, look each pick up in the pool's residency
629
+ // map. All resident → write the slots and carry on; any absent → record l+1 in `st`, after which
630
+ // every guarded kernel returns at once, so the rest of the token costs dispatches, not math.
631
+ const ROUTE_KERNEL = `
632
+ struct P { l: u32, E: u32, K: u32, pad: u32 }
633
+ @group(0) @binding(0) var<storage, read> sel: array<u32>;
634
+ @group(0) @binding(1) var<storage, read> rmap: array<u32>;
635
+ @group(0) @binding(2) var<storage, read_write> slots: array<u32>;
636
+ @group(0) @binding(3) var<storage, read_write> st: array<u32>;
637
+ @group(0) @binding(4) var<storage, read_write> used: array<u32>;
638
+ @group(0) @binding(5) var<uniform> p: P;
639
+ @compute @workgroup_size(1) fn main() {
640
+ if (st[0] != 0u) { return; }
641
+ var miss = false;
642
+ for (var j = 0u; j < p.K; j++) {
643
+ let e = sel[j]; let s = rmap[p.l * p.E + e];
644
+ used[p.l * p.K + j] = e;
645
+ if (s == 0xffffffffu) { miss = true; } else { slots[j] = s; }
646
+ }
647
+ if (miss) { st[0] = p.l + 1u; }
648
+ }`;
649
+ const ABSENT = 0xffffffff;
650
+
651
+ // Adds a read-only `stopF` binding after a kernel's last binding and an early return while it is
652
+ // set. A read-only storage value is uniform, so the return keeps barriers and subgroup ops legal.
653
+ function guarded(src) {
654
+ let last = -1;
655
+ for (const m of src.matchAll(/@binding\((\d+)\)/g)) last = Math.max(last, +m[1]);
656
+ const at = src.indexOf('@compute');
657
+ const out = `${src.slice(0, at)}@group(0) @binding(${last + 1}) var<storage, read> stopF: array<u32>;\n${src.slice(at)}`;
658
+ const withReturn = out.replace(/(@compute[^{]*?\bfn main\([\s\S]*?\)\s*\{)/, '$1\n if (stopF[0] != 0u) { return; }');
659
+ if (withReturn === out) throw new Error('guarded(): no compute entry point');
660
+ return withReturn;
661
+ }
662
+
663
+ export class Qwen35MoeSsd extends Qwen3MoeSsd {
664
+ // Without a url, the pinned Hugging Face file; with one (a local serve, a layer cut), the OPFS
665
+ // directory is named after the file, never after rung 2a's model.
666
+ static async load(modelId = null, opts = {}) {
667
+ const source = opts.source || (opts.url ? { repo: null, file: opts.url.split('/').pop(), revision: null } : QWEN36_35B_A3B);
668
+ return super.load(modelId, { ...opts, source });
669
+ }
670
+
671
+ constructor(device, manifest, gguf, opts) {
672
+ super(device, manifest, gguf, opts);
673
+ const c = this.cfg;
674
+ if (c.arch !== 'qwen35moe') throw new Error(`Qwen35MoeSsd needs a qwen35moe GGUF, got ${c.arch}`);
675
+ if (c.ssm.dConv !== 4) throw new Error(`gdnConv assumes a 4-tap conv, got ${c.ssm.dConv}`);
676
+ if (c.ssm.dState !== 128 || c.ssm.dInner !== c.ssm.vHeads * c.ssm.dState) throw new Error('gdnStep assumes 128-wide heads');
677
+ this.recurrent = Array.from({ length: c.layers }, (_, l) => (l + 1) % c.attnInterval !== 0);
678
+ // Prompts go through in chunks of up to chunkTokens tokens per layer (one routing readback per
679
+ // layer per chunk); decode stays one token at a time.
680
+ this.batchPrefill = opts.batchPrefill ?? true;
681
+ this.chunkTokens = Math.max(2, Math.floor(opts.prefillChunk || 256));
682
+ // gpuRouting: decode syncs with the CPU only at a layer whose experts are not all in the pool
683
+ // (and every routeWindow layers), instead of after every layer's routing. Opt-in at load: it
684
+ // compiles every one-token kernel with a stop guard. Off by default until it is timed on a
685
+ // quiet machine; at 2026-10-04's miss rate (~18 of 40 layers per token need a load) the counts
686
+ // put the gain near 5%. Output is token-identical either way.
687
+ this.routeGuards = !!opts.gpuRouting;
688
+ this.gpuRouting = this.routeGuards;
689
+ // A submit runs at most this many layers' experts. After a stop every kernel left in it still
690
+ // dispatches (about 7 ms for a whole token's worth), so a long submit mostly runs no-ops.
691
+ // Full model, 48 tokens: window 2 → 29.1 syncs and 17.6 no-op layers per token; 3 → 26.2, 36.6.
692
+ this.routeWindow = Math.max(1, Math.floor(opts.routeWindow ?? 2));
693
+ }
694
+
695
+ // Subgroup GEMVs replace the scalar ones when subgroups are exactly 32 wide (Apple, most others).
696
+ get subgroupKernels() {
697
+ const i = this.device.adapterInfo || this.opts.adapterInfo || {};
698
+ return this.opts.subgroups !== false && this.device.features.has('subgroups') && i.subgroupMinSize === 32 && i.subgroupMaxSize === 32;
699
+ }
700
+ // With gpuRouting, every one-token kernel carries the stop guard (chunked prefill's do not).
701
+ get kernels() {
702
+ if (!this._kernels) {
703
+ const k = { ...super.kernels, ...KERNELS, ...(this.subgroupKernels ? SG_KERNELS : {}) };
704
+ if (this.routeGuards) for (const n of Object.keys(k)) k[n] = guarded(k[n]);
705
+ this._kernels = { ...k, ...BATCH_KERNELS, route: ROUTE_KERNEL };
706
+ }
707
+ return this._kernels;
708
+ }
709
+ isGuarded(name) { return this.routeGuards && name !== 'route' && !(name in BATCH_KERNELS); }
710
+
711
+ bind(name, buffers) { return super.bind(name, this.isGuarded(name) ? [...buffers, this.a.stop] : buffers); }
712
+ // Like bind(), but a resource may be { buffer, offset, size }.
713
+ bindR(name, resources) {
714
+ const all = this.isGuarded(name) ? [...resources, this.a.stop] : resources;
715
+ return this.device.createBindGroup({
716
+ layout: this.pipeline(name).getBindGroupLayout(0),
717
+ entries: all.map((r, i) => ({ binding: i, resource: r instanceof GPUBuffer ? { buffer: r } : r })),
718
+ });
719
+ }
720
+
721
+ // For the GPU budget: f16 KV for the attention layers, conv + delta-rule state for the DeltaNet layers.
722
+ kvBytes() {
723
+ const c = this.cfg, s = c.ssm, convDim = 2 * s.kHeads * s.dState + s.dInner;
724
+ return this.recurrent.reduce((sum, r) => sum + (r ? 3 * convDim * 4 + s.vHeads * s.dState * s.dState * 4 : 2 * 2 * this.maxCtx * c.kvHeads * c.headDim), 0);
725
+ }
726
+
727
+ initBuffers() {
728
+ const c = this.cfg, s = c.ssm;
729
+ const H = c.hidden, hd = c.headDim, QN = c.heads * hd, kvn = c.kvHeads * hd, F = c.expertFf, SF = c.shexpFf, K = c.topK;
730
+ const keyDim = s.kHeads * s.dState, convDim = 2 * keyDim + s.dInner;
731
+ if (2 * K * 4 > 256) throw new Error(`top-k ${K} > 32: selections would overlap their 256-byte readback regions`);
732
+ if (c.experts > 256) throw new Error(`${c.experts} experts: the router top-k kernel handles at most 256`);
733
+ if ((SF * 4) % 256) throw new Error('shared expert gate/up halves must be 256-byte aligned');
734
+
735
+ // KV cache (f16) for the attention layers; conv + delta-rule state for the DeltaNet layers.
736
+ this.kc = []; this.vc = []; this.conv = []; this.ssm = [];
737
+ for (let l = 0; l < c.layers; l++) {
738
+ if (this.recurrent[l]) {
739
+ this.conv[l] = this.buffer(3 * convDim * 4, STORAGE | COPY_DST, 'state');
740
+ this.ssm[l] = this.buffer(s.vHeads * s.dState * s.dState * 4, STORAGE | COPY_DST, 'state');
741
+ } else {
742
+ this.kc[l] = this.buffer(this.maxCtx * kvn * 2, STORAGE, 'kv');
743
+ this.vc[l] = this.buffer(this.maxCtx * kvn * 2, STORAGE, 'kv');
744
+ }
745
+ }
746
+ const A = (n) => this.buffer(n * 4, STORAGE | COPY_SRC | COPY_DST);
747
+ this.a = {
748
+ x: A(H), xn: A(H), qg: A(2 * QN), q: A(QN), gate: A(QN), k: A(kvn), kr: A(kvn), v: A(kvn), att: A(QN), o: A(H),
749
+ sc: A(c.heads * this.maxCtx), mixed: A(convDim), cv: A(convDim), z: A(s.dInner), bb: A(s.vHeads), al: A(s.vHeads),
750
+ go: A(s.dInner), gn: A(s.dInner), dummy: A(4),
751
+ rl: A(c.experts), sel: A(2 * K), pxn: A(H), prl: A(c.experts), psel: A(64 * this.lookahead), slots: A(K),
752
+ gu: A(K * 2 * F), act: A(K * F), dn: A(K * H), shgu: A(2 * SF), shact: A(SF), sh: A(H), sg: A(4),
753
+ logits: A(c.vocab), am: A(4), rope: A(c.ropeDims),
754
+ stop: A(4), used: A(c.layers * K), rmap: A(c.layers * c.experts),
755
+ };
756
+ this.tok = this.buffer(16, UNIFORM | COPY_DST);
757
+ this.rbSel = this.buffer(256 * (1 + this.lookahead), MAP_READ | COPY_DST);
758
+ this.rbLogits = this.buffer(c.vocab * 4, MAP_READ | COPY_DST);
759
+ this.rbArg = this.buffer(16, MAP_READ | COPY_DST);
760
+ // One readback per decode sync: [stop | layer's sel | guesses | every layer's picks | argmax].
761
+ this.seg = { sel: 256, psel: 512, used: 512 + 256 * this.lookahead };
762
+ this.seg.am = this.seg.used + Math.ceil(c.layers * K * 4 / 16) * 16;
763
+ this.seg.bytes = this.seg.am + 16;
764
+ this.rbSeg = this.buffer(this.seg.bytes, MAP_READ | COPY_DST);
765
+ this.routeMap = new Uint32Array(c.layers * c.experts);
766
+
767
+ const U = (w) => this.uniform(w), eps = { f: c.eps };
768
+ const u = {
769
+ nH: U([H]), rmsH: U([H, eps]),
770
+ qg: U([2 * QN, H]), kv: U([kvn, H]), o: U([H, QN]),
771
+ ropeQ: U([hd, 2 * hd, c.ropeDims, 1, eps]), ropeK: U([hd, hd, c.ropeDims, 0, eps]), kvn: U([kvn]),
772
+ att: U([c.heads, c.kvHeads, hd, { f: 1 / Math.sqrt(hd) }]), attOut: U([c.heads, c.kvHeads, hd, 0]), gate: U([QN]),
773
+ qkv: U([convDim, H]), z: U([s.dInner, H]), ab: U([s.vHeads, H]), gout: U([H, s.dInner]),
774
+ conv: U([convDim]), step: U([s.kHeads, s.dState, keyDim, 0, eps, { f: 1 / Math.sqrt(s.dState) }]), ng: U([s.dState, eps]),
775
+ router: U([c.experts, H]), topk: U([c.experts, K]),
776
+ gu: U([2 * F, H, 2 * F, 0]), dn: U([H, F, H, 1]), silu: U([F, K]), acc: U([H, K]),
777
+ shgu: U([SF, H]), shsilu: U([SF, 1]), shd: U([H, SF]), shgate: U([1, H]),
778
+ lm: U([c.vocab, H]), am: U([c.vocab]),
779
+ };
780
+ const a = this.a, W = this.w;
781
+ const q8 = (name, x, y, uni) => { const t = W[name]; if (!t || !t.q) throw new Error(`missing Q8_0 tensor ${name}`); return this.bind('matmulQ8', [x, t.q, t.s, y, uni]); };
782
+ const raw = (name) => { const t = W[name]; if (!t || !t.raw) throw new Error(`missing F32 tensor ${name}`); return t.raw; };
783
+ this.g = { embed: this.bind('embedQ8', [W['token_embd.weight'].q, W['token_embd.weight'].s, a.x, this.tok, u.nH]) };
784
+ this.layers = [];
785
+ for (let l = 0; l < c.layers; l++) {
786
+ const n = (t) => `blk.${l}.${t}`;
787
+ const g = { recur: this.recurrent[l], rmsA: this.bind('rmsnorm', [a.x, raw(n('attn_norm.weight')), a.xn, u.rmsH]) };
788
+ if (g.recur) {
789
+ Object.assign(g, {
790
+ qkv: q8(n('attn_qkv.weight'), a.xn, a.mixed, u.qkv),
791
+ z: q8(n('attn_gate.weight'), a.xn, a.z, u.z),
792
+ beta: q8(n('ssm_beta.weight'), a.xn, a.bb, u.ab),
793
+ alpha: q8(n('ssm_alpha.weight'), a.xn, a.al, u.ab),
794
+ conv: this.bind('gdnConv', [a.mixed, raw(n('ssm_conv1d.weight')), this.conv[l], a.cv, u.conv]),
795
+ step: this.bind('gdnStep', [a.cv, a.bb, a.al, raw(n('ssm_dt.bias')), raw(n('ssm_a')), this.ssm[l], a.go, u.step]),
796
+ norm: this.bind('gdnNormGate', [a.go, raw(n('ssm_norm.weight')), a.z, a.gn, u.ng]),
797
+ out: q8(n('ssm_out.weight'), a.gn, a.o, u.gout),
798
+ });
799
+ } else {
800
+ Object.assign(g, {
801
+ q: q8(n('attn_q.weight'), a.xn, a.qg, u.qg),
802
+ k: q8(n('attn_k.weight'), a.xn, a.k, u.kv),
803
+ v: q8(n('attn_v.weight'), a.xn, a.v, u.kv),
804
+ ropeQ: this.bind('qkNormRopeP', [a.qg, raw(n('attn_q_norm.weight')), a.rope, a.q, a.gate, u.ropeQ]),
805
+ ropeK: this.bind('qkNormRopeP', [a.k, raw(n('attn_k_norm.weight')), a.rope, a.kr, a.dummy, u.ropeK]),
806
+ kv: this.bind('kvStore', [a.kr, a.v, this.kc[l], this.vc[l], this.tok, u.kvn]),
807
+ score: this.bind('attnScore', [a.q, this.kc[l], a.sc, this.tok, u.att]),
808
+ soft: this.bind('softmax', [a.sc, this.tok]),
809
+ attOut: this.bind('attnOut', [a.sc, this.vc[l], a.att, this.tok, u.attOut]),
810
+ gate: this.bind('sigmoidMul', [a.att, a.gate, u.gate]),
811
+ o: q8(n('attn_output.weight'), a.att, a.o, u.o),
812
+ });
813
+ }
814
+ Object.assign(g, {
815
+ addO: this.bind('addInPlace', [a.x, a.o, u.nH]),
816
+ rmsF: this.bind('rmsnorm', [a.x, raw(n('post_attention_norm.weight')), a.xn, u.rmsH]),
817
+ router: this.bind('matmulF32', [a.xn, raw(n('ffn_gate_inp.weight')), a.rl, u.router]),
818
+ topk: this.bind('topk', [a.rl, a.sel, u.topk]),
819
+ pf: Array.from({ length: Math.min(this.lookahead, c.layers - 1 - l) }, (_, i) => {
820
+ const at = (t) => `blk.${l + 1 + i}.${t}`;
821
+ return {
822
+ rms: this.bind('rmsnorm', [a.x, raw(at('post_attention_norm.weight')), a.pxn, u.rmsH]),
823
+ router: this.bind('matmulF32', [a.pxn, raw(at('ffn_gate_inp.weight')), a.prl, u.router]),
824
+ topk: this.bindR('topk', [a.prl, { buffer: a.psel, offset: 256 * i, size: 8 * K }, u.topk]),
825
+ };
826
+ }),
827
+ // Shared expert: [gate | up] into shgu, SiLU·up, down; its scalar gate logit into sg.
828
+ shg: this.bindR('matmulQ8', [a.xn, W[n('ffn_gate_shexp.weight')].q, W[n('ffn_gate_shexp.weight')].s, { buffer: a.shgu, offset: 0, size: SF * 4 }, u.shgu]),
829
+ shu: this.bindR('matmulQ8', [a.xn, W[n('ffn_up_shexp.weight')].q, W[n('ffn_up_shexp.weight')].s, { buffer: a.shgu, offset: SF * 4, size: SF * 4 }, u.shgu]),
830
+ shsilu: this.bind('siluMulMoe', [a.shgu, a.shact, u.shsilu]),
831
+ shd: q8(n('ffn_down_shexp.weight'), a.shact, a.sh, u.shd),
832
+ shgate: this.bind('matmulF32', [a.xn, raw(n('ffn_gate_inp_shexp.weight')), a.sg, u.shgate]),
833
+ gu: this.bind('expertQ8', [a.xn, this.pool.guQ, this.pool.guS, a.gu, u.gu, a.slots]),
834
+ silu: this.bind('siluMulMoe', [a.gu, a.act, u.silu]),
835
+ dn: this.bind('expertQ8', [a.act, this.pool.dQ, this.pool.dS, a.dn, u.dn, a.slots]),
836
+ acc: this.bind('moeAccumShared', [a.dn, a.sel, a.sh, a.sg, a.x, u.acc]),
837
+ route: this.bind('route', [a.sel, a.rmap, a.slots, a.stop, a.used, U([l, c.experts, K, 0])]),
838
+ });
839
+ this.layers.push(g);
840
+ }
841
+ this.g.rmsOut = this.bind('rmsnorm', [a.x, raw('output_norm.weight'), a.xn, u.rmsH]);
842
+ this.g.lm = this.bind('matmulQ8', [a.xn, W['output.weight'].q, W['output.weight'].s, a.logits, u.lm]);
843
+ this.g.am = this.bind('argmax', [a.logits, a.am, u.am]);
844
+ // A chunk pins min(experts, T·K) experts per layer at once; keep that under half the pool.
845
+ if (Math.min(c.experts, this.chunkTokens * K) > this.poolSlots / 2) this.chunkTokens = Math.max(2, Math.floor(this.poolSlots / 2 / K));
846
+ if (this.batchPrefill) this.initBatch();
847
+ }
848
+
849
+ // Buffers and bind groups for chunked prefill (rows = tokens of the chunk, token-major).
850
+ initBatch() {
851
+ const c = this.cfg, s = c.ssm, B = this.chunkTokens, a = this.a, W = this.w;
852
+ const H = c.hidden, hd = c.headDim, QN = c.heads * hd, kvn = c.kvHeads * hd, F = c.expertFf, SF = c.shexpFf, K = c.topK;
853
+ const keyDim = s.kHeads * s.dState, convDim = 2 * keyDim + s.dInner;
854
+ const A = (n) => this.buffer(n * 4, STORAGE | COPY_SRC | COPY_DST, 'batch');
855
+ const b = this.b = {
856
+ ids: A(B), x: A(B * H), xn: A(B * H), qg: A(B * 2 * QN), q: A(B * QN), gate: A(B * QN), k: A(B * kvn), kr: A(B * kvn), v: A(B * kvn),
857
+ att: A(B * QN), o: A(B * H), sc: A(c.heads * B * this.maxCtx), mixed: A(B * convDim), cv: A(B * convDim), z: A(B * s.dInner),
858
+ bb: A(B * s.vHeads), al: A(B * s.vHeads), go: A(B * s.dInner), gn: A(B * s.dInner),
859
+ rl: A(B * c.experts), sel: A(B * 2 * K), slots: A(B * K), gu: A(B * K * 2 * F), act: A(B * K * F), dn: A(B * K * H),
860
+ shg: A(B * SF), shu: A(B * SF), shact: A(B * SF), sh: A(B * H), sg: A(B), rope: A(B * c.ropeDims), dummy: A(4),
861
+ };
862
+ this.qb = this.buffer(16, UNIFORM | COPY_DST, 'batch');
863
+ this.rbSelB = this.buffer(B * 2 * K * 4, MAP_READ | COPY_DST, 'batch');
864
+ const U = (w) => this.uniform(w), eps = { f: c.eps }, qb = this.qb;
865
+ const u = {
866
+ nH: U([H]), rmsH: U([H, eps]), qg: U([2 * QN, H]), kv: U([kvn, H]), o: U([H, QN]),
867
+ ropeQ: U([hd, 2 * hd, 2 * QN, QN, c.ropeDims, 1, eps]), ropeK: U([hd, hd, kvn, kvn, c.ropeDims, 0, eps]), kvn: U([kvn]),
868
+ att: U([c.heads, c.kvHeads, hd, { f: 1 / Math.sqrt(hd) }]), attOut: U([c.heads, c.kvHeads, hd, 0]), gate: U([QN]),
869
+ qkv: U([convDim, H]), z: U([s.dInner, H]), ab: U([s.vHeads, H]), gout: U([H, s.dInner]), conv: U([convDim]),
870
+ step: U([s.kHeads, s.dState, keyDim, convDim, s.vHeads, 0, eps, { f: 1 / Math.sqrt(s.dState) }]), ng: U([s.dState, s.dInner, eps]),
871
+ router: U([c.experts, H]), topk: U([c.experts, K]),
872
+ gu: U([2 * F, H, 2 * F, 0, K]), dn: U([H, F, H, 1, K]), silu: U([F, B * K]), acc: U([H, K]),
873
+ shgu: U([SF, H]), shd: U([H, SF]), sh2: U([SF]), shgate: U([1, H]),
874
+ };
875
+ const raw = (name) => W[name].raw;
876
+ const q8 = (name, x, y, uni) => this.bind('matmulQ8T', [x, W[name].q, W[name].s, y, uni, qb]);
877
+ this.gB = { embed: this.bind('embedQ8B', [W['token_embd.weight'].q, W['token_embd.weight'].s, b.ids, b.x, u.nH, qb]) };
878
+ for (let l = 0; l < c.layers; l++) {
879
+ const n = (t) => `blk.${l}.${t}`, g = { rmsA: this.bind('rmsnormB', [b.x, raw(n('attn_norm.weight')), b.xn, u.rmsH]) };
880
+ if (this.recurrent[l]) {
881
+ Object.assign(g, {
882
+ qkv: q8(n('attn_qkv.weight'), b.xn, b.mixed, u.qkv), z: q8(n('attn_gate.weight'), b.xn, b.z, u.z),
883
+ beta: q8(n('ssm_beta.weight'), b.xn, b.bb, u.ab), alpha: q8(n('ssm_alpha.weight'), b.xn, b.al, u.ab),
884
+ conv: this.bind('gdnConvB', [b.mixed, raw(n('ssm_conv1d.weight')), this.conv[l], b.cv, u.conv, qb]),
885
+ step: this.bind('gdnStepB', [b.cv, b.bb, b.al, raw(n('ssm_dt.bias')), raw(n('ssm_a')), this.ssm[l], b.go, u.step, qb]),
886
+ norm: this.bind('gdnNormGateB', [b.go, raw(n('ssm_norm.weight')), b.z, b.gn, u.ng]),
887
+ out: q8(n('ssm_out.weight'), b.gn, b.o, u.gout),
888
+ });
889
+ } else {
890
+ Object.assign(g, {
891
+ q: q8(n('attn_q.weight'), b.xn, b.qg, u.qg), k: q8(n('attn_k.weight'), b.xn, b.k, u.kv), v: q8(n('attn_v.weight'), b.xn, b.v, u.kv),
892
+ ropeQ: this.bind('qkNormRopePB', [b.qg, raw(n('attn_q_norm.weight')), b.rope, b.q, b.gate, u.ropeQ]),
893
+ ropeK: this.bind('qkNormRopePB', [b.k, raw(n('attn_k_norm.weight')), b.rope, b.kr, b.dummy, u.ropeK]),
894
+ kv: this.bind('kvStoreB', [b.kr, b.v, this.kc[l], this.vc[l], u.kvn, qb]),
895
+ score: this.bind('attnScoreB', [b.q, this.kc[l], b.sc, u.att, qb]),
896
+ soft: this.bind('softmaxB', [b.sc, qb]),
897
+ attOut: this.bind('attnOutB', [b.sc, this.vc[l], b.att, u.attOut, qb]),
898
+ gate: this.bind('sigmoidMulB', [b.att, b.gate, u.gate, qb]),
899
+ o: q8(n('attn_output.weight'), b.att, b.o, u.o),
900
+ });
901
+ }
902
+ Object.assign(g, {
903
+ addO: this.bind('addInPlaceB', [b.x, b.o, u.nH, qb]),
904
+ rmsF: this.bind('rmsnormB', [b.x, raw(n('post_attention_norm.weight')), b.xn, u.rmsH]),
905
+ router: this.bind('matmulF32T', [b.xn, raw(n('ffn_gate_inp.weight')), b.rl, u.router]),
906
+ topk: this.bind('topkB', [b.rl, b.sel, u.topk]),
907
+ shg: q8(n('ffn_gate_shexp.weight'), b.xn, b.shg, u.shgu), shu: q8(n('ffn_up_shexp.weight'), b.xn, b.shu, u.shgu),
908
+ shact: this.bind('siluMul2B', [b.shg, b.shu, b.shact, u.sh2, qb]),
909
+ shd: q8(n('ffn_down_shexp.weight'), b.shact, b.sh, u.shd),
910
+ shgate: this.bind('matmulF32T', [b.xn, raw(n('ffn_gate_inp_shexp.weight')), b.sg, u.shgate]),
911
+ gu: this.bind('expertQ8B', [b.xn, this.pool.guQ, this.pool.guS, b.gu, u.gu, b.slots]),
912
+ silu: this.bind('siluMulMoe', [b.gu, b.act, u.silu]),
913
+ dn: this.bind('expertQ8B', [b.act, this.pool.dQ, this.pool.dS, b.dn, u.dn, b.slots]),
914
+ acc: this.bind('moeAccumSharedB', [b.dn, b.sel, b.sh, b.sg, b.x, u.acc, qb]),
915
+ });
916
+ this.layers[l].B = g;
917
+ }
918
+ }
919
+
920
+ // Layer l over the chunk's T tokens up to its routing (the batched twin of encodeAttention).
921
+ encodeTrunkB(pass, l, T, S) {
922
+ const c = this.cfg, s = c.ssm, g = this.layers[l].B, d = (nm, gr, x, y, z) => this.dispatch(pass, nm, gr, x, y, z);
923
+ const H = c.hidden, QN = c.heads * c.headDim, kvn = c.kvHeads * c.headDim, convDim = 2 * s.kHeads * s.dState + s.dInner;
924
+ const tt = Math.ceil(T / 16), mm = (n) => Math.ceil(n / 16);
925
+ d('rmsnormB', g.rmsA, T);
926
+ if (this.recurrent[l]) {
927
+ d('matmulQ8T', g.qkv, mm(convDim), tt); d('matmulQ8T', g.z, mm(s.dInner), tt);
928
+ d('matmulQ8T', g.beta, mm(s.vHeads), tt); d('matmulQ8T', g.alpha, mm(s.vHeads), tt);
929
+ d('gdnConvB', g.conv, Math.ceil(convDim / 256));
930
+ d('gdnStepB', g.step, s.vHeads);
931
+ d('gdnNormGateB', g.norm, s.vHeads, T);
932
+ d('matmulQ8T', g.out, mm(H), tt);
933
+ } else {
934
+ d('matmulQ8T', g.q, mm(2 * QN), tt); d('matmulQ8T', g.k, mm(kvn), tt); d('matmulQ8T', g.v, mm(kvn), tt);
935
+ d('qkNormRopePB', g.ropeQ, c.heads, T); d('qkNormRopePB', g.ropeK, c.kvHeads, T);
936
+ d('kvStoreB', g.kv, Math.ceil(kvn / 256), T);
937
+ d('attnScoreB', g.score, Math.ceil(c.heads * T * S / 256));
938
+ d('softmaxB', g.soft, c.heads * T);
939
+ d('attnOutB', g.attOut, Math.ceil(QN / 256), T);
940
+ d('sigmoidMulB', g.gate, Math.ceil(QN / 256), T);
941
+ d('matmulQ8T', g.o, mm(H), tt);
942
+ }
943
+ d('addInPlaceB', g.addO, Math.ceil(H / 256), T);
944
+ d('rmsnormB', g.rmsF, T);
945
+ d('matmulF32T', g.router, c.experts, T);
946
+ d('topkB', g.topk, T);
947
+ d('matmulQ8T', g.shg, mm(c.shexpFf), tt); d('matmulQ8T', g.shu, mm(c.shexpFf), tt);
948
+ d('siluMul2B', g.shact, Math.ceil(c.shexpFf / 256), T);
949
+ d('matmulQ8T', g.shd, mm(H), tt);
950
+ d('matmulF32T', g.shgate, 1, T);
951
+ }
952
+
953
+ encodeExpertsB(pass, l, T) {
954
+ const c = this.cfg, g = this.layers[l].B, d = (nm, gr, x, y, z) => this.dispatch(pass, nm, gr, x, y, z);
955
+ d('expertQ8B', g.gu, Math.ceil(2 * c.expertFf / 4), 1, T * c.topK);
956
+ d('siluMulMoe', g.silu, Math.ceil(T * c.topK * c.expertFf / 256));
957
+ d('expertQ8B', g.dn, Math.ceil(c.hidden / 4), 1, T * c.topK);
958
+ d('moeAccumSharedB', g.acc, Math.ceil(c.hidden / 256), T);
959
+ }
960
+
961
+ // Feeds T = ids.length tokens at this.position through every layer: one routing readback per
962
+ // layer, the union of the chunk's experts loaded once. want: 'none' | 'argmax' | 'logits' for
963
+ // the last token, as in step().
964
+ async prefillChunk(ids, want = 'none') {
965
+ const c = this.cfg, dev = this.device, K = c.topK, H = c.hidden, T = ids.length, pos0 = this.position, S = pos0 + T;
966
+ if (T > this.chunkTokens) throw new Error(`chunk of ${T} > ${this.chunkTokens}`);
967
+ if (S > this.maxCtx) throw new Error(`context full (${this.maxCtx} tokens)`);
968
+ const tStart = performance.now();
969
+ dev.queue.writeBuffer(this.a.stop, 0, new Uint32Array(4)); // the head below uses guarded kernels
970
+ dev.queue.writeBuffer(this.b.ids, 0, Uint32Array.from(ids));
971
+ dev.queue.writeBuffer(this.qb, 0, new Uint32Array([T, pos0, S, 0]));
972
+ const nRot = c.ropeDims, half = nRot / 2, r = new Float32Array(T * nRot), stepR = Math.fround(Math.pow(c.ropeTheta, -2 / nRot));
973
+ for (let t = 0; t < T; t++) {
974
+ let th = Math.fround(pos0 + t);
975
+ for (let i = 0; i < half; i++) { r[t * nRot + i] = Math.cos(th); r[t * nRot + half + i] = Math.sin(th); th = Math.fround(th * stepR); }
976
+ }
977
+ dev.queue.writeBuffer(this.b.rope, 0, r);
978
+ const selBytes = T * 2 * K * 4;
979
+ let enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
980
+ this.dispatch(pass, 'embedQ8B', this.gB.embed, Math.ceil(H / 256), T);
981
+ this.encodeTrunkB(pass, 0, T, S);
982
+ pass.end();
983
+ enc.copyBufferToBuffer(this.b.sel, 0, this.rbSelB, 0, selBytes);
984
+ dev.queue.submit([enc.finish()]);
985
+ for (let l = 0; l < c.layers; l++) {
986
+ const t0 = performance.now();
987
+ await this.rbSelB.mapAsync(1, 0, selBytes);
988
+ this.counters.gpuWaitMs += performance.now() - t0;
989
+ const sel = new Uint32Array(this.rbSelB.getMappedRange(0, selBytes).slice(0));
990
+ this.rbSelB.unmap();
991
+ const pairs = new Uint32Array(T * K);
992
+ for (let t = 0; t < T; t++) for (let j = 0; j < K; j++) pairs[t * K + j] = sel[t * 2 * K + j];
993
+ const unique = [...new Set(pairs)];
994
+ const tq = performance.now();
995
+ const slotsU = await this.xs.ensure(l, unique);
996
+ this.counters.ensureMs += performance.now() - tq;
997
+ const slotOf = new Map(unique.map((e, i) => [e, slotsU[i]]));
998
+ const pairSlots = pairs.map((e) => slotOf.get(e));
999
+ dev.queue.writeBuffer(this.b.slots, 0, pairSlots);
1000
+ enc = dev.createCommandEncoder(); pass = enc.beginComputePass();
1001
+ this.encodeExpertsB(pass, l, T);
1002
+ if (l + 1 < c.layers) {
1003
+ this.encodeTrunkB(pass, l + 1, T, S);
1004
+ pass.end();
1005
+ enc.copyBufferToBuffer(this.b.sel, 0, this.rbSelB, 0, selBytes);
1006
+ } else {
1007
+ pass.end();
1008
+ enc.copyBufferToBuffer(this.b.x, (T - 1) * H * 4, this.a.x, 0, H * 4);
1009
+ if (want !== 'none') {
1010
+ pass = enc.beginComputePass();
1011
+ this.dispatch(pass, 'rmsnorm', this.g.rmsOut, 1);
1012
+ this.dispatch(pass, 'matmulQ8', this.g.lm, Math.ceil(c.vocab / 4));
1013
+ if (want === 'argmax') this.dispatch(pass, 'argmax', this.g.am, 1);
1014
+ pass.end();
1015
+ if (want === 'argmax') enc.copyBufferToBuffer(this.a.am, 0, this.rbArg, 0, 4);
1016
+ if (want === 'logits') enc.copyBufferToBuffer(this.a.logits, 0, this.rbLogits, 0, c.vocab * 4);
1017
+ }
1018
+ }
1019
+ dev.queue.submit([enc.finish()]);
1020
+ this.xs.release(slotsU);
1021
+ }
1022
+ this.position += T;
1023
+ this.cached.push(...ids);
1024
+ let result = null;
1025
+ if (want === 'argmax') {
1026
+ await this.rbArg.mapAsync(1, 0, 4);
1027
+ result = new Uint32Array(this.rbArg.getMappedRange(0, 4))[0];
1028
+ this.rbArg.unmap();
1029
+ } else if (want === 'logits') {
1030
+ await this.rbLogits.mapAsync(1);
1031
+ result = new Float32Array(this.rbLogits.getMappedRange().slice(0));
1032
+ this.rbLogits.unmap();
1033
+ } else {
1034
+ await dev.queue.onSubmittedWorkDone();
1035
+ }
1036
+ this.counters.tokens += T;
1037
+ this.counters.wallMs += performance.now() - tStart;
1038
+ return result;
1039
+ }
1040
+
1041
+ // Layer l up to its routing: the trunk (DeltaNet or gated attention), residual, post-norm,
1042
+ // router + top-k (and the prefetch guesses), and the shared expert, which needs no routing.
1043
+ encodeAttention(pass, l, seqLen) {
1044
+ const c = this.cfg, s = c.ssm, g = this.layers[l], d = (nm, gr, x, y, z) => this.dispatch(pass, nm, gr, x, y, z);
1045
+ const H = c.hidden, QN = c.heads * c.headDim, kvn = c.kvHeads * c.headDim, convDim = 2 * s.kHeads * s.dState + s.dInner;
1046
+ d('rmsnorm', g.rmsA, 1);
1047
+ if (g.recur) {
1048
+ d('matmulQ8', g.qkv, Math.ceil(convDim / 4)); d('matmulQ8', g.z, Math.ceil(s.dInner / 4));
1049
+ d('matmulQ8', g.beta, Math.ceil(s.vHeads / 4)); d('matmulQ8', g.alpha, Math.ceil(s.vHeads / 4));
1050
+ d('gdnConv', g.conv, Math.ceil(convDim / 256));
1051
+ d('gdnStep', g.step, s.vHeads);
1052
+ d('gdnNormGate', g.norm, s.vHeads);
1053
+ d('matmulQ8', g.out, Math.ceil(H / 4));
1054
+ } else {
1055
+ d('matmulQ8', g.q, Math.ceil(2 * QN / 4)); d('matmulQ8', g.k, Math.ceil(kvn / 4)); d('matmulQ8', g.v, Math.ceil(kvn / 4));
1056
+ d('qkNormRopeP', g.ropeQ, c.heads); d('qkNormRopeP', g.ropeK, c.kvHeads);
1057
+ d('kvStore', g.kv, Math.ceil(kvn / 256));
1058
+ d('attnScore', g.score, Math.ceil(c.heads * seqLen / 256));
1059
+ d('softmax', g.soft, c.heads);
1060
+ d('attnOut', g.attOut, Math.ceil(QN / 256));
1061
+ d('sigmoidMul', g.gate, Math.ceil(QN / 256));
1062
+ d('matmulQ8', g.o, Math.ceil(H / 4));
1063
+ }
1064
+ d('addInPlace', g.addO, Math.ceil(H / 256));
1065
+ if (this.prefetch) for (const p of g.pf) { d('rmsnorm', p.rms, 1); d('matmulF32', p.router, c.experts); d('topk', p.topk, 1); }
1066
+ d('rmsnorm', g.rmsF, 1);
1067
+ d('matmulF32', g.router, c.experts);
1068
+ d('topk', g.topk, 1);
1069
+ d('matmulQ8', g.shg, Math.ceil(c.shexpFf / 4)); d('matmulQ8', g.shu, Math.ceil(c.shexpFf / 4));
1070
+ d('siluMulMoe', g.shsilu, Math.ceil(c.shexpFf / 256));
1071
+ d('matmulQ8', g.shd, Math.ceil(H / 4));
1072
+ d('matmulF32', g.shgate, 1);
1073
+ }
1074
+
1075
+ encodeExperts(pass, l) {
1076
+ const c = this.cfg, g = this.layers[l], d = (nm, gr, x, y, z) => this.dispatch(pass, nm, gr, x, y, z);
1077
+ d('expertQ8', g.gu, Math.ceil(2 * c.expertFf / 4), 1, c.topK);
1078
+ d('siluMulMoe', g.silu, Math.ceil(c.topK * c.expertFf / 256));
1079
+ d('expertQ8', g.dn, Math.ceil(c.hidden / 4), 1, c.topK);
1080
+ d('moeAccumShared', g.acc, Math.ceil(c.hidden / 256));
1081
+ }
1082
+
1083
+ // The pool's residency as the route kernel reads it: slot of (layer, expert), or ABSENT. Built
1084
+ // from the streamer's slot keys, which are set only once an expert's upload has been submitted.
1085
+ // Queue order keeps it safe: an upload into a slot runs after every submit made before it.
1086
+ writeRouteMap() {
1087
+ const map = this.routeMap, key = this.xs.slotKey;
1088
+ map.fill(ABSENT);
1089
+ for (let s = 0; s < key.length; s++) if (key[s] >= 0) map[key[s]] = s;
1090
+ this.device.queue.writeBuffer(this.a.rmap, 0, map);
1091
+ }
1092
+
1093
+ // One token with routing on the GPU. Each submit runs up to routeWindow layers and routes the
1094
+ // layer after them; the route kernel stops it early at the first layer with an expert outside
1095
+ // the pool. Either way the CPU then ensures the next layer's experts (loading any absent ones and
1096
+ // prefetching the guesses for the layers after it) and resumes at that layer's experts:
1097
+ // everything it computed before routing (x, xn, sel, shared expert) is still in place, since
1098
+ // every kernel after a stop returns without writing.
1099
+ async step(token, want = 'argmax') {
1100
+ if (!this.gpuRouting) return super.step(token, want);
1101
+ const c = this.cfg, dev = this.device, K = c.topK, L = c.layers, seg = this.seg;
1102
+ if (this.position >= this.maxCtx) throw new Error(`context full (${this.maxCtx} tokens)`);
1103
+ const tStart = performance.now();
1104
+ const pos = this.position, seqLen = pos + 1;
1105
+ this.writeTokenUniforms(token, pos);
1106
+ // resume: -1 starts at the embedding; l ≥ 0 resumes at layer l's experts (its trunk already ran,
1107
+ // and a DeltaNet trunk must not run twice: it advances the recurrent state).
1108
+ let resume = -1, slots = null, result = null;
1109
+ for (;;) {
1110
+ const te = performance.now();
1111
+ this.writeRouteMap();
1112
+ dev.queue.writeBuffer(this.a.stop, 0, new Uint32Array(4));
1113
+ if (slots) dev.queue.writeBuffer(this.a.slots, 0, slots);
1114
+ const enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
1115
+ if (resume < 0) {
1116
+ this.dispatch(pass, 'embedQ8', this.g.embed, Math.ceil(c.hidden / 256));
1117
+ this.encodeAttention(pass, 0, seqLen);
1118
+ this.dispatch(pass, 'route', this.layers[0].route, 1);
1119
+ }
1120
+ const first = Math.max(0, resume), last = Math.min(L - 1, first + this.routeWindow - 1);
1121
+ for (let l = first; l <= last; l++) {
1122
+ this.encodeExperts(pass, l);
1123
+ if (l + 1 < L) { this.encodeAttention(pass, l + 1, seqLen); this.dispatch(pass, 'route', this.layers[l + 1].route, 1); }
1124
+ }
1125
+ if (last === L - 1 && want !== 'none') {
1126
+ this.dispatch(pass, 'rmsnorm', this.g.rmsOut, 1);
1127
+ this.dispatch(pass, 'matmulQ8', this.g.lm, Math.ceil(c.vocab / 4));
1128
+ if (want === 'argmax') this.dispatch(pass, 'argmax', this.g.am, 1);
1129
+ }
1130
+ pass.end();
1131
+ enc.copyBufferToBuffer(this.a.stop, 0, this.rbSeg, 0, 16);
1132
+ enc.copyBufferToBuffer(this.a.sel, 0, this.rbSeg, seg.sel, 8 * K);
1133
+ if (this.prefetch) enc.copyBufferToBuffer(this.a.psel, 0, this.rbSeg, seg.psel, 256 * this.lookahead);
1134
+ enc.copyBufferToBuffer(this.a.used, 0, this.rbSeg, seg.used, L * K * 4);
1135
+ if (last === L - 1 && want === 'argmax') enc.copyBufferToBuffer(this.a.am, 0, this.rbSeg, seg.am, 4);
1136
+ if (last === L - 1 && want === 'logits') enc.copyBufferToBuffer(this.a.logits, 0, this.rbLogits, 0, c.vocab * 4);
1137
+ dev.queue.submit([enc.finish()]);
1138
+ if (slots) { this.xs.release(slots); slots = null; }
1139
+ this.counters.encodeMs += performance.now() - te;
1140
+ this.counters.syncs++;
1141
+ this.counters.layersEncoded += last - first + 1;
1142
+
1143
+ const tw = performance.now();
1144
+ await this.rbSeg.mapAsync(1, 0, seg.bytes);
1145
+ this.counters.gpuWaitMs += performance.now() - tw;
1146
+ const u = new Uint32Array(this.rbSeg.getMappedRange(0, seg.bytes).slice(0));
1147
+ this.rbSeg.unmap();
1148
+ const stopL = u[0] ? u[0] - 1 : last + 1;
1149
+ // Layers this submit routed and ran on the GPU: refresh them in the LRU. (A resumed layer went
1150
+ // through ensure() already.)
1151
+ for (let l = resume + 1; l < stopL; l++) this.xs.touchUsed(l, u.subarray(seg.used / 4 + l * K, seg.used / 4 + (l + 1) * K));
1152
+ if (stopL === L) {
1153
+ if (want === 'argmax') result = u[seg.am / 4];
1154
+ break;
1155
+ }
1156
+ const ids = u.slice(seg.sel / 4, seg.sel / 4 + K);
1157
+ const tq = performance.now();
1158
+ const ensuring = this.xs.ensure(stopL, ids);
1159
+ if (this.prefetch) {
1160
+ const n = this.layers[stopL].pf.length;
1161
+ for (let i = 0; i < n; i++) this.xs.prefetch(stopL + 1 + i, u.subarray(seg.psel / 4 + 64 * i, seg.psel / 4 + 64 * i + K));
1162
+ }
1163
+ slots = await ensuring;
1164
+ this.counters.ensureMs += performance.now() - tq;
1165
+ resume = stopL;
1166
+ }
1167
+ if (want === 'logits') {
1168
+ const tw = performance.now();
1169
+ await this.rbLogits.mapAsync(1);
1170
+ this.counters.gpuWaitMs += performance.now() - tw;
1171
+ result = new Float32Array(this.rbLogits.getMappedRange().slice(0));
1172
+ this.rbLogits.unmap();
1173
+ }
1174
+ this.position++;
1175
+ this.cached.push(token);
1176
+ this.counters.tokens++;
1177
+ this.counters.wallMs += performance.now() - tStart;
1178
+ return result;
1179
+ }
1180
+
1181
+ resetCounters() { super.resetCounters(); this.counters.syncs = 0; this.counters.layersEncoded = 0; }
1182
+ // layersEncodedPerToken − layers = expert blocks dispatched after a stop (no-ops).
1183
+ stats() {
1184
+ const n = Math.max(1, this.counters.tokens);
1185
+ return { ...super.stats(), gpuRouting: this.gpuRouting, routeWindow: this.routeWindow, syncsPerToken: this.counters.syncs / n, layersEncodedPerToken: this.counters.layersEncoded / n };
1186
+ }
1187
+
1188
+ // RoPE angles as ggml's CPU rope cache forms them for text positions (all M-RoPE sections at
1189
+ // the token position): θ = pos, multiplied by base^(-2/nRot) once per pair, in f32.
1190
+ writeTokenUniforms(token, pos) {
1191
+ const c = this.cfg, nRot = c.ropeDims, half = nRot / 2;
1192
+ this.device.queue.writeBuffer(this.tok, 0, new Uint32Array([token, pos, pos + 1, 0]));
1193
+ const r = new Float32Array(nRot), step = Math.fround(Math.pow(c.ropeTheta, -2 / nRot));
1194
+ let th = Math.fround(pos);
1195
+ for (let i = 0; i < half; i++) { r[i] = Math.cos(th); r[half + i] = Math.sin(th); th = Math.fround(th * step); }
1196
+ this.device.queue.writeBuffer(this.a.rope, 0, r);
1197
+ }
1198
+
1199
+ // Qwen3.5/3.6's template opens the reasoning block in the generation prompt when thinking is on.
1200
+ chatPrompt(messages, { enableThinking = true } = {}) {
1201
+ let p = '';
1202
+ for (const m of messages) p += `<|im_start|>${m.role}\n${m.content}<|im_end|>\n`;
1203
+ return p + '<|im_start|>assistant\n' + (enableThinking ? '<think>\n' : '<think>\n\n</think>\n\n');
1204
+ }
1205
+
1206
+ // The recurrent state cannot rewind to an earlier position, so any change before the cached
1207
+ // end restarts from token 0.
1208
+ reset() {
1209
+ super.reset();
1210
+ if (!this.conv) return;
1211
+ const enc = this.device.createCommandEncoder();
1212
+ for (let l = 0; l < this.cfg.layers; l++) if (this.recurrent[l]) { enc.clearBuffer(this.conv[l]); enc.clearBuffer(this.ssm[l]); }
1213
+ this.device.queue.submit([enc.finish()]);
1214
+ }
1215
+
1216
+ async prefill(ids, want = 'argmax') {
1217
+ let common = 0;
1218
+ while (common < ids.length - 1 && common < this.cached.length && this.cached[common] === ids[common]) common++;
1219
+ if (common < this.cached.length) this.reset();
1220
+ const rest = ids.slice(this.cached.length);
1221
+ if (!this.batchPrefill || rest.length < 2) return super.prefill(ids, want);
1222
+ let r = null;
1223
+ for (let i = 0; i < rest.length; i += this.chunkTokens) {
1224
+ r = await this.prefillChunk(rest.slice(i, i + this.chunkTokens), i + this.chunkTokens >= rest.length ? want : 'none');
1225
+ }
1226
+ return r;
1227
+ }
1228
+ }
1229
+
1230
+ // For kernel tests (scripts): the WGSL this module adds.
1231
+ export { KERNELS as QWEN35_KERNELS, BATCH_KERNELS as QWEN35_BATCH_KERNELS, SG_KERNELS as QWEN35_SG_KERNELS };
engines/qwen3_moe_ssd.js ADDED
@@ -0,0 +1,1035 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* qwen3_moe_ssd.js — Qwen3-MoE (Qwen3-30B-A3B) in a browser tab, larger than RAM: the dense
2
+ * weights, routers and KV cache live on the GPU; the 6,144 routed experts (≈31 GB at Q8_0)
3
+ * live in OPFS and stream into a GPU slot pool as the router asks for them.
4
+ *
5
+ * Rung 2a of the SSD-streaming ladder (plan/pending.md). From-scratch WebGPU engine; the
6
+ * kernels are ports of browser-big-fast-lab's custom-kernels engine (an Apache-2.0 fork of
7
+ * tylerstraub/gemma4-webgpu): in-shader Q8_0 GEMV (int8 + f16 scale per 32-block, f32
8
+ * accumulate), GPU softmax + top-k router, slot-indexed batched expert GEMVs, MoE combine.
9
+ *
10
+ * const m = await Qwen3MoeSsd.load('Qwen/Qwen3-30B-A3B-GGUF', { fetch, onProgress, poolBytes });
11
+ * for await (const { text } of m.generate(messages, { maxNewTokens })) …
12
+ *
13
+ * Load = (first time only) ingest the GGUF into OPFS in the engine's own layout, then upload
14
+ * the dense part to the GPU. OPFS layout under localmind-ssd/<key>/:
15
+ * header.bin the GGUF header bytes (metadata + tokenizer), re-parsed at every load
16
+ * dense.bin every non-expert tensor; Q8_0 split losslessly into int8 + f16-scale planes
17
+ * experts.bin one 5,013,504-byte record per (layer, expert), gate‖up‖down contiguous
18
+ * manifest.json source (repo, revision, sha256, size) + every offset above
19
+ * Q8_0 is never re-quantized: each block's f16 scale and 32 int8 values are copied bit-exact,
20
+ * so the engine computes the same dot products llama.cpp does, up to summation order.
21
+ */
22
+
23
+ import { OpfsReaderPool, readOpfsText } from '../src/opfs-reader.js';
24
+ import { ExpertStreamer } from '../src/expert-stream.js';
25
+ // The store layer is diskformer.js's (github.com/NakliTechie/diskformer.js), byte-identical copies: gguf.js, ingest.js.
26
+ import { GGML, Q8_BLOCK, Q4_BLOCK, Q6K_BLOCK, parseGguf, tensorBytes } from '../src/gguf.js';
27
+ import { ingestGguf as ingestStore, ingestProgress, storeKey, removeStore } from '../src/ingest.js';
28
+ export { GGML, Q4_BLOCK, Q6K_BLOCK, parseGguf, tensorBytes, splitQ4 } from '../src/gguf.js';
29
+ export { ingestProgress } from '../src/ingest.js';
30
+
31
+ export const QWEN3_30B_A3B = {
32
+ repo: 'Qwen/Qwen3-30B-A3B-GGUF',
33
+ file: 'Qwen3-30B-A3B-Q8_0.gguf',
34
+ revision: 'e4d4bafdfb96a411a163846265362aceb0b9c63a',
35
+ sha256: '4ad960d180b16f56024f5b704697e5dd5b0837167c2e515ef0569abfc599743c',
36
+ size: 32483931648,
37
+ };
38
+ const FORMAT = 'localmind-qwen3moe-ssd/1';
39
+ const OPFS_ROOT = 'localmind-ssd';
40
+
41
+ // ── GGUF ────────────────────────────────────────────────────────────────────
42
+ // qwen3moe (rung 2a) and qwen35moe (Qwen3.5/3.6 MoE, rung 2b: qwen35_moe_ssd.js) share the
43
+ // expert layout; qwen35moe adds the Gated DeltaNet (ssm.*), partial RoPE and a shared expert.
44
+ export function configFromGguf(kv) {
45
+ const a = kv['general.architecture'];
46
+ if (a !== 'qwen3moe' && a !== 'qwen35moe') throw new Error(`expected a qwen3moe or qwen35moe GGUF, got ${a}`);
47
+ const g = (k) => kv[`${a}.${k}`];
48
+ const extra = a !== 'qwen35moe' ? {} : {
49
+ ropeDims: g('rope.dimension_count'), shexpFf: g('expert_shared_feed_forward_length'),
50
+ attnInterval: g('full_attention_interval') || 4,
51
+ ssm: { dConv: g('ssm.conv_kernel'), dInner: g('ssm.inner_size'), dState: g('ssm.state_size'), vHeads: g('ssm.time_step_rank'), kHeads: g('ssm.group_count') },
52
+ };
53
+ return {
54
+ ...extra,
55
+ arch: a,
56
+ layers: g('block_count'), hidden: g('embedding_length'),
57
+ heads: g('attention.head_count'), kvHeads: g('attention.head_count_kv'),
58
+ headDim: g('attention.key_length') || g('embedding_length') / g('attention.head_count'),
59
+ experts: g('expert_count'), topK: g('expert_used_count'), expertFf: g('expert_feed_forward_length'),
60
+ ropeTheta: g('rope.freq_base'), eps: g('attention.layer_norm_rms_epsilon'),
61
+ contextLength: g('context_length'),
62
+ vocab: kv['tokenizer.ggml.tokens'].length,
63
+ bos: kv['tokenizer.ggml.bos_token_id'], eos: kv['tokenizer.ggml.eos_token_id'],
64
+ };
65
+ }
66
+
67
+ // The engine's OPFS layout, derived from the header alone (ingest and load agree on it).
68
+ export function planLayout(gguf) {
69
+ const cfg = configFromGguf(gguf.kv);
70
+ const H = cfg.hidden, F = cfg.expertFf;
71
+ const qBytes = F * H; // one projection's int8 plane
72
+ const sBytes = (F * H / 32) * 2; // its f16 scale plane
73
+ const parts = {
74
+ guQ: { off: 0, bytes: 2 * qBytes }, // gate rows 0..F-1, up rows F..2F-1
75
+ guS: { off: 2 * qBytes, bytes: 2 * sBytes },
76
+ dQ: { off: 2 * qBytes + 2 * sBytes, bytes: qBytes },
77
+ dS: { off: 3 * qBytes + 2 * sBytes, bytes: sBytes },
78
+ };
79
+ const record = 3 * qBytes + 3 * sBytes;
80
+ const dense = {};
81
+ let off = 0;
82
+ const take = (n) => { const o = off; off += Math.ceil(n / 256) * 256; return o; };
83
+ for (const t of gguf.tensors) {
84
+ if (/_exps\.weight$/.test(t.name)) continue;
85
+ const n = t.dims.reduce((a, b) => a * b, 1);
86
+ if (t.type === GGML.Q8_0) {
87
+ dense[t.name] = { type: 'q8', dims: t.dims, q: { off: take(n), bytes: n }, s: { off: take(n / 16), bytes: n / 16 } };
88
+ } else if (t.type === GGML.F32) {
89
+ dense[t.name] = { type: 'f32', dims: t.dims, raw: { off: take(n * 4), bytes: n * 4 } };
90
+ } else throw new Error(`dense tensor ${t.name}: unsupported type ${t.type}`);
91
+ }
92
+ return {
93
+ config: cfg,
94
+ experts: { file: 'experts.bin', record, parts, layers: cfg.layers, perLayer: cfg.experts, bytes: record * cfg.layers * cfg.experts },
95
+ dense: { file: 'dense.bin', bytes: off, tensors: dense },
96
+ };
97
+ }
98
+
99
+ // ── Ingest: GGUF (over HTTP Range) → OPFS, in the engine layout ─────────────
100
+ // Every byte range of the GGUF the engine keeps, as units small enough to transform in
101
+ // memory: one unit per expert slice, row bands for dense Q8_0, whole F32 tensors.
102
+ function planUnits(gguf, layout) {
103
+ const units = [];
104
+ const { record, parts, perLayer } = layout.experts;
105
+ const qBytes = parts.dQ.bytes, sBytes = parts.dS.bytes;
106
+ for (const t of gguf.tensors) {
107
+ const abs = gguf.dataStart + t.offset;
108
+ const m = /^blk\.(\d+)\.ffn_(gate|up|down)_exps\.weight$/.exec(t.name);
109
+ if (m) {
110
+ const layer = Number(m[1]), which = m[2];
111
+ if (t.type !== GGML.Q8_0) throw new Error(`${t.name}: experts must be Q8_0`);
112
+ const slice = tensorBytes(t) / perLayer;
113
+ const qOff = which === 'gate' ? parts.guQ.off : which === 'up' ? parts.guQ.off + qBytes : parts.dQ.off;
114
+ const sOff = which === 'gate' ? parts.guS.off : which === 'up' ? parts.guS.off + sBytes : parts.dS.off;
115
+ for (let e = 0; e < perLayer; e++) {
116
+ const rec = (layer * perLayer + e) * record;
117
+ units.push({ src: abs + e * slice, len: slice, file: 'experts', q: rec + qOff, s: rec + sOff });
118
+ }
119
+ continue;
120
+ }
121
+ const d = layout.dense.tensors[t.name];
122
+ if (d.type === 'f32') { units.push({ src: abs, len: d.raw.bytes, file: 'dense', raw: d.raw.off }); continue; }
123
+ const cols = t.dims[0], rows = d.q.bytes / cols, rowBytes = (cols / 32) * Q8_BLOCK;
124
+ const band = Math.max(1, Math.floor((8 << 20) / rowBytes));
125
+ for (let r = 0; r < rows; r += band) {
126
+ const nr = Math.min(band, rows - r);
127
+ units.push({ src: abs + r * rowBytes, len: nr * rowBytes, file: 'dense', q: d.q.off + r * cols, s: d.s.off + r * cols / 16 });
128
+ }
129
+ }
130
+ units.sort((a, b) => a.src - b.src);
131
+ return units;
132
+ }
133
+
134
+ // The ingest is diskformer's ingestGguf with this engine's store folder and layout tag (kept, so stores written
135
+ // before the split stay valid); `plan` lets another architecture supply its own layout (Gemma 4: gemma4_moe_ssd.js).
136
+ export function ingestGguf(opts) {
137
+ return ingestStore({ ...opts, root: opts.root || OPFS_ROOT, format: FORMAT, plan: opts.plan || { layout: planLayout, units: planUnits } });
138
+ }
139
+
140
+ // ── Tokenizer: byte-level BPE from the GGUF vocab (tokenizer.ggml.model = gpt2, pre = qwen2) ──
141
+ // Same pre-tokenizer regex llama.cpp uses for LLAMA_VOCAB_PRE_TYPE_QWEN2.
142
+ const QWEN2_PRE = /(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
143
+ // LLAMA_VOCAB_PRE_TYPE_QWEN35 (Qwen3.5/3.6): combining marks (\p{M}) count as part of a word.
144
+ const QWEN35_PRE = /(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
145
+
146
+ function byteToUnicode() {
147
+ const bs = [];
148
+ for (let i = 33; i <= 126; i++) bs.push(i);
149
+ for (let i = 161; i <= 172; i++) bs.push(i);
150
+ for (let i = 174; i <= 255; i++) bs.push(i);
151
+ const cs = bs.slice();
152
+ let n = 0;
153
+ for (let b = 0; b < 256; b++) if (!bs.includes(b)) { bs.push(b); cs.push(256 + n); n++; }
154
+ const enc = new Array(256), dec = new Map();
155
+ bs.forEach((b, i) => { enc[b] = String.fromCodePoint(cs[i]); dec.set(String.fromCodePoint(cs[i]), b); });
156
+ return { enc, dec };
157
+ }
158
+
159
+ export class BpeTokenizer {
160
+ constructor(kv) {
161
+ this.tokens = kv['tokenizer.ggml.tokens'];
162
+ this.preRe = kv['tokenizer.ggml.pre'] === 'qwen35' ? QWEN35_PRE : QWEN2_PRE;
163
+ const types = kv['tokenizer.ggml.token_type'] || [];
164
+ this.ids = new Map(this.tokens.map((t, i) => [t, i]));
165
+ this.ranks = new Map((kv['tokenizer.ggml.merges'] || []).map((m, i) => [m, i]));
166
+ this.special = [];
167
+ this.isSpecial = new Uint8Array(this.tokens.length);
168
+ for (let i = 0; i < this.tokens.length; i++) if (types[i] === 3 || types[i] === 4) { this.special.push(this.tokens[i]); this.isSpecial[i] = 1; }
169
+ this.special.sort((a, b) => b.length - a.length);
170
+ const esc = (s) => s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
171
+ this.specialRe = this.special.length ? new RegExp(`(${this.special.map(esc).join('|')})`) : null;
172
+ const { enc, dec } = byteToUnicode();
173
+ this.byteEnc = enc; this.byteDec = dec;
174
+ this.cache = new Map();
175
+ this.utf8 = new TextEncoder();
176
+ }
177
+ bpe(word) {
178
+ const hit = this.cache.get(word);
179
+ if (hit) return hit;
180
+ let parts = Array.from(word);
181
+ while (parts.length > 1) {
182
+ let best = -1, bestRank = Infinity;
183
+ for (let i = 0; i < parts.length - 1; i++) {
184
+ const r = this.ranks.get(parts[i] + ' ' + parts[i + 1]);
185
+ if (r !== undefined && r < bestRank) { bestRank = r; best = i; }
186
+ }
187
+ if (best < 0) break;
188
+ parts = [...parts.slice(0, best), parts[best] + parts[best + 1], ...parts.slice(best + 2)];
189
+ }
190
+ const ids = parts.map((p) => {
191
+ const id = this.ids.get(p);
192
+ if (id === undefined) throw new Error(`BPE: no token for ${JSON.stringify(p)}`);
193
+ return id;
194
+ });
195
+ if (this.cache.size < 50000) this.cache.set(word, ids);
196
+ return ids;
197
+ }
198
+ // parseSpecial: special-token text (e.g. <|im_start|>) maps to its single id.
199
+ encode(text, { parseSpecial = true } = {}) {
200
+ const out = [];
201
+ const pieces = parseSpecial && this.specialRe ? text.split(this.specialRe) : [text];
202
+ for (let i = 0; i < pieces.length; i++) {
203
+ const piece = pieces[i];
204
+ if (!piece) continue;
205
+ if (parseSpecial && i % 2 === 1) { out.push(this.ids.get(piece)); continue; }
206
+ for (const m of piece.matchAll(this.preRe)) {
207
+ let s = '';
208
+ for (const b of this.utf8.encode(m[0])) s += this.byteEnc[b];
209
+ out.push(...this.bpe(s));
210
+ }
211
+ }
212
+ return out;
213
+ }
214
+ tokenBytes(id) {
215
+ const t = this.tokens[id];
216
+ if (this.isSpecial[id]) return this.utf8.encode(t);
217
+ const bytes = [];
218
+ for (const ch of t) { const b = this.byteDec.get(ch); if (b !== undefined) bytes.push(b); }
219
+ return Uint8Array.from(bytes);
220
+ }
221
+ decode(ids) {
222
+ const chunks = ids.map((id) => this.tokenBytes(id));
223
+ const n = chunks.reduce((a, c) => a + c.length, 0), all = new Uint8Array(n);
224
+ let o = 0; for (const c of chunks) { all.set(c, o); o += c.length; }
225
+ return new TextDecoder().decode(all);
226
+ }
227
+ }
228
+
229
+ export function chatPrompt(messages, { enableThinking = true } = {}) {
230
+ let s = '';
231
+ for (const m of messages) s += `<|im_start|>${m.role}\n${m.content}<|im_end|>\n`;
232
+ s += '<|im_start|>assistant\n';
233
+ if (!enableThinking) s += '<think>\n\n</think>\n\n';
234
+ return s;
235
+ }
236
+
237
+ // ── WGSL ────────────────────────────────────────────────────────────────────
238
+ // Activations f32 throughout; Q8_0 weights as an int8 plane (u32-packed) + an f16 scale per
239
+ // 32 values, dequantized exactly in f32 inside the GEMV; KV cache f16 (llama.cpp's default).
240
+ export const TOK = 'struct Tok { token: u32, pos: u32, seqLen: u32, pad: u32 }';
241
+ export const I8 = `fn i8at(w: u32, lane: u32) -> f32 { let b = (w >> (lane * 8u)) & 0xffu; return f32(select(i32(b), i32(b) - 256, b >= 128u)); }`;
242
+ const WGSL = {
243
+ embedQ8: `enable f16;
244
+ ${TOK}
245
+ struct P { n: u32 }
246
+ @group(0) @binding(0) var<storage, read> q: array<u32>;
247
+ @group(0) @binding(1) var<storage, read> s: array<f16>;
248
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
249
+ @group(0) @binding(3) var<uniform> tok: Tok;
250
+ @group(0) @binding(4) var<uniform> p: P;
251
+ ${I8}
252
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
253
+ let i = g.x; if (i >= p.n) { return; }
254
+ let e = tok.token * p.n + i;
255
+ y[i] = i8at(q[e >> 2u], e & 3u) * f32(s[e >> 5u]);
256
+ }`,
257
+ rmsnorm: `
258
+ struct P { n: u32, eps: f32 }
259
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
260
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
261
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
262
+ @group(0) @binding(3) var<uniform> p: P;
263
+ var<workgroup> red: array<f32, 256>;
264
+ @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
265
+ let t = l.x; var acc = 0.0;
266
+ for (var i = t; i < p.n; i += 256u) { let v = x[i]; acc += v * v; }
267
+ red[t] = acc; workgroupBarrier();
268
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
269
+ let scale = 1.0 / sqrt(red[0] / f32(p.n) + p.eps);
270
+ for (var i = t; i < p.n; i += 256u) { y[i] = (x[i] * scale) * w[i]; }
271
+ }`,
272
+ // 4 output rows per workgroup; input read once per k. Dispatch ceil(M/4) workgroups (2-D ok).
273
+ matmulQ8: `enable f16;
274
+ struct P { M: u32, N: u32 }
275
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
276
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
277
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
278
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
279
+ @group(0) @binding(4) var<uniform> p: P;
280
+ ${I8}
281
+ var<workgroup> part: array<f32, 1024>;
282
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
283
+ let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let N = p.N; let M = p.M; let nb = N / 32u;
284
+ if (m0 >= M) { return; }
285
+ let r0 = min(m0, M - 1u); let r1 = min(m0 + 1u, M - 1u); let r2 = min(m0 + 2u, M - 1u); let r3 = min(m0 + 3u, M - 1u);
286
+ let q0 = (r0 * N) >> 2u; let q1 = (r1 * N) >> 2u; let q2 = (r2 * N) >> 2u; let q3 = (r3 * N) >> 2u;
287
+ var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
288
+ for (var k = t; k < N; k += 256u) {
289
+ let xv = x[k]; let wi = k >> 2u; let ln = k & 3u; let b = k >> 5u;
290
+ a0 += i8at(qw[q0 + wi], ln) * f32(qs[r0 * nb + b]) * xv;
291
+ a1 += i8at(qw[q1 + wi], ln) * f32(qs[r1 * nb + b]) * xv;
292
+ a2 += i8at(qw[q2 + wi], ln) * f32(qs[r2 * nb + b]) * xv;
293
+ a3 += i8at(qw[q3 + wi], ln) * f32(qs[r3 * nb + b]) * xv;
294
+ }
295
+ part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
296
+ workgroupBarrier();
297
+ for (var s = 128u; s > 0u; s >>= 1u) {
298
+ if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
299
+ workgroupBarrier();
300
+ }
301
+ if (t < 4u && m0 + t < M) { y[m0 + t] = part[t * 256u]; }
302
+ }`,
303
+ matmulF32: `
304
+ struct P { M: u32, N: u32 }
305
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
306
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
307
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
308
+ @group(0) @binding(3) var<uniform> p: P;
309
+ var<workgroup> red: array<f32, 256>;
310
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
311
+ let m = wg.x; let t = l.x; var acc = 0.0;
312
+ for (var k = t; k < p.N; k += 256u) { acc += w[m * p.N + k] * x[k]; }
313
+ red[t] = acc; workgroupBarrier();
314
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
315
+ if (t == 0u) { y[m] = red[0]; }
316
+ }`,
317
+ // Per-head RMSNorm (QK-norm) then NeoX RoPE (pairs i, i+hd/2), in place. One workgroup per head.
318
+ // rope[0..hd/2) = cos, rope[hd/2..hd) = sin for this token's position (computed on the CPU).
319
+ qkNormRope: `
320
+ struct P { hd: u32, eps: f32 }
321
+ @group(0) @binding(0) var<storage, read_write> x: array<f32>;
322
+ @group(0) @binding(1) var<storage, read> w: array<f32>;
323
+ @group(0) @binding(2) var<storage, read> rope: array<f32>;
324
+ @group(0) @binding(3) var<uniform> p: P;
325
+ var<workgroup> red: array<f32, 64>;
326
+ @compute @workgroup_size(64) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
327
+ let base = wg.x * p.hd; let t = l.x; let half = p.hd / 2u;
328
+ var acc = 0.0;
329
+ for (var i = t; i < p.hd; i += 64u) { let v = x[base + i]; acc += v * v; }
330
+ red[t] = acc; workgroupBarrier();
331
+ for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
332
+ let scale = 1.0 / sqrt(red[0] / f32(p.hd) + p.eps);
333
+ for (var i = t; i < half; i += 64u) {
334
+ let a = (x[base + i] * scale) * w[i];
335
+ let b = (x[base + i + half] * scale) * w[i + half];
336
+ let c = rope[i]; let sn = rope[half + i];
337
+ x[base + i] = a * c - b * sn;
338
+ x[base + i + half] = a * sn + b * c;
339
+ }
340
+ }`,
341
+ kvStore: `enable f16;
342
+ ${TOK}
343
+ struct P { n: u32 }
344
+ @group(0) @binding(0) var<storage, read> k: array<f32>;
345
+ @group(0) @binding(1) var<storage, read> v: array<f32>;
346
+ @group(0) @binding(2) var<storage, read_write> kc: array<f16>;
347
+ @group(0) @binding(3) var<storage, read_write> vc: array<f16>;
348
+ @group(0) @binding(4) var<uniform> tok: Tok;
349
+ @group(0) @binding(5) var<uniform> p: P;
350
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
351
+ let i = g.x; if (i >= p.n) { return; }
352
+ kc[tok.pos * p.n + i] = f16(k[i]);
353
+ vc[tok.pos * p.n + i] = f16(v[i]);
354
+ }`,
355
+ attnScore: `enable f16;
356
+ ${TOK}
357
+ struct P { heads: u32, kvHeads: u32, hd: u32, scale: f32 }
358
+ @group(0) @binding(0) var<storage, read> q: array<f32>;
359
+ @group(0) @binding(1) var<storage, read> kc: array<f16>;
360
+ @group(0) @binding(2) var<storage, read_write> sc: array<f32>;
361
+ @group(0) @binding(3) var<uniform> tok: Tok;
362
+ @group(0) @binding(4) var<uniform> p: P;
363
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
364
+ let n = tok.seqLen; let idx = g.x; if (idx >= p.heads * n) { return; }
365
+ let h = idx / n; let pos = idx % n; let kvh = h / (p.heads / p.kvHeads);
366
+ let qo = h * p.hd; let ko = pos * p.kvHeads * p.hd + kvh * p.hd;
367
+ var acc = 0.0;
368
+ for (var d = 0u; d < p.hd; d++) { acc += q[qo + d] * f32(kc[ko + d]); }
369
+ sc[h * n + pos] = acc * p.scale;
370
+ }`,
371
+ softmax: `
372
+ ${TOK}
373
+ @group(0) @binding(0) var<storage, read_write> sc: array<f32>;
374
+ @group(0) @binding(1) var<uniform> tok: Tok;
375
+ var<workgroup> red: array<f32, 256>;
376
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
377
+ let n = tok.seqLen; let base = wg.x * n; let t = l.x;
378
+ var mx = -3.4e38;
379
+ for (var i = t; i < n; i += 256u) { mx = max(mx, sc[base + i]); }
380
+ red[t] = mx; workgroupBarrier();
381
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
382
+ let m = red[0]; workgroupBarrier();
383
+ var sum = 0.0;
384
+ for (var i = t; i < n; i += 256u) { let e = exp(sc[base + i] - m); sc[base + i] = e; sum += e; }
385
+ red[t] = sum; workgroupBarrier();
386
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
387
+ let inv = 1.0 / red[0];
388
+ for (var i = t; i < n; i += 256u) { sc[base + i] = sc[base + i] * inv; }
389
+ }`,
390
+ attnOut: `enable f16;
391
+ ${TOK}
392
+ struct P { heads: u32, kvHeads: u32, hd: u32, pad: u32 }
393
+ @group(0) @binding(0) var<storage, read> pr: array<f32>;
394
+ @group(0) @binding(1) var<storage, read> vc: array<f16>;
395
+ @group(0) @binding(2) var<storage, read_write> y: array<f32>;
396
+ @group(0) @binding(3) var<uniform> tok: Tok;
397
+ @group(0) @binding(4) var<uniform> p: P;
398
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
399
+ let idx = g.x; if (idx >= p.heads * p.hd) { return; }
400
+ let n = tok.seqLen; let h = idx / p.hd; let d = idx % p.hd; let kvh = h / (p.heads / p.kvHeads);
401
+ let stride = p.kvHeads * p.hd;
402
+ var acc = 0.0;
403
+ for (var pos = 0u; pos < n; pos++) { acc += pr[h * n + pos] * f32(vc[pos * stride + kvh * p.hd + d]); }
404
+ y[idx] = acc;
405
+ }`,
406
+ addInPlace: `
407
+ struct P { n: u32 }
408
+ @group(0) @binding(0) var<storage, read_write> x: array<f32>;
409
+ @group(0) @binding(1) var<storage, read> y: array<f32>;
410
+ @group(0) @binding(2) var<uniform> p: P;
411
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
412
+ let i = g.x; if (i < p.n) { x[i] = x[i] + y[i]; }
413
+ }`,
414
+ // Router: softmax over n logits (n ≤ 256), greedy top-k (ties → lowest index), weights
415
+ // renormalized to sum 1 (Qwen3-MoE norm_topk_prob). sel = k ids then k f32 weights (bitcast).
416
+ // One workgroup; each of the k picks is a parallel argmax over the per-thread probabilities.
417
+ // Every probability is the same f32 expression the serial version evaluated, so the picks and
418
+ // weights are the same; the serial version cost ~0.1 ms per call on an M4 Pro.
419
+ topk: `
420
+ struct P { n: u32, k: u32 }
421
+ @group(0) @binding(0) var<storage, read> lg: array<f32>;
422
+ @group(0) @binding(1) var<storage, read_write> sel: array<u32>;
423
+ @group(0) @binding(2) var<uniform> p: P;
424
+ var<workgroup> red: array<f32, 256>;
425
+ var<workgroup> bi: array<u32, 256>;
426
+ var<workgroup> gmax: f32;
427
+ var<workgroup> gsum: f32;
428
+ var<workgroup> picks: array<u32, 32>;
429
+ var<workgroup> probs: array<f32, 32>;
430
+ @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
431
+ let t = l.x; let n = p.n;
432
+ red[t] = select(-3.4e38, lg[min(t, n - 1u)], t < n); workgroupBarrier();
433
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] = max(red[t], red[t + s]); } workgroupBarrier(); }
434
+ if (t == 0u) { gmax = red[0]; } workgroupBarrier();
435
+ red[t] = select(0.0, exp(lg[min(t, n - 1u)] - gmax), t < n); workgroupBarrier();
436
+ for (var s = 128u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
437
+ if (t == 0u) { gsum = red[0]; } workgroupBarrier();
438
+ var mine = select(-1.0, exp(lg[min(t, n - 1u)] - gmax) / gsum, t < n);
439
+ for (var j = 0u; j < p.k; j++) {
440
+ red[t] = mine; bi[t] = t; workgroupBarrier();
441
+ for (var s = 128u; s > 0u; s >>= 1u) {
442
+ if (t < s) { let o = red[t + s]; if (o > red[t] || (o == red[t] && bi[t + s] < bi[t])) { red[t] = o; bi[t] = bi[t + s]; } }
443
+ workgroupBarrier();
444
+ }
445
+ if (t == 0u) { picks[j] = bi[0]; probs[j] = red[0]; }
446
+ workgroupBarrier();
447
+ if (t == picks[j]) { mine = -1.0; }
448
+ workgroupBarrier();
449
+ }
450
+ if (t == 0u) {
451
+ var wsum = 0.0;
452
+ for (var j = 0u; j < p.k; j++) { wsum += probs[j]; }
453
+ for (var j = 0u; j < p.k; j++) { sel[j] = picks[j]; sel[p.k + j] = bitcast<u32>(probs[j] / wsum); }
454
+ }
455
+ }`,
456
+ // Routed experts, all k at once: workgroup z = top-k position, slot = slots[z] picks the
457
+ // expert's rows in the pool buffer (rows_per_expert per slot). Input is shared (gate/up) or
458
+ // per-position (down). Output region z × M.
459
+ expertQ8: `enable f16;
460
+ struct P { M: u32, N: u32, rowsPerExpert: u32, inputPerSlot: u32 }
461
+ @group(0) @binding(0) var<storage, read> x: array<f32>;
462
+ @group(0) @binding(1) var<storage, read> qw: array<u32>;
463
+ @group(0) @binding(2) var<storage, read> qs: array<f16>;
464
+ @group(0) @binding(3) var<storage, read_write> y: array<f32>;
465
+ @group(0) @binding(4) var<uniform> p: P;
466
+ @group(0) @binding(5) var<storage, read> slots: array<u32>;
467
+ ${I8}
468
+ var<workgroup> part: array<f32, 1024>;
469
+ @compute @workgroup_size(256) fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) ng: vec3<u32>, @builtin(local_invocation_id) l: vec3<u32>) {
470
+ let z = wg.z; let m0 = (wg.y * ng.x + wg.x) * 4u; let t = l.x; let N = p.N; let M = p.M; let nb = N / 32u;
471
+ if (m0 >= M) { return; }
472
+ let rb = slots[z] * p.rowsPerExpert;
473
+ let xb = select(0u, z * N, p.inputPerSlot == 1u);
474
+ let r0 = rb + min(m0, M - 1u); let r1 = rb + min(m0 + 1u, M - 1u); let r2 = rb + min(m0 + 2u, M - 1u); let r3 = rb + min(m0 + 3u, M - 1u);
475
+ let q0 = (r0 * N) >> 2u; let q1 = (r1 * N) >> 2u; let q2 = (r2 * N) >> 2u; let q3 = (r3 * N) >> 2u;
476
+ var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
477
+ for (var k = t; k < N; k += 256u) {
478
+ let xv = x[xb + k]; let wi = k >> 2u; let ln = k & 3u; let b = k >> 5u;
479
+ a0 += i8at(qw[q0 + wi], ln) * f32(qs[r0 * nb + b]) * xv;
480
+ a1 += i8at(qw[q1 + wi], ln) * f32(qs[r1 * nb + b]) * xv;
481
+ a2 += i8at(qw[q2 + wi], ln) * f32(qs[r2 * nb + b]) * xv;
482
+ a3 += i8at(qw[q3 + wi], ln) * f32(qs[r3 * nb + b]) * xv;
483
+ }
484
+ part[t] = a0; part[256u + t] = a1; part[512u + t] = a2; part[768u + t] = a3;
485
+ workgroupBarrier();
486
+ for (var s = 128u; s > 0u; s >>= 1u) {
487
+ if (t < s) { part[t] += part[t + s]; part[256u + t] += part[256u + t + s]; part[512u + t] += part[512u + t + s]; part[768u + t] += part[768u + t + s]; }
488
+ workgroupBarrier();
489
+ }
490
+ if (t < 4u && m0 + t < M) { y[z * M + m0 + t] = part[t * 256u]; }
491
+ }`,
492
+ // gu = k × [gate F | up F] → act = k × F of SiLU(gate)·up.
493
+ siluMulMoe: `
494
+ struct P { F: u32, k: u32 }
495
+ @group(0) @binding(0) var<storage, read> gu: array<f32>;
496
+ @group(0) @binding(1) var<storage, read_write> act: array<f32>;
497
+ @group(0) @binding(2) var<uniform> p: P;
498
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
499
+ let i = g.x; if (i >= p.k * p.F) { return; }
500
+ let z = i / p.F; let j = i % p.F; let g0 = gu[z * 2u * p.F + j];
501
+ act[i] = (g0 / (1.0 + exp(-g0))) * gu[z * 2u * p.F + p.F + j];
502
+ }`,
503
+ // x[h] += Σ_j w_j · out[j·H + h] (the expert sum, then the residual add — llama.cpp's order)
504
+ moeAccum: `
505
+ struct P { H: u32, k: u32 }
506
+ @group(0) @binding(0) var<storage, read> out: array<f32>;
507
+ @group(0) @binding(1) var<storage, read> sel: array<u32>;
508
+ @group(0) @binding(2) var<storage, read_write> x: array<f32>;
509
+ @group(0) @binding(3) var<uniform> p: P;
510
+ @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) g: vec3<u32>) {
511
+ let h = g.x; if (h >= p.H) { return; }
512
+ var acc = 0.0;
513
+ for (var j = 0u; j < p.k; j++) { acc += bitcast<f32>(sel[p.k + j]) * out[j * p.H + h]; }
514
+ x[h] = x[h] + acc;
515
+ }`,
516
+ argmax: `
517
+ struct P { n: u32 }
518
+ @group(0) @binding(0) var<storage, read> lg: array<f32>;
519
+ @group(0) @binding(1) var<storage, read_write> res: array<u32>;
520
+ @group(0) @binding(2) var<uniform> p: P;
521
+ var<workgroup> bv: array<f32, 256>;
522
+ var<workgroup> bi: array<u32, 256>;
523
+ @compute @workgroup_size(256) fn main(@builtin(local_invocation_id) l: vec3<u32>) {
524
+ let t = l.x; var v = -3.4e38; var ix = 0u;
525
+ for (var i = t; i < p.n; i += 256u) { let x = lg[i]; if (x > v) { v = x; ix = i; } }
526
+ bv[t] = v; bi[t] = ix; workgroupBarrier();
527
+ for (var s = 128u; s > 0u; s >>= 1u) {
528
+ if (t < s) { let o = bv[t + s]; if (o > bv[t] || (o == bv[t] && bi[t + s] < bi[t])) { bv[t] = o; bi[t] = bi[t + s]; } }
529
+ workgroupBarrier();
530
+ }
531
+ if (t == 0u) { res[0] = bi[0]; }
532
+ }`,
533
+ };
534
+
535
+ export const STORAGE = 0x80, COPY_SRC = 0x04, COPY_DST = 0x08, UNIFORM = 0x40, MAP_READ = 0x01;
536
+
537
+ // ── Engine ──────────────────────────────────────────────────────────────────
538
+ export class Qwen3MoeSsd {
539
+ // A subclass for another architecture overrides these (gemma4_moe_ssd.js).
540
+ static configFrom(kv) { return configFromGguf(kv); }
541
+ static get ingestPlan() { return undefined; }
542
+ static tokenizerFrom(kv) { return new BpeTokenizer(kv); }
543
+
544
+ static async load(modelId = QWEN3_30B_A3B.repo, opts = {}) {
545
+ const { fetch: fetchFn = (u, i) => fetch(u, i), onProgress = () => {}, signal } = opts;
546
+ const src = opts.source || (modelId === QWEN3_30B_A3B.repo || !modelId ? QWEN3_30B_A3B : { repo: modelId, file: opts.file, revision: opts.revision || 'main' });
547
+ const url = opts.url || `https://huggingface.co/${src.repo}/resolve/${src.revision}/${src.file}`;
548
+ const key = opts.key || storeKey(opts.file || src.file || url.split('/').pop());
549
+ const root = opts.root || OPFS_ROOT; // the OPFS folder for this engine's stores
550
+ const dir = `${root}/${key}`;
551
+
552
+ onProgress({ status: 'init' });
553
+ if (!navigator.gpu) throw new Error('WebGPU is not available');
554
+ const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' });
555
+ if (!adapter) throw new Error('no WebGPU adapter');
556
+ if (!adapter.features.has('shader-f16')) throw new Error('this GPU lacks shader-f16');
557
+ const L = adapter.limits;
558
+ // subgroups: optional; an engine may use subgroup kernels when the adapter offers them.
559
+ const features = ['shader-f16', ...(adapter.features.has('subgroups') ? ['subgroups'] : [])];
560
+ const device = await adapter.requestDevice({
561
+ requiredFeatures: features,
562
+ requiredLimits: { maxBufferSize: L.maxBufferSize, maxStorageBufferBindingSize: L.maxStorageBufferBindingSize, maxStorageBuffersPerShaderStage: L.maxStorageBuffersPerShaderStage, maxComputeWorkgroupsPerDimension: L.maxComputeWorkgroupsPerDimension },
563
+ });
564
+ device.lost.then((info) => console.warn('qwen3-moe-ssd: GPU device lost:', info.message));
565
+
566
+ let manifest = null;
567
+ try { manifest = JSON.parse(await readOpfsText(`${dir}/manifest.json`) || 'null'); } catch (_) { manifest = null; }
568
+ if (opts.reingest || !manifest || !manifest.complete || manifest.format !== FORMAT) {
569
+ manifest = await ingestGguf({
570
+ url, key, root, fetch: fetchFn, signal, source: { repo: src.repo, file: src.file, revision: src.revision, sha256: src.sha256, size: src.size },
571
+ onProgress: (e) => onProgress(ingestProgress(e)),
572
+ plan: this.ingestPlan,
573
+ });
574
+ }
575
+ const headerFile = await (await (await navigator.storage.getDirectory()).getDirectoryHandle(root)).getDirectoryHandle(key);
576
+ const headerBytes = new Uint8Array(await (await (await headerFile.getFileHandle('header.bin')).getFile()).arrayBuffer());
577
+ const gguf = parseGguf(headerBytes);
578
+ const engine = new this(device, manifest, gguf, { ...opts, dir, key, url, adapterInfo: adapter.info || {} });
579
+ await engine.init(onProgress);
580
+ return engine;
581
+ }
582
+
583
+ constructor(device, manifest, gguf, opts) {
584
+ this.device = device;
585
+ this.manifest = manifest;
586
+ this.cfg = this.constructor.configFrom(gguf.kv);
587
+ this.tokenizer = this.constructor.tokenizerFrom(gguf.kv);
588
+ this.opts = opts;
589
+ this.maxCtx = Math.min(opts.maxCtx || 4096, this.cfg.contextLength);
590
+ this.gpuBytes = { dense: 0, pool: 0, kv: 0, act: 0 };
591
+ this.pipelines = {};
592
+ this.position = 0;
593
+ this.cached = []; // token ids whose K/V are in the cache, in order
594
+ // Prefetch: routers of the next `lookahead` layers guess their experts from this layer's
595
+ // post-attention state. Measured 2026-10-04 (full model, 4 GB pool, M4 Pro): off 4.41,
596
+ // lookahead 1 5.02, 2 5.48, 4 4.56, 6 3.90 tok/s — guesses beyond 2 layers waste SSD bandwidth.
597
+ this.prefetch = opts.prefetch ?? true;
598
+ this.lookahead = Math.max(1, Math.floor(opts.lookahead ?? 2));
599
+ this.resetCounters();
600
+ }
601
+
602
+ buffer(size, usage, cat = 'act') {
603
+ const b = this.device.createBuffer({ size: Math.ceil(size / 4) * 4, usage });
604
+ this.gpuBytes[cat] = (this.gpuBytes[cat] || 0) + b.size;
605
+ return b;
606
+ }
607
+ uniform(words) {
608
+ const b = this.buffer(Math.max(16, words.length * 4), UNIFORM | COPY_DST, 'act');
609
+ const ab = new ArrayBuffer(b.size), u = new Uint32Array(ab), f = new Float32Array(ab);
610
+ words.forEach((w, i) => { if (typeof w === 'object') f[i] = w.f; else u[i] = w; });
611
+ this.device.queue.writeBuffer(b, 0, ab);
612
+ return b;
613
+ }
614
+ // The kernel sources; a subclass adds its own.
615
+ get kernels() { return WGSL; }
616
+ pipeline(name) {
617
+ if (!this.pipelines[name]) {
618
+ const code = this.kernels[name];
619
+ if (!code) throw new Error(`no kernel ${name}`);
620
+ const module = this.device.createShaderModule({ code, label: name });
621
+ this.pipelines[name] = this.device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' }, label: name });
622
+ }
623
+ return this.pipelines[name];
624
+ }
625
+ bind(name, buffers) {
626
+ return this.device.createBindGroup({
627
+ layout: this.pipeline(name).getBindGroupLayout(0),
628
+ entries: buffers.map((b, i) => ({ binding: i, resource: { buffer: b } })),
629
+ });
630
+ }
631
+
632
+ async init(onProgress) {
633
+ const dev = this.device;
634
+ // Fail early on a shader that does not compile.
635
+ dev.pushErrorScope('validation');
636
+ for (const k of Object.keys(this.kernels)) this.pipeline(k);
637
+ const err = await dev.popErrorScope();
638
+ if (err) throw new Error(`WGSL: ${err.message}`);
639
+ await this.uploadDense(onProgress);
640
+ await this.initPool();
641
+ this.initBuffers();
642
+ await dev.queue.onSubmittedWorkDone();
643
+ }
644
+
645
+ // Dense weights: OPFS → GPU (this.w[name] = { q, s, rows, cols } or { raw }).
646
+ async uploadDense(onProgress) {
647
+ const dev = this.device, m = this.manifest;
648
+ const t0 = performance.now();
649
+ const reader = await OpfsReaderPool.open(`${this.opts.dir}/${m.dense.file}`, { workers: 4 });
650
+ const names = Object.keys(m.dense.tensors);
651
+ this.w = {};
652
+ let doneT = 0;
653
+ const CH = 64 << 20;
654
+ const upload = async (gbuf, part) => {
655
+ for (let o = 0; o < part.bytes; o += CH) {
656
+ const n = Math.min(CH, part.bytes - o);
657
+ const r = await reader.read(part.off + o, n);
658
+ if (r.got !== n) throw new Error(`dense read short at ${part.off + o}`);
659
+ dev.queue.writeBuffer(gbuf, o, r.buf, 0, n);
660
+ }
661
+ };
662
+ const queue = [];
663
+ for (const name of names) {
664
+ const d = m.dense.tensors[name];
665
+ if (d.type === 'q8' || d.type === 'q4') { // q4: Gemma 4 (two values per byte)
666
+ const q = this.buffer(d.q.bytes, STORAGE | COPY_DST, 'dense'), s = this.buffer(d.s.bytes, STORAGE | COPY_DST, 'dense');
667
+ this.w[name] = { q, s, rows: (d.type === 'q4' ? 2 * d.q.bytes : d.q.bytes) / d.dims[0], cols: d.dims[0] };
668
+ queue.push(() => upload(q, d.q), () => upload(s, d.s));
669
+ } else {
670
+ const raw = this.buffer(d.raw.bytes, STORAGE | COPY_DST, 'dense');
671
+ this.w[name] = { raw };
672
+ queue.push(() => upload(raw, d.raw));
673
+ }
674
+ }
675
+ let qi = 0;
676
+ await Promise.all([0, 1, 2, 3].map(async () => {
677
+ while (qi < queue.length) { await queue[qi++](); doneT++; onProgress({ status: 'weights', kind: 'tensors', loaded: doneT, total: queue.length }); }
678
+ }));
679
+ await reader.close();
680
+ await dev.queue.onSubmittedWorkDone();
681
+ this.denseUploadSecs = (performance.now() - t0) / 1000;
682
+ }
683
+
684
+ // Bytes for the expert cache. With opts.gpuBudgetBytes, the cache gets what the budget leaves after the weights
685
+ // that stay on the GPU, the KV cache and working buffers; a budget too small for those is refused. Otherwise
686
+ // opts.poolBytes, default 4 GiB.
687
+ poolBudget() {
688
+ const o = this.opts;
689
+ if (!o.gpuBudgetBytes) return o.poolBytes || 4 * 2 ** 30;
690
+ const gb = (n) => (n / 1e9).toFixed(2) + ' GB';
691
+ const fixed = this.gpuBytes.dense + this.kvBytes() + (o.workBytes ?? 512 * 2 ** 20);
692
+ const rest = o.gpuBudgetBytes - fixed;
693
+ const min = this.cfg.topK * 2 * this.manifest.experts.record;
694
+ if (rest < min) throw new Error(`GPU budget ${gb(o.gpuBudgetBytes)} is too small: weights that stay on the GPU, the KV cache and working buffers need ${gb(fixed)}, and the expert cache at least ${gb(min)} more`);
695
+ return rest;
696
+ }
697
+ // KV cache bytes at this.maxCtx (f16 K and V for every layer); engines with other layouts override it.
698
+ kvBytes() { const c = this.cfg; return 2 * 2 * c.layers * this.maxCtx * c.kvHeads * c.headDim; }
699
+
700
+ // Expert slot pool + the streamer that feeds it.
701
+ async initPool() {
702
+ const c = this.cfg, dev = this.device, xp = this.manifest.experts;
703
+ const maxSlotsByBinding = Math.floor(dev.limits.maxStorageBufferBindingSize / xp.parts.guQ.bytes);
704
+ this.poolSlots = Math.max(c.topK * 2, Math.min(maxSlotsByBinding, Math.floor(this.poolBudget() / xp.record), c.layers * c.experts));
705
+ this.pool = {};
706
+ for (const [k, part] of Object.entries(xp.parts)) this.pool[k] = this.buffer(this.poolSlots * part.bytes, STORAGE | COPY_DST, 'pool');
707
+ this.expertReader = await OpfsReaderPool.open(`${this.opts.dir}/${xp.file}`, { workers: this.opts.readers || 4 });
708
+ this.xs = new ExpertStreamer({
709
+ device: dev, reader: this.expertReader, recordBytes: xp.record,
710
+ recordOffset: (layer, e) => (layer * xp.perLayer + e) * xp.record,
711
+ parts: ['guQ', 'guS', 'dQ', 'dS'].map((k) => ({ buffer: this.pool[k], srcOffset: xp.parts[k].off, bytes: xp.parts[k].bytes })),
712
+ slots: this.poolSlots, numLayers: c.layers, numExperts: c.experts,
713
+ ...(this.opts.uploadRing !== undefined ? { uploadRing: this.opts.uploadRing } : {}),
714
+ ...(this.opts.evict ? { evict: this.opts.evict, hotHalfLife: this.opts.hotHalfLife } : {}),
715
+ });
716
+ this.gpuBytes.staging = this.xs.ringBytes();
717
+ }
718
+
719
+ // KV cache (f16), activations, uniforms and every layer's bind groups.
720
+ initBuffers() {
721
+ const c = this.cfg;
722
+ const kvn = c.kvHeads * c.headDim;
723
+ this.kc = []; this.vc = [];
724
+ for (let l = 0; l < c.layers; l++) {
725
+ this.kc.push(this.buffer(this.maxCtx * kvn * 2, STORAGE, 'kv'));
726
+ this.vc.push(this.buffer(this.maxCtx * kvn * 2, STORAGE, 'kv'));
727
+ }
728
+ const H = c.hidden, QN = c.heads * c.headDim, F = c.expertFf, K = c.topK;
729
+ if (2 * K * 4 > 256) throw new Error(`top-k ${K} > 32: selections would overlap their 256-byte readback regions`);
730
+ if (c.experts > 256) throw new Error(`${c.experts} experts: the router top-k kernel handles at most 256`);
731
+ const A = (n) => this.buffer(n * 4, STORAGE | COPY_SRC | COPY_DST);
732
+ this.a = {
733
+ x: A(H), xn: A(H), q: A(QN), k: A(kvn), v: A(kvn), att: A(QN), o: A(H), sc: A(c.heads * this.maxCtx),
734
+ rl: A(c.experts), sel: A(2 * K), pxn: A(H), prl: A(c.experts), psel: A(64 * this.lookahead), slots: A(K),
735
+ gu: A(K * 2 * F), act: A(K * F), dn: A(K * H), logits: A(c.vocab), am: A(4),
736
+ rope: A(c.headDim),
737
+ };
738
+ this.tok = this.buffer(16, UNIFORM | COPY_DST);
739
+ // Readback: this layer's selection at 0, guess d (layers ahead) at 256 × d bytes.
740
+ this.rbSel = this.buffer(256 * (1 + this.lookahead), MAP_READ | COPY_DST);
741
+ this.rbLogits = this.buffer(c.vocab * 4, MAP_READ | COPY_DST);
742
+ this.rbArg = this.buffer(16, MAP_READ | COPY_DST);
743
+
744
+ const u = {
745
+ nH: this.uniform([H]), rmsH: this.uniform([H, { f: c.eps }]), mmQ: this.uniform([QN, H]), mmKV: this.uniform([kvn, H]),
746
+ mmO: this.uniform([H, QN]), qk: this.uniform([c.headDim, { f: c.eps }]), kvn: this.uniform([kvn]),
747
+ att: this.uniform([c.heads, c.kvHeads, c.headDim, { f: 1 / Math.sqrt(c.headDim) }]), attOut: this.uniform([c.heads, c.kvHeads, c.headDim, 0]),
748
+ router: this.uniform([c.experts, H]), topk: this.uniform([c.experts, K]),
749
+ gu: this.uniform([2 * F, H, 2 * F, 0]), dn: this.uniform([H, F, H, 1]), silu: this.uniform([F, K]), acc: this.uniform([H, K]),
750
+ lm: this.uniform([c.vocab, H]), am: this.uniform([c.vocab]),
751
+ };
752
+ const a = this.a, W = this.w;
753
+ this.g = { embed: this.bind('embedQ8', [W['token_embd.weight'].q, W['token_embd.weight'].s, a.x, this.tok, u.nH]) };
754
+ this.layers = [];
755
+ for (let l = 0; l < c.layers; l++) {
756
+ const p = (n) => W[`blk.${l}.${n}.weight`];
757
+ this.layers.push({
758
+ rmsA: this.bind('rmsnorm', [a.x, p('attn_norm').raw, a.xn, u.rmsH]),
759
+ q: this.bind('matmulQ8', [a.xn, p('attn_q').q, p('attn_q').s, a.q, u.mmQ]),
760
+ k: this.bind('matmulQ8', [a.xn, p('attn_k').q, p('attn_k').s, a.k, u.mmKV]),
761
+ v: this.bind('matmulQ8', [a.xn, p('attn_v').q, p('attn_v').s, a.v, u.mmKV]),
762
+ ropeQ: this.bind('qkNormRope', [a.q, p('attn_q_norm').raw, a.rope, u.qk]),
763
+ ropeK: this.bind('qkNormRope', [a.k, p('attn_k_norm').raw, a.rope, u.qk]),
764
+ kv: this.bind('kvStore', [a.k, a.v, this.kc[l], this.vc[l], this.tok, u.kvn]),
765
+ score: this.bind('attnScore', [a.q, this.kc[l], a.sc, this.tok, u.att]),
766
+ soft: this.bind('softmax', [a.sc, this.tok]),
767
+ attOut: this.bind('attnOut', [a.sc, this.vc[l], a.att, this.tok, u.attOut]),
768
+ o: this.bind('matmulQ8', [a.att, p('attn_output').q, p('attn_output').s, a.o, u.mmO]),
769
+ addO: this.bind('addInPlace', [a.x, a.o, u.nH]),
770
+ rmsF: this.bind('rmsnorm', [a.x, p('ffn_norm').raw, a.xn, u.rmsH]),
771
+ router: this.bind('matmulF32', [a.xn, p('ffn_gate_inp').raw, a.rl, u.router]),
772
+ topk: this.bind('topk', [a.rl, a.sel, u.topk]),
773
+ // Routers of layers l+1..l+lookahead on this layer's post-attention state: the prefetch
774
+ // guesses, each written to its own 256-byte region of psel.
775
+ pf: Array.from({ length: Math.min(this.lookahead, c.layers - 1 - l) }, (_, i) => {
776
+ const at = (n) => W[`blk.${l + 1 + i}.${n}.weight`];
777
+ return {
778
+ rms: this.bind('rmsnorm', [a.x, at('ffn_norm').raw, a.pxn, u.rmsH]),
779
+ router: this.bind('matmulF32', [a.pxn, at('ffn_gate_inp').raw, a.prl, u.router]),
780
+ topk: this.device.createBindGroup({ layout: this.pipeline('topk').getBindGroupLayout(0), entries: [
781
+ { binding: 0, resource: { buffer: a.prl } },
782
+ { binding: 1, resource: { buffer: a.psel, offset: 256 * i, size: 8 * K } },
783
+ { binding: 2, resource: { buffer: u.topk } },
784
+ ] }),
785
+ };
786
+ }),
787
+ gu: this.bind('expertQ8', [a.xn, this.pool.guQ, this.pool.guS, a.gu, u.gu, a.slots]),
788
+ silu: this.bind('siluMulMoe', [a.gu, a.act, u.silu]),
789
+ dn: this.bind('expertQ8', [a.act, this.pool.dQ, this.pool.dS, a.dn, u.dn, a.slots]),
790
+ acc: this.bind('moeAccum', [a.dn, a.sel, a.x, u.acc]),
791
+ });
792
+ }
793
+ this.g.rmsOut = this.bind('rmsnorm', [a.x, W['output_norm.weight'].raw, a.xn, u.rmsH]);
794
+ this.g.lm = this.bind('matmulQ8', [a.xn, W['output.weight'].q, W['output.weight'].s, a.logits, u.lm]);
795
+ this.g.am = this.bind('argmax', [a.logits, a.am, u.am]);
796
+ }
797
+
798
+ // ── forward ───────────────────────────────────────────────────────────────
799
+ dispatch(pass, name, group, x, y = 1, z = 1) {
800
+ pass.setPipeline(this.pipeline(name));
801
+ pass.setBindGroup(0, group);
802
+ if (x > 65535) { y = Math.ceil(x / 32768); x = 32768; }
803
+ pass.dispatchWorkgroups(x, y, z);
804
+ }
805
+ encodeAttention(pass, l, seqLen) {
806
+ const c = this.cfg, g = this.layers[l], d = (n, gr, x, y, z) => this.dispatch(pass, n, gr, x, y, z);
807
+ const QN = c.heads * c.headDim, kvn = c.kvHeads * c.headDim;
808
+ d('rmsnorm', g.rmsA, 1);
809
+ d('matmulQ8', g.q, Math.ceil(QN / 4)); d('matmulQ8', g.k, Math.ceil(kvn / 4)); d('matmulQ8', g.v, Math.ceil(kvn / 4));
810
+ d('qkNormRope', g.ropeQ, c.heads); d('qkNormRope', g.ropeK, c.kvHeads);
811
+ d('kvStore', g.kv, Math.ceil(kvn / 256));
812
+ d('attnScore', g.score, Math.ceil(c.heads * seqLen / 256));
813
+ d('softmax', g.soft, c.heads);
814
+ d('attnOut', g.attOut, Math.ceil(QN / 256));
815
+ d('matmulQ8', g.o, Math.ceil(c.hidden / 4));
816
+ d('addInPlace', g.addO, Math.ceil(c.hidden / 256));
817
+ if (this.prefetch) for (const p of g.pf) { d('rmsnorm', p.rms, 1); d('matmulF32', p.router, c.experts); d('topk', p.topk, 1); }
818
+ d('rmsnorm', g.rmsF, 1);
819
+ d('matmulF32', g.router, c.experts);
820
+ d('topk', g.topk, 1);
821
+ }
822
+ // The token's embedding into a.x, and the head (final norm, logits, argmax when wanted): a subclass
823
+ // with another embedding or head format overrides these two.
824
+ encodeEmbed(pass) { this.dispatch(pass, 'embedQ8', this.g.embed, Math.ceil(this.cfg.hidden / 256)); }
825
+ encodeHead(pass, want) {
826
+ this.dispatch(pass, 'rmsnorm', this.g.rmsOut, 1);
827
+ this.dispatch(pass, 'matmulQ8', this.g.lm, Math.ceil(this.cfg.vocab / 4));
828
+ if (want === 'argmax') this.dispatch(pass, 'argmax', this.g.am, 1);
829
+ }
830
+ encodeExperts(pass, l) {
831
+ const c = this.cfg, g = this.layers[l], d = (n, gr, x, y, z) => this.dispatch(pass, n, gr, x, y, z);
832
+ d('expertQ8', g.gu, Math.ceil(2 * c.expertFf / 4), 1, c.topK);
833
+ d('siluMulMoe', g.silu, Math.ceil(c.topK * c.expertFf / 256));
834
+ d('expertQ8', g.dn, Math.ceil(c.hidden / 4), 1, c.topK);
835
+ d('moeAccum', g.acc, Math.ceil(c.hidden / 256));
836
+ }
837
+
838
+ writeTokenUniforms(token, pos) {
839
+ const c = this.cfg;
840
+ this.device.queue.writeBuffer(this.tok, 0, new Uint32Array([token, pos, pos + 1, 0]));
841
+ // NeoX RoPE angles the way llama.cpp's Metal kernel forms them: f32 pos × base^(-2i/hd).
842
+ const half = c.headDim / 2, r = new Float32Array(c.headDim), inv = Math.fround(-1 / c.headDim);
843
+ for (let i = 0; i < half; i++) {
844
+ const th = Math.fround(pos * Math.fround(Math.pow(c.ropeTheta, Math.fround(inv * 2 * i))));
845
+ r[i] = Math.cos(th); r[half + i] = Math.sin(th);
846
+ }
847
+ this.device.queue.writeBuffer(this.a.rope, 0, r);
848
+ }
849
+
850
+ resetCounters() {
851
+ this.counters = { tokens: 0, gpuWaitMs: 0, ensureMs: 0, encodeMs: 0, wallMs: 0, predHits: 0, predTotal: 0 };
852
+ if (this.xs) this.xs.resetStats();
853
+ }
854
+
855
+ // → { ids, preds }: this layer's top-k ids and `nPred` guesses for the layers after it.
856
+ async readSel(nPred) {
857
+ const K = this.cfg.topK, bytes = 256 * (1 + nPred);
858
+ const t0 = performance.now();
859
+ await this.rbSel.mapAsync(1, 0, bytes);
860
+ this.counters.gpuWaitMs += performance.now() - t0;
861
+ const u = new Uint32Array(this.rbSel.getMappedRange(0, bytes).slice(0));
862
+ this.rbSel.unmap();
863
+ const preds = [];
864
+ for (let i = 1; i <= nPred; i++) preds.push(u.subarray(64 * i, 64 * i + K));
865
+ return { ids: u.subarray(0, K), preds };
866
+ }
867
+ // Encodes the readback of layer l's selection and its guesses.
868
+ copySel(enc, l) {
869
+ enc.copyBufferToBuffer(this.a.sel, 0, this.rbSel, 0, 8 * this.cfg.topK);
870
+ const n = this.prefetch ? this.layers[l].pf.length : 0;
871
+ if (n) enc.copyBufferToBuffer(this.a.psel, 0, this.rbSel, 256, 256 * n);
872
+ return n;
873
+ }
874
+
875
+ // One token through all layers at position this.position. want: 'none' | 'argmax' | 'logits'.
876
+ async step(token, want = 'argmax') {
877
+ const c = this.cfg, dev = this.device, K = c.topK;
878
+ if (this.position >= this.maxCtx) throw new Error(`context full (${this.maxCtx} tokens)`);
879
+ const tStart = performance.now();
880
+ const pos = this.position, seqLen = pos + 1;
881
+ this.writeTokenUniforms(token, pos);
882
+ let te = performance.now();
883
+ let enc = dev.createCommandEncoder(), pass = enc.beginComputePass();
884
+ this.encodeEmbed(pass);
885
+ this.encodeAttention(pass, 0, seqLen);
886
+ pass.end();
887
+ let nPred = this.copySel(enc, 0);
888
+ dev.queue.submit([enc.finish()]);
889
+ this.counters.encodeMs += performance.now() - te;
890
+ let predIds = null;
891
+ for (let l = 0; l < c.layers; l++) {
892
+ const { ids, preds } = await this.readSel(nPred);
893
+ if (predIds) { // accuracy of the one-layer-ahead guess made one layer earlier
894
+ let hit = 0; for (const e of ids) if (predIds.includes(e)) hit++;
895
+ this.counters.predHits += hit; this.counters.predTotal += K;
896
+ }
897
+ const tq = performance.now();
898
+ const ensuring = this.xs.ensure(l, ids);
899
+ preds.forEach((p, i) => this.xs.prefetch(l + 1 + i, p)); // nearest layer first
900
+ predIds = preds.length ? Array.from(preds[0]) : null;
901
+ const slots = await ensuring;
902
+ this.counters.ensureMs += performance.now() - tq;
903
+ dev.queue.writeBuffer(this.a.slots, 0, slots);
904
+ te = performance.now();
905
+ enc = dev.createCommandEncoder(); pass = enc.beginComputePass();
906
+ this.encodeExperts(pass, l);
907
+ if (l + 1 < c.layers) {
908
+ this.encodeAttention(pass, l + 1, seqLen);
909
+ pass.end();
910
+ nPred = this.copySel(enc, l + 1);
911
+ } else {
912
+ if (want !== 'none') this.encodeHead(pass, want);
913
+ pass.end();
914
+ if (want === 'argmax') enc.copyBufferToBuffer(this.a.am, 0, this.rbArg, 0, 4);
915
+ if (want === 'logits') enc.copyBufferToBuffer(this.a.logits, 0, this.rbLogits, 0, c.vocab * 4);
916
+ }
917
+ dev.queue.submit([enc.finish()]);
918
+ this.counters.encodeMs += performance.now() - te;
919
+ this.xs.release(slots);
920
+ }
921
+ this.position++;
922
+ this.cached.push(token);
923
+ let result = null;
924
+ const tw = performance.now();
925
+ if (want === 'argmax') {
926
+ await this.rbArg.mapAsync(1, 0, 4);
927
+ result = new Uint32Array(this.rbArg.getMappedRange(0, 4))[0];
928
+ this.rbArg.unmap();
929
+ } else if (want === 'logits') {
930
+ await this.rbLogits.mapAsync(1);
931
+ result = new Float32Array(this.rbLogits.getMappedRange().slice(0));
932
+ this.rbLogits.unmap();
933
+ } else {
934
+ await dev.queue.onSubmittedWorkDone();
935
+ }
936
+ this.counters.gpuWaitMs += performance.now() - tw;
937
+ this.counters.tokens++;
938
+ this.counters.wallMs += performance.now() - tStart;
939
+ return result;
940
+ }
941
+
942
+ reset() { this.position = 0; this.cached = []; }
943
+
944
+ // Feeds `ids` after the cached prefix they share (re-prefilling only what changed); returns
945
+ // the logits/argmax of the last token.
946
+ async prefill(ids, want = 'argmax') {
947
+ let common = 0;
948
+ while (common < ids.length - 1 && common < this.cached.length && this.cached[common] === ids[common]) common++;
949
+ if (common < this.cached.length) { this.position = common; this.cached.length = common; }
950
+ let r = null;
951
+ for (let i = common; i < ids.length; i++) r = await this.step(ids[i], i === ids.length - 1 ? want : 'none');
952
+ return r;
953
+ }
954
+
955
+ // Diagnostics for the llama.cpp comparison: greedy continuation of raw token ids, with the
956
+ // top-n logits at every generated position.
957
+ async greedy(ids, n, { top = 5 } = {}) {
958
+ this.reset();
959
+ const out = [], tops = [];
960
+ let logits = await this.prefill(ids, 'logits');
961
+ for (let i = 0; i < n; i++) {
962
+ const order = topN(logits, top);
963
+ let mx = -Infinity, sum = 0;
964
+ for (let j = 0; j < logits.length; j++) if (logits[j] > mx) mx = logits[j];
965
+ for (let j = 0; j < logits.length; j++) sum += Math.exp(logits[j] - mx);
966
+ const lse = mx + Math.log(sum);
967
+ for (const o of order) o.logprob = o.logit - lse;
968
+ tops.push(order);
969
+ const next = order[0].id;
970
+ out.push(next);
971
+ if (i + 1 < n) logits = await this.step(next, 'logits');
972
+ }
973
+ return { ids: out, tops };
974
+ }
975
+
976
+ async *generate(messages, { maxNewTokens = 512, signal, enableThinking = true } = {}) {
977
+ const ids = this.tokenizer.encode(this.chatPrompt(messages, { enableThinking }));
978
+ const stops = new Set(this.stopTokenIds);
979
+ let next = await this.prefill(ids, 'argmax');
980
+ const outIds = [];
981
+ for (let i = 0; i < maxNewTokens; i++) {
982
+ if (signal && signal.aborted) break;
983
+ if (stops.has(next)) break;
984
+ outIds.push(next);
985
+ yield { text: this.tokenizer.decode(outIds), token: next, tokens: outIds.length };
986
+ if (this.position >= this.maxCtx) break;
987
+ next = await this.step(next, 'argmax');
988
+ }
989
+ }
990
+
991
+ chatPrompt(messages, opts) { return chatPrompt(messages, opts); }
992
+ get stopTokenIds() { return [this.cfg.eos, this.tokenizer.ids.get('<|im_end|>'), this.tokenizer.ids.get('<|endoftext|>')].filter((t) => t !== undefined); }
993
+
994
+ // Raw ids of the reasoning delimiters, for hosts that re-mark the thought block.
995
+ get thinkOpenTokenId() { return this.tokenizer.ids.get('<think>') ?? null; }
996
+ get thinkCloseTokenId() { return this.tokenizer.ids.get('</think>') ?? null; }
997
+
998
+ stats() {
999
+ const s = this.xs.stats, c = this.counters, n = Math.max(1, c.tokens);
1000
+ return {
1001
+ tokens: c.tokens, tokPerSec: c.tokens / (c.wallMs / 1000),
1002
+ msPerToken: c.wallMs / n, gpuWaitMsPerToken: c.gpuWaitMs / n, ensureMsPerToken: c.ensureMs / n, encodeMsPerToken: c.encodeMs / n,
1003
+ readMsPerToken: s.readMs / n, uploadMsPerToken: s.uploadMs / n,
1004
+ hitRate: (s.hits + s.lateHits) / Math.max(1, s.hits + s.lateHits + s.misses),
1005
+ hits: s.hits, lateHits: s.lateHits, misses: s.misses, evictions: s.evictions,
1006
+ bytesReadPerToken: s.bytesRead / n, prefetchIssued: s.prefetchIssued, prefetchUsed: s.prefetchUsed,
1007
+ prefetchAccuracy: c.predTotal ? c.predHits / c.predTotal : null,
1008
+ hitRateByLayer: Array.from(s.hitsByLayer, (h, l) => +(h / Math.max(1, h + s.missesByLayer[l])).toFixed(3)),
1009
+ poolSlots: this.poolSlots, resident: this.xs.resident(), gpuBytes: { ...this.gpuBytes },
1010
+ };
1011
+ }
1012
+
1013
+ async warmup() {}
1014
+
1015
+ async dispose() {
1016
+ try { await this.xs.drain(); } catch (_) {}
1017
+ try { await this.expertReader.close(); } catch (_) {}
1018
+ try { this.device.destroy(); } catch (_) {}
1019
+ }
1020
+ }
1021
+
1022
+ export function topN(logits, n) {
1023
+ const best = [];
1024
+ for (let i = 0; i < logits.length; i++) {
1025
+ const v = logits[i];
1026
+ if (best.length < n || v > best[best.length - 1].logit) {
1027
+ best.push({ id: i, logit: v });
1028
+ best.sort((a, b) => b.logit - a.logit || a.id - b.id);
1029
+ if (best.length > n) best.pop();
1030
+ }
1031
+ }
1032
+ return best;
1033
+ }
1034
+
1035
+ export async function removeIngest(key) { await removeStore(key, { root: OPFS_ROOT }); }
examples/chat/index.html ADDED
@@ -0,0 +1,236 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!doctype html>
2
+ <html lang="en">
3
+ <head>
4
+ <meta charset="utf-8">
5
+ <meta name="viewport" content="width=device-width, initial-scale=1">
6
+ <title>diskformer chat</title>
7
+ <style>
8
+ :root { --bg: #fafaf9; --fg: #1c1917; --muted: #78716c; --line: #e7e5e4; --card: #fff; --accent: #2563eb; --bad: #b91c1c; }
9
+ @media (prefers-color-scheme: dark) { :root { --bg: #161514; --fg: #e7e5e4; --muted: #a8a29e; --line: #2e2b29; --card: #1f1d1b; --accent: #60a5fa; --bad: #f87171; } }
10
+ * { box-sizing: border-box; }
11
+ body { margin: 0; background: var(--bg); color: var(--fg); font: 15px/1.5 system-ui, -apple-system, Segoe UI, sans-serif; }
12
+ main { max-width: 760px; margin: 0 auto; padding: 24px 16px 48px; }
13
+ h1 { font-size: 22px; margin: 0 0 6px; }
14
+ p { margin: 0 0 12px; }
15
+ .muted { color: var(--muted); }
16
+ .card { background: var(--card); border: 1px solid var(--line); border-radius: 10px; padding: 14px; margin: 14px 0; }
17
+ .row { display: flex; flex-wrap: wrap; gap: 10px 16px; align-items: end; }
18
+ label { display: flex; flex-direction: column; gap: 4px; font-size: 13px; color: var(--muted); }
19
+ select, input, textarea, button { font: inherit; color: var(--fg); background: var(--bg); border: 1px solid var(--line); border-radius: 7px; padding: 6px 9px; }
20
+ input[type=number] { width: 90px; }
21
+ button { cursor: pointer; }
22
+ button.primary { background: var(--accent); border-color: var(--accent); color: #fff; }
23
+ button:disabled { opacity: .5; cursor: default; }
24
+ .bar { height: 6px; background: var(--line); border-radius: 3px; overflow: hidden; margin: 10px 0 6px; }
25
+ .bar span { display: block; height: 100%; width: 0; background: var(--accent); transition: width .3s; }
26
+ #status.bad { color: var(--bad); }
27
+ .meter { display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr)); gap: 10px; }
28
+ .meter div { border: 1px solid var(--line); border-radius: 8px; padding: 8px 10px; }
29
+ .meter b { display: block; font-size: 20px; }
30
+ .meter span { font-size: 12px; color: var(--muted); }
31
+ #log { display: flex; flex-direction: column; gap: 10px; margin-bottom: 12px; }
32
+ .msg { white-space: pre-wrap; overflow-wrap: anywhere; padding: 9px 12px; border-radius: 9px; border: 1px solid var(--line); }
33
+ .msg.user { align-self: flex-end; background: var(--bg); max-width: 85%; }
34
+ .msg.model { background: var(--card); }
35
+ form { display: flex; gap: 8px; align-items: end; }
36
+ textarea { flex: 1; resize: vertical; min-height: 44px; }
37
+ footer { font-size: 12px; color: var(--muted); margin-top: 24px; }
38
+ </style>
39
+ </head>
40
+ <body>
41
+ <main>
42
+ <h1>diskformer chat</h1>
43
+ <p>A mixture-of-experts model larger than the GPU memory you give it, answering in this tab. The model downloads once
44
+ into this browser's private storage (OPFS). The GPU holds the dense layers and a small cache of experts; the other
45
+ experts are read from disk as each token needs them.</p>
46
+ <p class="muted">Chromium with WebGPU. Nothing leaves this tab except the model download from Hugging Face.</p>
47
+ <p id="framed" hidden>This page is inside a frame, where browser storage can be partitioned or limited.
48
+ <a id="ownTab" target="_blank" rel="noopener">Open it in its own tab</a> to keep the model in first-party storage.</p>
49
+
50
+ <section class="card">
51
+ <div class="row">
52
+ <label>Model <select id="model"></select></label>
53
+ <label>GPU memory budget (GB) <input id="budget" type="number" min="1" step="0.5"></label>
54
+ <label>Context (tokens) <select id="ctx"><option>1024</option><option selected>2048</option><option>4096</option></select></label>
55
+ <button id="load" class="primary">Download and load</button>
56
+ <button id="remove">Remove from disk</button>
57
+ </div>
58
+ <div class="bar"><span id="fill"></span></div>
59
+ <p id="status" class="muted">Pick a model and a budget. The first load downloads it; later loads take seconds.</p>
60
+ </section>
61
+
62
+ <section class="card" id="meterCard" hidden>
63
+ <div class="meter">
64
+ <div><b id="mModel">–</b><span>model in this browser's storage</span></div>
65
+ <div><b id="mGpu">–</b><span id="mGpuNote">GPU memory used</span></div>
66
+ <div><b id="mCache">–</b><span>experts cached on the GPU</span></div>
67
+ <div><b id="mSpeed">–</b><span id="mSpeedNote">decode speed</span></div>
68
+ </div>
69
+ </section>
70
+
71
+ <section id="chat" hidden>
72
+ <div id="log"></div>
73
+ <form id="ask">
74
+ <textarea id="q" rows="2" placeholder="Ask something"></textarea>
75
+ <button class="primary" id="send">Send</button>
76
+ <button type="button" id="stop" disabled>Stop</button>
77
+ </form>
78
+ </section>
79
+
80
+ <footer>Greedy decoding, so the same question gets the same answer. Engines and store:
81
+ <a href="https://github.com/NakliTechie/diskformer.js">diskformer.js</a>. Model terms: Gemma (ai.google.dev/gemma/terms), Qwen3.6 (Apache-2.0).</footer>
82
+ </main>
83
+
84
+ <script type="module">
85
+ import { Gemma4MoeSsd, GEMMA4_26B_A4B } from '../../engines/gemma4_moe_ssd.js';
86
+ import { Qwen35MoeSsd, QWEN36_35B_A3B } from '../../engines/qwen35_moe_ssd.js';
87
+ import { removeStore, storeKey } from '../../src/ingest.js';
88
+
89
+ // ?model=gemma|qwen ?budget=GB ?ctx=tokens ?url=GGUF (a local copy) ?root=OPFS folder ?gate=refs.json
90
+ const MODELS = {
91
+ gemma: { label: 'Gemma 4 26B-A4B · 14.4 GB', Engine: Gemma4MoeSsd, source: GEMMA4_26B_A4B, budget: 3 },
92
+ qwen: { label: 'Qwen3.6 35B-A3B · 36.9 GB', Engine: Qwen35MoeSsd, source: QWEN36_35B_A3B, budget: 4 },
93
+ };
94
+ const q = new URLSearchParams(location.search);
95
+ const ROOT = q.get('root') || 'diskformer';
96
+ const $ = (id) => document.getElementById(id);
97
+ const GB = (n) => (n / 1e9).toFixed(2) + ' GB';
98
+ let engine = null, messages = [], ctrl = null, busy = ''; // busy: '' | 'load' | 'remove' | 'reply'
99
+
100
+ for (const [k, m] of Object.entries(MODELS)) $('model').add(new Option(m.label, k));
101
+ if (MODELS[q.get('model')]) $('model').value = q.get('model');
102
+ const pick = () => MODELS[$('model').value];
103
+ $('budget').value = q.get('budget') || pick().budget;
104
+ if (q.get('ctx')) $('ctx').value = q.get('ctx');
105
+ $('model').onchange = () => { if (!q.get('budget')) $('budget').value = pick().budget; };
106
+
107
+ const status = (text, bad = false) => { $('status').textContent = text; $('status').className = bad ? 'bad' : 'muted'; window.progressText = text; };
108
+ const fill = (f) => { $('fill').style.width = `${Math.round(100 * f)}%`; };
109
+ addEventListener('error', (e) => status(`Error: ${e.message}`, true));
110
+ addEventListener('unhandledrejection', (e) => status(`Error: ${e.reason && e.reason.message || e.reason}`, true));
111
+ const controls = () => {
112
+ for (const id of ['load', 'remove', 'model']) $(id).disabled = !!busy;
113
+ $('send').disabled = !!busy || !engine; $('stop').disabled = busy !== 'reply';
114
+ };
115
+ // The OPFS folder a model's store lives in: the loaded engine's, else the name the engine would give it.
116
+ const keyOf = (m) => (engine && engine.constructor === m.Engine ? engine.opts.key : storeKey(m === MODELS.qwen && q.get('url') ? q.get('url').split('/').pop() : m.source.file));
117
+ controls();
118
+ if (window.self !== window.top) { $('framed').hidden = false; $('ownTab').href = location.href; }
119
+ if (!navigator.gpu) status('This page needs WebGPU: Chrome or Edge 121 or later, on a GPU with shader-f16.', true);
120
+
121
+ async function load() {
122
+ if (busy) return null;
123
+ busy = 'load'; controls();
124
+ const m = pick(), budgetGB = Number($('budget').value), maxCtx = Number($('ctx').value);
125
+ $('chat').hidden = true;
126
+ const t0 = performance.now();
127
+ try {
128
+ if (engine) { const old = engine; engine = null; await old.dispose(); }
129
+ engine = await m.Engine.load(null, {
130
+ url: q.get('url') || undefined, root: ROOT, gpuBudgetBytes: budgetGB * 1e9, maxCtx,
131
+ onProgress: (e) => {
132
+ if (e.status === 'ingest-plan' && e.quota && e.quota - (e.usage || 0) < e.needed)
133
+ status(`Needs ${GB(e.needed)} of storage; the browser grants ${GB(e.quota - (e.usage || 0))} so far and raises that as the download writes, while the disk has room.`);
134
+ else if (e.status === 'ingest-resume') { fill(e.loaded / e.total); status(`Resuming the download at ${GB(e.loaded)} of ${GB(e.total)}`); }
135
+ else if (e.status === 'weights' && e.kind === 'tensors') { fill(e.loaded / e.total); status(`Loading the dense layers onto the GPU: ${e.loaded} of ${e.total} tensors`); }
136
+ else if (e.status === 'weights') {
137
+ fill(e.loaded / e.total);
138
+ const rate = e.secs ? (e.loaded - (e.resumedFrom || 0)) / e.secs : 0, left = rate ? (e.total - e.loaded) / rate : 0;
139
+ status(`Downloading into this browser's storage (once): ${GB(e.loaded)} of ${GB(e.total)}${rate ? ` · ${(rate / 1e6).toFixed(0)} MB/s · ${Math.ceil(left / 60)} min left` : ''}`);
140
+ } else if (e.status === 'init') status('Starting WebGPU');
141
+ },
142
+ });
143
+ engine.device.lost.then((info) => status(`The GPU device was lost: ${info.message}. Reload the page.`, true));
144
+ fill(1);
145
+ const used = Object.values(engine.gpuBytes).reduce((a, b) => a + b, 0), c = engine.cfg;
146
+ const disk = engine.manifest.experts.bytes + engine.manifest.dense.bytes;
147
+ $('mModel').textContent = GB(disk);
148
+ $('mGpu').textContent = GB(used);
149
+ $('mGpuNote').textContent = `GPU memory used, budget ${budgetGB} GB`;
150
+ $('mCache').textContent = `${engine.poolSlots.toLocaleString()} of ${(c.layers * c.experts).toLocaleString()}`;
151
+ $('meterCard').hidden = false; $('chat').hidden = false;
152
+ status(`Loaded in ${((performance.now() - t0) / 1000).toFixed(1)} s. ${GB(disk)} in OPFS; ${GB(used)} on the GPU.`);
153
+ messages = []; $('log').replaceChildren();
154
+ return { model: $('model').value, budgetGB, maxCtx, modelBytes: m.source.size, diskBytes: disk, gpuBytes: used, gpuBytesByKind: { ...engine.gpuBytes }, poolSlots: engine.poolSlots, experts: c.layers * c.experts };
155
+ } catch (e) { status(`Could not load: ${e.message}`, true); throw e; }
156
+ finally { busy = ''; controls(); }
157
+ }
158
+
159
+ // One reply, streamed into onText. Returns the text, its timing, and whether `signal` stopped it.
160
+ async function reply(history, { maxNewTokens = 1024, onText = () => {}, signal } = {}) {
161
+ const t0 = performance.now();
162
+ let text = '', tokens = 0, tFirst = 0;
163
+ for await (const r of engine.generate(history, { maxNewTokens, signal, enableThinking: false })) {
164
+ if (!tokens) tFirst = performance.now();
165
+ text = r.text; tokens = r.tokens; onText(text);
166
+ }
167
+ const decodeSecs = (performance.now() - tFirst) / 1000;
168
+ return { text, tokens, stopped: !!(signal && signal.aborted), firstTokenMs: tokens ? Math.round(tFirst - t0) : null, decodeTokPerSec: tokens > 1 ? +((tokens - 1) / decodeSecs).toFixed(2) : null };
169
+ }
170
+
171
+ function bubble(role, text) { const d = document.createElement('div'); d.className = `msg ${role}`; d.textContent = text; $('log').append(d); return d; }
172
+
173
+ $('load').onclick = () => load().catch(() => {});
174
+ $('remove').onclick = async () => {
175
+ if (busy) return;
176
+ busy = 'remove'; controls();
177
+ const m = pick(), key = keyOf(m), name = m.label.split(' · ')[0];
178
+ try {
179
+ if (engine) { const old = engine; engine = null; await old.dispose(); }
180
+ $('chat').hidden = true; $('meterCard').hidden = true; fill(0);
181
+ await removeStore(key, { root: ROOT });
182
+ status(`Removed ${name} from this browser's storage.`);
183
+ } catch (e) {
184
+ if (e.name === 'NotFoundError') status(`${name} is not in this browser's storage.`);
185
+ else status(`Could not remove ${name}: ${e.message}`, true);
186
+ } finally { busy = ''; controls(); }
187
+ };
188
+ $('stop').onclick = () => { if (ctrl) ctrl.abort(); };
189
+ $('q').onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); $('ask').requestSubmit(); } };
190
+ $('ask').onsubmit = async (e) => {
191
+ e.preventDefault();
192
+ const text = $('q').value.trim();
193
+ if (!text || !engine || busy) return;
194
+ busy = 'reply'; ctrl = new AbortController(); controls();
195
+ $('q').value = '';
196
+ messages.push({ role: 'user', content: text });
197
+ bubble('user', text);
198
+ const out = bubble('model', '…');
199
+ try {
200
+ const r = await reply(messages, { signal: ctrl.signal, onText: (t) => { out.textContent = t; } });
201
+ if (r.stopped) out.textContent = r.text ? `${r.text} [stopped]` : '[stopped before the reply]';
202
+ if (r.text) messages.push({ role: 'assistant', content: r.text });
203
+ else messages.pop(); // no reply: drop the turn so the next one is not two user turns
204
+ if (r.decodeTokPerSec) { $('mSpeed').textContent = `${r.decodeTokPerSec} tok/s`; $('mSpeedNote').textContent = `decode speed · first token ${(r.firstTokenMs / 1000).toFixed(1)} s`; }
205
+ } catch (err) { out.textContent = `Error: ${err.message}`; messages.pop(); }
206
+ finally { busy = ''; ctrl = null; controls(); }
207
+ };
208
+
209
+ // Gate: load, then replay each reference conversation and compare the text with llama.cpp's.
210
+ async function gate(url) {
211
+ const ref = await (await fetch(url, { cache: 'no-store' })).json();
212
+ const info = await load();
213
+ const PREFIX = '<|channel>thought\n<channel|>';
214
+ const rows = [];
215
+ for (const c of ref.conversations) {
216
+ status(`Gate: ${c.name} (${rows.length + 1} of ${ref.conversations.length})`);
217
+ const r = await reply(c.messages, { maxNewTokens: c.tokens || 64 });
218
+ if (c.content === undefined) { rows.push({ name: c.name, text: r.text, tokens: r.tokens, decodeTokPerSec: r.decodeTokPerSec, firstTokenMs: r.firstTokenMs }); continue; }
219
+ const theirs = c.content.startsWith(PREFIX) ? c.content.slice(PREFIX.length) : c.content;
220
+ let at = 0; while (at < r.text.length && at < theirs.length && r.text[at] === theirs[at]) at++;
221
+ rows.push({ name: c.name, identical: r.text === theirs, sameUpTo: at, chars: r.text.length, tokens: r.tokens, decodeTokPerSec: r.decodeTokPerSec, firstTokenMs: r.firstTokenMs, ...(r.text === theirs ? {} : { ours: r.text, theirs }) });
222
+ }
223
+ // A conversation without reference text (no llama.cpp run to compare with) only has to produce a reply.
224
+ const compared = rows.filter((r) => 'identical' in r), identical = compared.filter((r) => r.identical).length;
225
+ const speeds = rows.map((r) => r.decodeTokPerSec).filter(Boolean).sort((a, b) => a - b);
226
+ const ok = identical === compared.length && rows.every((r) => r.tokens > 0);
227
+ window.result = { ok, identical, compared: compared.length, of: rows.length, ...info, medianDecodeTokPerSec: speeds[speeds.length >> 1], rows };
228
+ status(`Gate: ${identical} of ${compared.length} replies identical to the reference; ${rows.length - compared.length} without one.`);
229
+ }
230
+
231
+ window.demo = { load, reply, get engine() { return engine; } };
232
+ if (q.get('gate')) gate(q.get('gate')).catch((e) => { window.result = { ok: false, error: e.message, stack: e.stack }; });
233
+ else if (q.get('auto')) load().catch(() => {});
234
+ </script>
235
+ </body>
236
+ </html>
examples/chat/refs/gemma.json ADDED
@@ -0,0 +1,137 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "about": "llama.cpp b9830 (c1a1c8ee9), Metal, -ngl 99: llama-server /v1/chat/completions, temperature 0, max_tokens 64, on google/gemma-4-26B-A4B-it-qat-q4_0-gguf @ d1c082be (gemma-4-26B_q4_0-it.gguf, sha256 3eca3b8f...eca51d). Recorded 2026-10-05 for LocalMind's parity revalidation.",
3
+ "conversations": [
4
+ {
5
+ "name": "explain-tcp",
6
+ "messages": [
7
+ {
8
+ "role": "user",
9
+ "content": "Explain the difference between TCP and UDP to a new programmer."
10
+ }
11
+ ],
12
+ "prompt": "<|turn>user\nExplain the difference between TCP and UDP to a new programmer.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
13
+ "content": "<|channel>thought\n<channel|>To understand the difference between **TCP** and **UDP**, forget about code for a second and imagine you are sending messages to a friend.\n\n### 1. TCP (Transmission Control Protocol)\n**The \"Certified Mail\" Approach**\n\nImagine you are sending a very important legal document via certified mail. You want",
14
+ "finish": "length",
15
+ "tokens": 64
16
+ },
17
+ {
18
+ "name": "json-extract",
19
+ "messages": [
20
+ {
21
+ "role": "user",
22
+ "content": "Extract the name, city and age from this sentence as JSON: \"Priya, 34, moved from Pune to Lisbon last year.\""
23
+ }
24
+ ],
25
+ "prompt": "<|turn>user\nExtract the name, city and age from this sentence as JSON: \"Priya, 34, moved from Pune to Lisbon last year.\"<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
26
+ "content": "<|channel>thought\n<channel|>```json\n{\n \"name\": \"Priya\",\n \"city\": \"Lisbon\",\n \"age\": 34\n}\n```",
27
+ "finish": "stop",
28
+ "tokens": 35
29
+ },
30
+ {
31
+ "name": "translate-hi",
32
+ "messages": [
33
+ {
34
+ "role": "user",
35
+ "content": "Translate into Hindi: \"The library opens at nine and closes at six.\""
36
+ }
37
+ ],
38
+ "prompt": "<|turn>user\nTranslate into Hindi: \"The library opens at nine and closes at six.\"<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
39
+ "content": "<|channel>thought\n<channel|>The translation of \"The library opens at nine and closes at six\" in Hindi is:\n\n**पुस्तकालय नौ बजे खुलता है और छह बजे बंद होता है।**\n\n**Transliteration:** \n*Pustakalay nau baje khulta hai aur chhah baje band hota hai.*",
40
+ "finish": "length",
41
+ "tokens": 64
42
+ },
43
+ {
44
+ "name": "word-problem",
45
+ "messages": [
46
+ {
47
+ "role": "user",
48
+ "content": "A train leaves at 14:10 and arrives at 17:45. How long is the journey? Show the steps briefly."
49
+ }
50
+ ],
51
+ "prompt": "<|turn>user\nA train leaves at 14:10 and arrives at 17:45. How long is the journey? Show the steps briefly.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
52
+ "content": "<|channel>thought\n<channel|>To find the duration of the journey, follow these steps:\n\n1. **Subtract the hours:**\n 17:00 - 14:00 = **3 hours**\n\n2. **Subtract the minutes:**\n 45 minutes - 10 minutes = **35 minutes",
53
+ "finish": "length",
54
+ "tokens": 64
55
+ },
56
+ {
57
+ "name": "sql",
58
+ "messages": [
59
+ {
60
+ "role": "user",
61
+ "content": "Write a SQL query that returns the five customers with the highest total order value from tables customers(id, name) and orders(id, customer_id, amount)."
62
+ }
63
+ ],
64
+ "prompt": "<|turn>user\nWrite a SQL query that returns the five customers with the highest total order value from tables customers(id, name) and orders(id, customer_id, amount).<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
65
+ "content": "<|channel>thought\n<channel|>To get the five customers with the highest total order value, you need to `JOIN` the two tables, group the results by the customer, and then sum their order amounts.\n\nHere is the SQL query:\n\n```sql\nSELECT \n c.name, \n SUM(o.amount)",
66
+ "finish": "length",
67
+ "tokens": 64
68
+ },
69
+ {
70
+ "name": "list-rivers",
71
+ "messages": [
72
+ {
73
+ "role": "user",
74
+ "content": "List five major rivers of Africa, one per line, with the country where each reaches the sea."
75
+ }
76
+ ],
77
+ "prompt": "<|turn>user\nList five major rivers of Africa, one per line, with the country where each reaches the sea.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
78
+ "content": "<|channel>thought\n<channel|>Nile, Egypt\nCongo, Democratic Republic of the Congo\nNiger, Nigeria\nZambezi, Mozambique\nOrange, South Africa",
79
+ "finish": "stop",
80
+ "tokens": 29
81
+ },
82
+ {
83
+ "name": "unicode",
84
+ "messages": [
85
+ {
86
+ "role": "user",
87
+ "content": "What do these emoji usually mean in a text message: 🙏 🔥 💀 ?"
88
+ }
89
+ ],
90
+ "prompt": "<|turn>user\nWhat do these emoji usually mean in a text message: 🙏 🔥 💀 ?<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
91
+ "content": "<|channel>thought\n<channel|>The meaning of these emojis can change depending on the context and the age of the person sending them, but here is the general breakdown of how they are used in text messages:\n\n### 1. 🙏 (Folded Hands)\nThis emoji is the most versatile and its meaning depends entirely on the tone of the conversation",
92
+ "finish": "length",
93
+ "tokens": 64
94
+ },
95
+ {
96
+ "name": "haiku-autumn",
97
+ "messages": [
98
+ {
99
+ "role": "user",
100
+ "content": "Write a haiku about autumn leaves in a city."
101
+ }
102
+ ],
103
+ "prompt": "<|turn>user\nWrite a haiku about autumn leaves in a city.<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
104
+ "content": "<|channel>thought\n<channel|>Gold leaves drift on stone,\nDancing through the concrete maze,\nStreetlights catch the glow.",
105
+ "finish": "stop",
106
+ "tokens": 21
107
+ },
108
+ {
109
+ "name": "system-pirate",
110
+ "messages": [
111
+ {
112
+ "role": "system",
113
+ "content": "Answer like a cheerful pirate, in two sentences."
114
+ },
115
+ {
116
+ "role": "user",
117
+ "content": "What is the boiling point of water at sea level?"
118
+ }
119
+ ],
120
+ "prompt": "<|turn>system\nAnswer like a cheerful pirate, in two sentences.<turn|>\n<|turn>user\nWhat is the boiling point of water at sea level?<turn|>\n<|turn>model\n<|channel>thought\n<channel|>",
121
+ "content": "<|channel>thought\n<channel|>Ahoy there, matey, the waters reach a boil at a sweltering 212 degrees Fahrenheit! Keep a weather eye on that pot so ye don't scald yer treasure!",
122
+ "finish": "stop",
123
+ "tokens": 41
124
+ }
125
+ ],
126
+ "speed": {
127
+ "cache_n": 0,
128
+ "prompt_n": 24,
129
+ "prompt_ms": 117.416,
130
+ "prompt_per_token_ms": 4.892333333333333,
131
+ "prompt_per_second": 204.40144443687404,
132
+ "predicted_n": 128,
133
+ "predicted_ms": 1812.329,
134
+ "predicted_per_token_ms": 14.1588203125,
135
+ "predicted_per_second": 70.62735298061224
136
+ }
137
+ }
examples/chat/refs/qwen.json ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "about": "No reference text: llama.cpp's Metal path cannot hold the 36.9 GB model on the 24 GB Mac these runs used, and its CPU path rounds differently. The gate only checks that each conversation gets a reply, and records speed.",
3
+ "conversations": [
4
+ {
5
+ "name": "explain-tcp",
6
+ "messages": [
7
+ {
8
+ "role": "user",
9
+ "content": "Explain the difference between TCP and UDP to a new programmer."
10
+ }
11
+ ],
12
+ "tokens": 64
13
+ },
14
+ {
15
+ "name": "json-extract",
16
+ "messages": [
17
+ {
18
+ "role": "user",
19
+ "content": "Extract the name, city and age from this sentence as JSON: \"Priya, 34, moved from Pune to Lisbon last year.\""
20
+ }
21
+ ],
22
+ "tokens": 64
23
+ },
24
+ {
25
+ "name": "translate-hi",
26
+ "messages": [
27
+ {
28
+ "role": "user",
29
+ "content": "Translate into Hindi: \"The library opens at nine and closes at six.\""
30
+ }
31
+ ],
32
+ "tokens": 64
33
+ }
34
+ ]
35
+ }
index.html CHANGED
@@ -1,19 +1,5 @@
1
  <!doctype html>
2
- <html>
3
- <head>
4
- <meta charset="utf-8" />
5
- <meta name="viewport" content="width=device-width" />
6
- <title>My static Space</title>
7
- <link rel="stylesheet" href="style.css" />
8
- </head>
9
- <body>
10
- <div class="card">
11
- <h1>Welcome to your static Space!</h1>
12
- <p>You can modify this app directly by editing <i>index.html</i> in the Files and versions tab.</p>
13
- <p>
14
- Also don't forget to check the
15
- <a href="https://huggingface.co/docs/hub/spaces" target="_blank">Spaces documentation</a>.
16
- </p>
17
- </div>
18
- </body>
19
- </html>
 
1
  <!doctype html>
2
+ <meta charset="utf-8">
3
+ <meta http-equiv="refresh" content="0; url=examples/chat/">
4
+ <title>diskformer chat</title>
5
+ <a href="examples/chat/">diskformer chat</a>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/expert-stream.js ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* moe-expert-stream.js — keep a mixture-of-experts model's routed experts on SSD (OPFS) and
2
+ * page them into a fixed GPU slot pool on demand. Model-agnostic: an engine describes one
3
+ * expert record (its byte layout in the OPFS file and which GPU pool buffer each part goes
4
+ * to); this module owns the slot LRU, the reads, the uploads, the pins and the counters.
5
+ *
6
+ * const xs = new ExpertStreamer({ device, reader, recordBytes, recordOffset, parts, slots,
7
+ * numLayers, numExperts });
8
+ * const slots = await xs.ensure(layer, expertIds); // resident + pinned, in id order
9
+ * …write `slots` into the kernel's slot table, encode, queue.submit()…
10
+ * xs.release(slots); // after the submit
11
+ * xs.prefetch(layer + 1, predictedIds); // optional, never blocks
12
+ *
13
+ * `parts`: [{ buffer: GPUBuffer, srcOffset, bytes }]. Slot s of a part lives at s × bytes in
14
+ * that part's buffer, so a kernel finds expert row r of slot s at s × rowsPerExpert + r.
15
+ *
16
+ * Safety of reuse: an upload (a staging-ring copy, or queue.writeBuffer) is ordered after every
17
+ * earlier queue.submit, so a slot may be overwritten as soon as the work that reads it has been
18
+ * submitted, even if the GPU has not run it yet. Until then the slot is pinned and never chosen
19
+ * for eviction.
20
+ */
21
+
22
+ export class ExpertStreamer {
23
+ constructor({ device, reader, recordBytes, recordOffset, parts, slots, numLayers, numExperts, maxStaging = 32, uploadRing = 16, evict = 'lru', hotHalfLife = 8, maxPrefetchInflight = 8 }) {
24
+ Object.assign(this, { device, reader, recordBytes, recordOffset, parts, slots, numLayers, numExperts, maxStaging, evict, maxPrefetchInflight });
25
+ // Prefetches wait in pfQueue (nearest layer first) and at most maxPrefetchInflight of them
26
+ // read at once, so a demand read never queues behind a long run of guesses in the workers.
27
+ this.pfQueue = [];
28
+ this.pfInflight = 0;
29
+ // evict 'hot': drop the unpinned expert with the lowest decaying route count, oldest first
30
+ // on ties (llama.cpp PR #25294's policy). One tick per ensure() call, so the half-life is in
31
+ // tokens × layers. 'lru' (default): plain least recently used.
32
+ this.heat = new Float32Array(numLayers * numExperts);
33
+ this.heatTick = new Float64Array(numLayers * numExperts);
34
+ this.tick = 0;
35
+ this.decayPerTick = Math.pow(0.5, 1 / Math.max(1, hotHalfLife * numLayers));
36
+ // Uploads go through a ring of mapped MAP_WRITE buffers + copyBufferToBuffer: measured
37
+ // 2026-10-04 on an M4 Pro (Chrome 154) at 27.5 GB/s for 5 MB records with depth 4, against
38
+ // 1.88 GB/s for queue.writeBuffer. uploadRing = 0 keeps writeBuffer (A/B measurement only).
39
+ this.ring = [];
40
+ for (let i = 0; i < uploadRing; i++) {
41
+ this.ring.push({ buf: device.createBuffer({ size: Math.ceil(recordBytes / 4) * 4, usage: 0x02 | 0x04, mappedAtCreation: true }), ready: null });
42
+ }
43
+ this.ringNext = 0;
44
+ this.slotKey = new Int32Array(slots).fill(-1);
45
+ this.pinCount = new Int32Array(slots);
46
+ this.free = [];
47
+ for (let s = slots - 1; s >= 0; s--) this.free.push(s);
48
+ this.lru = new Map(); // key → slot, least recently used first
49
+ this.inflight = new Map(); // key → { p: Promise<slot>, pins }
50
+ this.staging = [];
51
+ this.resetStats();
52
+ }
53
+
54
+ resetStats() {
55
+ const L = this.numLayers;
56
+ this.stats = {
57
+ hits: 0, misses: 0, lateHits: 0, prefetchIssued: 0, prefetchUsed: 0, evictions: 0,
58
+ bytesRead: 0, readMs: 0, uploadMs: 0, waitMs: 0,
59
+ hitsByLayer: new Uint32Array(L), missesByLayer: new Uint32Array(L),
60
+ };
61
+ this.prefetched = new Set(); // keys loaded by prefetch() and not yet demanded
62
+ }
63
+
64
+ key(layer, expert) { return layer * this.numExperts + expert; }
65
+ resident() { return this.lru.size; }
66
+
67
+ pin(s) { this.pinCount[s]++; }
68
+ release(slots) { for (const s of slots) if (this.pinCount[s] > 0) this.pinCount[s]--; }
69
+
70
+ takeStaging() { return this.staging.pop() || new ArrayBuffer(this.recordBytes); }
71
+ giveStaging(buf) { if (this.staging.length < this.maxStaging && buf.byteLength >= this.recordBytes) this.staging.push(buf); }
72
+
73
+ heatOf(key) { return this.heat[key] * Math.pow(this.decayPerTick, this.tick - this.heatTick[key]); }
74
+ touch(key) { this.heat[key] = this.heatOf(key) + 1; this.heatTick[key] = this.tick; }
75
+
76
+ allocSlot() {
77
+ if (this.free.length) return this.free.pop();
78
+ if (this.evict === 'hot') {
79
+ let bestK = -1, bestS = -1, bestH = Infinity;
80
+ for (const [k, s] of this.lru) { // LRU order, so the first minimum is the oldest
81
+ if (this.pinCount[s] !== 0) continue;
82
+ const h = this.heatOf(k);
83
+ if (h < bestH) { bestH = h; bestK = k; bestS = s; }
84
+ }
85
+ if (bestS >= 0) {
86
+ this.lru.delete(bestK); this.prefetched.delete(bestK); this.slotKey[bestS] = -1; this.stats.evictions++;
87
+ return bestS;
88
+ }
89
+ throw new Error(`expert pool exhausted: all ${this.slots} slots pinned — raise the pool size`);
90
+ }
91
+ for (const [k, s] of this.lru) {
92
+ if (this.pinCount[s] === 0) {
93
+ this.lru.delete(k);
94
+ this.prefetched.delete(k);
95
+ this.slotKey[s] = -1;
96
+ this.stats.evictions++;
97
+ return s;
98
+ }
99
+ }
100
+ throw new Error(`expert pool exhausted: all ${this.slots} slots pinned — raise the pool size`);
101
+ }
102
+
103
+ // Starts the read + upload of one expert into a newly allocated slot. The slot stays pinned
104
+ // while the load is in flight; `entry.pins` demands registered meanwhile are converted into
105
+ // pins in the same synchronous step that publishes the slot, so no eviction can slip
106
+ // between "loaded" and "pinned by its user". Returns the in-flight entry { p, pins }.
107
+ load(layer, expert, pins = 0) {
108
+ const key = this.key(layer, expert);
109
+ const s = this.allocSlot();
110
+ this.pin(s);
111
+ const entry = { pins, p: null };
112
+ entry.p = (async () => {
113
+ let ok = false;
114
+ try {
115
+ const r = await this.reader.read(this.recordOffset(layer, expert), this.recordBytes, this.takeStaging());
116
+ if (r.got !== this.recordBytes) throw new Error(`short expert read L${layer} E${expert}: ${r.got} of ${this.recordBytes} bytes`);
117
+ this.stats.readMs += r.ms;
118
+ this.stats.bytesRead += r.got;
119
+ const t0 = performance.now();
120
+ await this.upload(r.buf, s);
121
+ this.stats.uploadMs += performance.now() - t0;
122
+ this.giveStaging(r.buf);
123
+ this.slotKey[s] = key;
124
+ this.lru.set(key, s);
125
+ this.pinCount[s] += entry.pins;
126
+ ok = true;
127
+ return s;
128
+ } catch (err) {
129
+ if (err && err.buf) this.giveStaging(err.buf);
130
+ throw err;
131
+ } finally {
132
+ this.inflight.delete(key);
133
+ this.pinCount[s]--;
134
+ if (!ok) { this.pinCount[s] = 0; this.free.push(s); }
135
+ }
136
+ })();
137
+ this.inflight.set(key, entry);
138
+ return entry;
139
+ }
140
+
141
+ // Copies one record into slot `s` of every part. Queue order makes this as safe as
142
+ // writeBuffer: the copy is submitted after the work that last read the slot and before the
143
+ // work that will read the new expert.
144
+ async upload(buf, s) {
145
+ if (!this.ring.length) {
146
+ for (const part of this.parts) this.device.queue.writeBuffer(part.buffer, s * part.bytes, buf, part.srcOffset, part.bytes);
147
+ return;
148
+ }
149
+ const r = this.ring[this.ringNext];
150
+ this.ringNext = (this.ringNext + 1) % this.ring.length;
151
+ while (r.ready) { const p = r.ready; await p; if (r.ready === p) r.ready = null; }
152
+ new Uint8Array(r.buf.getMappedRange(0, r.buf.size)).set(new Uint8Array(buf, 0, this.recordBytes));
153
+ r.buf.unmap();
154
+ const enc = this.device.createCommandEncoder();
155
+ for (const part of this.parts) enc.copyBufferToBuffer(r.buf, part.srcOffset, part.buffer, s * part.bytes, part.bytes);
156
+ this.device.queue.submit([enc.finish()]);
157
+ r.ready = r.buf.mapAsync(2);
158
+ }
159
+
160
+ // Makes every expert in `ids` resident for `layer` and pins its slot. Returns the slots in
161
+ // the same order as `ids`. The caller must release() them after submitting the GPU work.
162
+ async ensure(layer, ids) {
163
+ this.tick++;
164
+ if (this.pfQueue.length) this.pfQueue = this.pfQueue.filter((q) => q.layer > layer); // not sorted by layer once lookahead >= 3
165
+ const out = new Uint32Array(ids.length);
166
+ const waits = [];
167
+ const st = this.stats;
168
+ for (let i = 0; i < ids.length; i++) {
169
+ const key = this.key(layer, ids[i]);
170
+ this.touch(key);
171
+ if (this.prefetched.delete(key)) st.prefetchUsed++;
172
+ const s = this.lru.get(key);
173
+ if (s !== undefined) {
174
+ this.lru.delete(key); this.lru.set(key, s); // most recently used
175
+ this.pin(s);
176
+ out[i] = s;
177
+ st.hits++; st.hitsByLayer[layer]++;
178
+ continue;
179
+ }
180
+ const flying = this.inflight.get(key);
181
+ if (flying) {
182
+ st.lateHits++; st.hitsByLayer[layer]++;
183
+ flying.pins++;
184
+ waits.push(flying.p.then((slot) => { out[i] = slot; }));
185
+ continue;
186
+ }
187
+ st.misses++; st.missesByLayer[layer]++;
188
+ waits.push(this.load(layer, ids[i], 1).p.then((slot) => { out[i] = slot; }));
189
+ }
190
+ if (waits.length) {
191
+ const t0 = performance.now();
192
+ await Promise.all(waits);
193
+ st.waitMs += performance.now() - t0;
194
+ }
195
+ return out;
196
+ }
197
+
198
+ // For an engine that routes on the GPU against a residency map: `ids` of `layer` were found
199
+ // resident and used without an ensure(). Refreshes their recency and counts them as hits; an
200
+ // expert evicted since (its slot was reused after that work was submitted) is skipped.
201
+ touchUsed(layer, ids) {
202
+ this.tick++;
203
+ const st = this.stats;
204
+ for (const e of ids) {
205
+ const key = this.key(layer, e);
206
+ const s = this.lru.get(key);
207
+ if (s === undefined) continue;
208
+ this.touch(key);
209
+ if (this.prefetched.delete(key)) st.prefetchUsed++;
210
+ this.lru.delete(key); this.lru.set(key, s);
211
+ st.hits++; st.hitsByLayer[layer]++;
212
+ }
213
+ }
214
+
215
+ // Queues loads for experts that are neither resident, in flight nor queued; never awaits.
216
+ // Call with the nearest layer first: the queue is served in order, and entries for a layer
217
+ // are dropped once that layer's demand (ensure) arrives. Returns how many loads it started.
218
+ prefetch(layer, ids) {
219
+ const before = this.stats.prefetchIssued;
220
+ for (const e of ids) {
221
+ const key = this.key(layer, e);
222
+ if (this.lru.has(key) || this.inflight.has(key) || this.pfQueue.some((q) => q.key === key)) continue;
223
+ this.pfQueue.push({ layer, e, key });
224
+ }
225
+ this.pump();
226
+ return this.stats.prefetchIssued - before;
227
+ }
228
+
229
+ // Starts queued prefetches while fewer than maxPrefetchInflight are reading. Stops when the
230
+ // pool has nothing evictable left.
231
+ pump() {
232
+ while (this.pfInflight < this.maxPrefetchInflight && this.pfQueue.length) {
233
+ const { layer, e, key } = this.pfQueue.shift();
234
+ if (this.lru.has(key) || this.inflight.has(key)) continue;
235
+ let entry;
236
+ try { entry = this.load(layer, e); } catch (_) { this.pfQueue.length = 0; break; }
237
+ this.prefetched.add(key);
238
+ this.stats.prefetchIssued++;
239
+ this.pfInflight++;
240
+ entry.p.catch(() => {}).finally(() => { this.pfInflight--; this.pump(); });
241
+ }
242
+ }
243
+
244
+ // Waits for every in-flight load and staging-ring map (before teardown and between runs).
245
+ async drain() {
246
+ await Promise.allSettled([...this.inflight.values()].map((f) => f.p));
247
+ await Promise.allSettled(this.ring.map((r) => r.ready).filter(Boolean));
248
+ }
249
+ ringBytes() { return this.ring.reduce((a, r) => a + r.buf.size, 0); }
250
+
251
+ // Empties the pool. Waits for loads still reading first: a load that finished after the
252
+ // reset would claim a slot that is back on the free list.
253
+ async clear() {
254
+ this.pfQueue = [];
255
+ await this.drain();
256
+ this.lru.clear();
257
+ this.inflight.clear();
258
+ this.prefetched.clear();
259
+ this.slotKey.fill(-1);
260
+ this.pinCount.fill(0);
261
+ this.free = [];
262
+ for (let s = this.slots - 1; s >= 0; s--) this.free.push(s);
263
+ }
264
+
265
+ }
src/gguf.js ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* gguf.js — read a GGUF header and split quantized blocks into GPU-friendly planes. No dependencies.
2
+ * Extracted unchanged from LocalMind's qwen3_moe_ssd.js (2026-10-05).
3
+ *
4
+ * const { kv, tensors, dataStart } = parseGguf(headerBytes); // throws { needBytes } if too short
5
+ * tensorBytes(tensor) // bytes of one tensor in the file
6
+ * splitQ8(src, q, s) / splitQ4(src, q, s) // blocks → value plane + f16-scale plane, bit-exact
7
+ */
8
+
9
+ export const GGML = { F32: 0, F16: 1, Q4_0: 2, Q8_0: 8, Q6_K: 14 };
10
+ export const Q8_BLOCK = 34; // f16 scale + 32 × int8
11
+ export const Q4_BLOCK = 18; // f16 scale + 32 × 4-bit (16 bytes)
12
+ export const Q6K_BLOCK = 210; // 256 values: ql[128] + qh[64] + scales[16] + f16 d
13
+
14
+ // Parses a GGUF header. Throws { needBytes } when `u8` stops before the header ends.
15
+ export function parseGguf(u8) {
16
+ const dv = new DataView(u8.buffer, u8.byteOffset, u8.byteLength);
17
+ let p = 0;
18
+ const need = (n) => { if (p + n > u8.byteLength) { const e = new Error('GGUF header truncated'); e.needBytes = Math.max(u8.byteLength * 2, p + n + (1 << 20)); throw e; } };
19
+ const u32 = () => { need(4); const v = dv.getUint32(p, true); p += 4; return v; };
20
+ const u64 = () => { need(8); const v = Number(dv.getBigUint64(p, true)); p += 8; return v; };
21
+ // ignoreBOM: keep a leading U+FEFF (Gemma 4's vocab has '#', '//' and '<?' twice, once with a byte-order mark).
22
+ const dec = new TextDecoder('utf-8', { ignoreBOM: true });
23
+ const str = () => { const n = u64(); need(n); const s = dec.decode(u8.subarray(p, p + n)); p += n; return s; };
24
+ const scalar = {
25
+ 0: () => { need(1); return dv.getUint8(p++); }, 1: () => { need(1); return dv.getInt8(p++); },
26
+ 2: () => { need(2); const v = dv.getUint16(p, true); p += 2; return v; }, 3: () => { need(2); const v = dv.getInt16(p, true); p += 2; return v; },
27
+ 4: u32, 5: () => { need(4); const v = dv.getInt32(p, true); p += 4; return v; },
28
+ 6: () => { need(4); const v = dv.getFloat32(p, true); p += 4; return v; },
29
+ 7: () => { need(1); return dv.getUint8(p++) !== 0; },
30
+ 10: u64, 11: () => { need(8); const v = Number(dv.getBigInt64(p, true)); p += 8; return v; },
31
+ 12: () => { need(8); const v = dv.getFloat64(p, true); p += 8; return v; },
32
+ };
33
+ const value = (t) => {
34
+ if (t === 8) return str();
35
+ if (t === 9) { const at = u32(), n = u64(); const a = new Array(n); for (let i = 0; i < n; i++) a[i] = value(at); return a; }
36
+ const f = scalar[t];
37
+ if (!f) throw new Error(`GGUF: unknown value type ${t}`);
38
+ return f();
39
+ };
40
+ need(4);
41
+ if (dec.decode(u8.subarray(0, 4)) !== 'GGUF') throw new Error('not a GGUF file');
42
+ p = 4;
43
+ const version = u32(), nTensors = u64(), nKv = u64();
44
+ const kv = {};
45
+ for (let i = 0; i < nKv; i++) { const k = str(); kv[k] = value(u32()); }
46
+ const tensors = [];
47
+ for (let i = 0; i < nTensors; i++) {
48
+ const name = str(), nd = u32(), dims = [];
49
+ for (let d = 0; d < nd; d++) dims.push(u64());
50
+ const type = u32(), offset = u64();
51
+ tensors.push({ name, dims, type, offset });
52
+ }
53
+ const align = kv['general.alignment'] || 32;
54
+ const dataStart = Math.ceil(p / align) * align;
55
+ return { version, kv, tensors, headerBytes: p, dataStart };
56
+ }
57
+
58
+ export function tensorBytes(t) {
59
+ const n = t.dims.reduce((a, b) => a * b, 1);
60
+ if (t.type === GGML.F32) return n * 4;
61
+ if (t.type === GGML.F16) return n * 2;
62
+ if (t.type === GGML.Q8_0) return (n / 32) * Q8_BLOCK;
63
+ if (t.type === GGML.Q4_0) return (n / 32) * Q4_BLOCK;
64
+ if (t.type === GGML.Q6_K) return (n / 256) * Q6K_BLOCK;
65
+ throw new Error(`unsupported tensor type ${t.type} (${t.name})`);
66
+ }
67
+
68
+ // Q8_0 blocks → an int8 plane + an f16-scale plane, bit-exact. `src` holds whole blocks and
69
+ // starts at an even byte offset, so both planes copy as u16 lanes.
70
+ export function splitQ8(src, q, s) {
71
+ const nb = src.byteLength / Q8_BLOCK;
72
+ const a = new Uint16Array(src.buffer, src.byteOffset, nb * 17);
73
+ const qd = new Uint16Array(q.buffer, q.byteOffset, nb * 16);
74
+ const sd = new Uint16Array(s.buffer, s.byteOffset, nb);
75
+ for (let b = 0, i = 0, o = 0; b < nb; b++, i += 17, o += 16) {
76
+ sd[b] = a[i];
77
+ qd[o] = a[i + 1]; qd[o + 1] = a[i + 2]; qd[o + 2] = a[i + 3]; qd[o + 3] = a[i + 4];
78
+ qd[o + 4] = a[i + 5]; qd[o + 5] = a[i + 6]; qd[o + 6] = a[i + 7]; qd[o + 7] = a[i + 8];
79
+ qd[o + 8] = a[i + 9]; qd[o + 9] = a[i + 10]; qd[o + 10] = a[i + 11]; qd[o + 11] = a[i + 12];
80
+ qd[o + 12] = a[i + 13]; qd[o + 13] = a[i + 14]; qd[o + 14] = a[i + 15]; qd[o + 15] = a[i + 16];
81
+ }
82
+ }
83
+
84
+ // Q4_0 blocks → a nibble plane (16 bytes per 32 values, ggml's order: low nibbles are values 0..15,
85
+ // high nibbles 16..31) + an f16-scale plane, bit-exact.
86
+ export function splitQ4(src, q, s) {
87
+ const nb = src.byteLength / Q4_BLOCK;
88
+ const a = new Uint16Array(src.buffer, src.byteOffset, nb * 9);
89
+ const qd = new Uint16Array(q.buffer, q.byteOffset, nb * 8);
90
+ const sd = new Uint16Array(s.buffer, s.byteOffset, nb);
91
+ for (let b = 0, i = 0, o = 0; b < nb; b++, i += 9, o += 8) {
92
+ sd[b] = a[i];
93
+ qd[o] = a[i + 1]; qd[o + 1] = a[i + 2]; qd[o + 2] = a[i + 3]; qd[o + 3] = a[i + 4];
94
+ qd[o + 4] = a[i + 5]; qd[o + 5] = a[i + 6]; qd[o + 6] = a[i + 7]; qd[o + 7] = a[i + 8];
95
+ }
96
+ }
src/ingest.js ADDED
@@ -0,0 +1,177 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* ingest.js — copy a GGUF from a URL into OPFS in an engine's own layout, streaming: HTTP Range requests of
2
+ * ≤1 GiB, each unit transformed and written while the next bytes arrive (at most 8 writes in flight). OPFS quota
3
+ * grows with what is written, so files grow by writes rather than one up-front truncate. The manifest is written
4
+ * incomplete first and complete last, so an interrupted ingest is never taken as valid. After each range request the
5
+ * manifest records the first unit not yet written; a later call with the same key and the same GGUF header resumes
6
+ * from there instead of from byte 0.
7
+ * Extracted from LocalMind's qwen3_moe_ssd.js (2026-10-05); `root` and `format` are new options (LocalMind
8
+ * passes 'localmind-ssd' and its own format string).
9
+ *
10
+ * const manifest = await ingestGguf({ url, key, plan: { layout, units }, onProgress });
11
+ * // OPFS: <root>/<key>/{header.bin, dense.bin, experts.bin, manifest.json}
12
+ */
13
+
14
+ import { OpfsWriter, readOpfsText, writeOpfsText, removeOpfs } from './opfs-reader.js';
15
+ import { parseGguf, splitQ8, splitQ4, Q8_BLOCK, Q4_BLOCK } from './gguf.js';
16
+
17
+ export const ROOT = 'diskformer';
18
+ export const FORMAT = 'diskformer-gguf/1';
19
+
20
+ // ── Ingest: GGUF (over HTTP Range) → OPFS, in the engine layout ─────────────
21
+ export const storeKey = (file) => file.replace(/\.gguf$/i, '').toLowerCase().replace(/[^a-z0-9._-]+/g, '-');
22
+
23
+ export async function readHeader(url, fetchFn, signal) {
24
+ let n = 16 << 20;
25
+ for (;;) {
26
+ const r = await fetchFn(url, { headers: { Range: `bytes=0-${n - 1}` }, signal, cache: 'no-store' });
27
+ if (!r.ok) throw new Error(`GGUF header fetch: HTTP ${r.status}`);
28
+ const h = (name) => (r.headers && r.headers.get ? r.headers.get(name) : null);
29
+ const total = /\/(\d+)\s*$/.exec(h('content-range') || '');
30
+ const file = { size: total ? Number(total[1]) : null, etag: h('etag') }; // for the resume identity
31
+ const u8 = new Uint8Array(await r.arrayBuffer());
32
+ try { return { gguf: parseGguf(u8), bytes: u8, file }; } catch (e) { if (!e.needBytes) throw e; n = e.needBytes; }
33
+ }
34
+ }
35
+
36
+ // FNV-1a, for the resume identity: over bytes (the GGUF header) or a string (the plan's output layout).
37
+ const fnv = (at, n) => { let h = 0x811c9dc5; for (let i = 0; i < n; i++) h = Math.imul(h ^ at(i), 0x01000193) >>> 0; return `${n}:${h.toString(16)}`; };
38
+ const bytesId = (u8) => fnv((i) => u8[i], u8.length);
39
+ const textId = (s) => fnv((i) => s.charCodeAt(i), s.length);
40
+
41
+ // plan.layout(gguf) → { experts: { file, bytes, … }, dense: { file, bytes, tensors } }: where every kept byte goes.
42
+ // plan.units(gguf, layout) → sorted units { src, len, file: 'dense' | 'experts', raw } or { …, kind?: 'q4', q, s }:
43
+ // byte ranges of the GGUF, each small enough to transform in memory (raw copies, or Q8_0 / Q4_0 blocks split
44
+ // into a value plane at q and an f16-scale plane at s).
45
+ // An interrupted ingest resumes only when the GGUF header, the plan's output layout, and the source's size and ETag
46
+ // (when the server sends them) match its checkpoint; otherwise it starts over. `rangeBytes` and `io` are for tests:
47
+ // smaller range requests, and fake OPFS files.
48
+ export async function ingestGguf({ url, key, plan, root = ROOT, format = FORMAT, fetch: fetchFn = fetch, onProgress = () => {}, signal, source = {},
49
+ rangeBytes = 1 << 30, io = { OpfsWriter, readOpfsText, writeOpfsText } }) {
50
+ if (!plan || !plan.layout || !plan.units) throw new Error('ingestGguf: plan { layout, units } is required');
51
+ const t0 = performance.now();
52
+ const { gguf, bytes: headerBuf, file } = await readHeader(url, fetchFn, signal);
53
+ const layout = plan.layout(gguf);
54
+ const dir = `${root}/${key}`;
55
+ const storage = typeof navigator !== 'undefined' && navigator.storage;
56
+ const est = storage && storage.estimate ? await storage.estimate() : {};
57
+ const persisted = storage && storage.persist ? await storage.persist().catch(() => false) : false;
58
+ const needed = layout.experts.bytes + layout.dense.bytes;
59
+ onProgress({ status: 'ingest-plan', needed, quota: est.quota, usage: est.usage, persisted });
60
+ const quotaLog = [{ written: 0, quota: est.quota, usage: est.usage }];
61
+ const units = plan.units(gguf, layout);
62
+ const identity = [bytesId(headerBuf.subarray(0, gguf.dataStart)), textId(JSON.stringify([layout, units])), file.size ?? '', file.etag ?? ''].join('|');
63
+ let prior = null;
64
+ try { prior = JSON.parse(await io.readOpfsText(`${dir}/manifest.json`) || 'null'); } catch (_) { prior = null; }
65
+ const r0 = prior && !prior.complete && prior.format === format && prior.resume;
66
+ const startUnit = r0 && r0.identity === identity && r0.unit > 0 && r0.unit < units.length ? r0.unit : 0;
67
+ const checkpoint = (unit) => io.writeOpfsText(`${dir}/manifest.json`, JSON.stringify({ format, complete: false, resume: { identity, unit } }));
68
+
69
+ // Stream the data section in order. Units are filled from the response chunks, transformed,
70
+ // and written while the next bytes arrive (at most `maxWrites` writes in flight).
71
+ const first = units[0].src, last = units[units.length - 1].src + units[units.length - 1].len;
72
+ const resumedFrom = units[startUnit].src - first;
73
+ const maxUnit = units.reduce((m, u) => Math.max(m, u.len), 0);
74
+ let ui = startUnit, fill = 0, unitBuf = new Uint8Array(maxUnit);
75
+ const writers = {};
76
+ const release = async (name) => { const w = writers[name]; delete writers[name]; await w.close(); };
77
+ const inflight = new Set(); const maxWrites = 8;
78
+ const spare = [];
79
+ const getBuf = (n) => { const i = spare.findIndex((b) => b.byteLength >= n); return i >= 0 ? spare.splice(i, 1)[0] : new ArrayBuffer(n); };
80
+ const queue = async (w, buf, len, at) => {
81
+ while (inflight.size >= maxWrites) await Promise.race(inflight);
82
+ const p = w.writeOwned(buf, len, at).then((b) => { inflight.delete(p); if (spare.length < 16) spare.push(b); });
83
+ p.catch(() => {}); // a failed write surfaces through the race above or the final Promise.all
84
+ inflight.add(p);
85
+ };
86
+ let done = resumedFrom, written = 0;
87
+ const flushUnit = async (u) => {
88
+ const src = unitBuf.subarray(0, u.len);
89
+ const w = writers[u.file];
90
+ if (u.raw !== undefined) {
91
+ const b = getBuf(u.len); new Uint8Array(b, 0, u.len).set(src);
92
+ await queue(w, b, u.len, u.raw); written += u.len; return;
93
+ }
94
+ const q4 = u.kind === 'q4', blk = q4 ? Q4_BLOCK : Q8_BLOCK, qPer = q4 ? 16 : 32;
95
+ const nb = u.len / blk;
96
+ const qb = getBuf(nb * qPer), sb = getBuf(nb * 2);
97
+ (q4 ? splitQ4 : splitQ8)(src, new Uint8Array(qb, 0, nb * qPer), new Uint8Array(sb, 0, nb * 2));
98
+ await queue(w, qb, nb * qPer, u.q);
99
+ await queue(w, sb, nb * 2, u.s);
100
+ written += nb * blk;
101
+ };
102
+ try {
103
+ if (!startUnit) {
104
+ await checkpoint(0);
105
+ // Header bytes, for the tokenizer and metadata at every later load.
106
+ writers.header = await io.OpfsWriter.open(`${dir}/header.bin`, { truncate: true });
107
+ await writers.header.write(headerBuf.subarray(0, gguf.dataStart), 0);
108
+ await release('header');
109
+ } else {
110
+ onProgress({ status: 'ingest-resume', unit: startUnit, units: units.length, loaded: resumedFrom, total: last - first });
111
+ }
112
+ // No up-front truncate to full size: a fresh origin's quota is ~10 GiB and grows with what is
113
+ // actually written (measured 2026-10-02), so a single 30.8 GB extension would be refused. The
114
+ // writes below grow each file by at most one layer's records (~642 MB) at a time.
115
+ writers.dense = await io.OpfsWriter.open(`${dir}/dense.bin`, { truncate: !startUnit });
116
+ writers.experts = await io.OpfsWriter.open(`${dir}/experts.bin`, { truncate: !startUnit });
117
+ // Sequential range requests of rangeBytes (1 GiB) keep any one response bounded.
118
+ let pos = units[startUnit].src;
119
+ for (let start = pos; start < last; start += rangeBytes) {
120
+ const end = Math.min(last, start + rangeBytes) - 1;
121
+ const r = await fetchFn(url, { headers: { Range: `bytes=${start}-${end}` }, signal, cache: 'no-store' });
122
+ if (r.status !== 206) throw new Error(`GGUF range ${start}-${end}: HTTP ${r.status}`);
123
+ const reader = r.body.getReader();
124
+ for (;;) {
125
+ const { done: eof, value } = await reader.read();
126
+ if (eof) break;
127
+ let off = 0;
128
+ while (off < value.byteLength && ui < units.length) {
129
+ const u = units[ui];
130
+ if (pos < u.src) { const skip = Math.min(u.src - pos, value.byteLength - off); pos += skip; off += skip; continue; }
131
+ const take = Math.min(u.len - fill, value.byteLength - off);
132
+ unitBuf.set(value.subarray(off, off + take), fill);
133
+ fill += take; off += take; pos += take;
134
+ if (fill === u.len) { await flushUnit(u); ui++; fill = 0; }
135
+ }
136
+ done = pos - first;
137
+ onProgress({ status: 'ingest', loaded: done, total: last - first, resumedFrom, written, secs: (performance.now() - t0) / 1000 });
138
+ if (written - quotaLog[quotaLog.length - 1].written > 4 * 2 ** 30 && storage && storage.estimate) {
139
+ const e = await storage.estimate();
140
+ quotaLog.push({ written, quota: e.quota, usage: e.usage });
141
+ }
142
+ }
143
+ // Every unit before ui is written: flush, then record ui as the resume point.
144
+ if (ui < units.length) {
145
+ await Promise.all(inflight);
146
+ await writers.dense.flush(); await writers.experts.flush();
147
+ await checkpoint(ui);
148
+ }
149
+ }
150
+ if (ui !== units.length) throw new Error(`ingest ended early: ${ui} of ${units.length} units`);
151
+ await Promise.all(inflight);
152
+ await release('dense');
153
+ await release('experts');
154
+ } catch (e) {
155
+ // Release every file still open, so a retry in this page can reopen them and resume.
156
+ await Promise.allSettled([...inflight]);
157
+ await Promise.allSettled(Object.keys(writers).map(release));
158
+ throw e;
159
+ }
160
+ const manifest = {
161
+ format, complete: true, ingestedAt: new Date().toISOString(),
162
+ source: { url, ...source, dataStart: gguf.dataStart, headerBytes: gguf.headerBytes },
163
+ ingestSecs: (performance.now() - t0) / 1000, resumedAtUnit: startUnit || undefined, quotaBefore: est.quota, persisted, quotaLog,
164
+ ...layout,
165
+ };
166
+ await io.writeOpfsText(`${dir}/manifest.json`, JSON.stringify(manifest));
167
+ return manifest;
168
+ }
169
+
170
+ // Ingest progress as hosts read it: a 'weights' event with byte counts. The spread goes first, so the
171
+ // event's own status ('ingest') cannot overwrite 'weights' (LocalMind's 2026-10-04 bug: a host's load
172
+ // watchdog, seeing no progress, killed a 37 GB first download).
173
+ export function ingestProgress(e) {
174
+ return e && e.status === 'ingest' ? { ...e, status: 'weights', loaded: e.loaded, total: e.total } : e;
175
+ }
176
+
177
+ export async function removeStore(key, { root = ROOT } = {}) { await removeOpfs(`${root}/${key}`); }
src/measure.js ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* measure.js — how fast this machine reads an OPFS file the way a disk tier does. Run it where the tier will run
2
+ * (a page or a worker); OPFS speed varies a lot between machines and browsers.
3
+ *
4
+ * const r = await measure(); // writes a 256 MB probe, measures, removes it
5
+ * const r = await measure({ path: 'diskformer/<key>/experts.bin', recordBytes: 3345408 }); // an existing store
6
+ * // → { bytes, sequentialMBps, random: { recordBytes, reads, p50Ms, p90Ms, MBps }, burst8: { p50Ms, MBps } }
7
+ *
8
+ * A probe the call just wrote is likely in the OS page cache, so its numbers are an upper bound; point `path` at a
9
+ * store larger than RAM to see the disk itself.
10
+ */
11
+ import { OpfsReaderPool, OpfsWriter, removeOpfs } from './opfs-reader.js';
12
+
13
+ const pct = (a, q) => { const s = [...a].sort((x, y) => x - y); return s[Math.min(s.length - 1, Math.floor(q * s.length))]; };
14
+
15
+ export async function measure({ path = null, bytes = 256 << 20, recordBytes = 3 << 20, reads = 64, workers = 4, chunkBytes = 64 << 20 } = {}) {
16
+ const probe = !path;
17
+ if (probe) {
18
+ path = 'diskformer/.measure/probe.bin';
19
+ const w = await OpfsWriter.open(path, { truncate: true });
20
+ const chunk = new Uint8Array(Math.min(chunkBytes, bytes));
21
+ for (let i = 0; i < chunk.length; i += 4096) chunk[i] = i & 255;
22
+ for (let at = 0; at < bytes; at += chunk.length) await w.write(chunk.subarray(0, Math.min(chunk.length, bytes - at)), at);
23
+ await w.close();
24
+ }
25
+ const pool = await OpfsReaderPool.open(path, { workers });
26
+ try {
27
+ const size = pool.size;
28
+ // Sequential: the load pattern (64 MB chunks, one at a time per worker).
29
+ let t0 = performance.now(), done = 0, inflight = [];
30
+ for (let at = 0; at < size; at += chunkBytes) {
31
+ inflight.push(pool.read(at, Math.min(chunkBytes, size - at)).then((r) => { done += r.got; }));
32
+ if (inflight.length >= workers) { await Promise.all(inflight); inflight = []; }
33
+ }
34
+ await Promise.all(inflight);
35
+ const sequentialMBps = done / 1e6 / ((performance.now() - t0) / 1000);
36
+ // Random whole records: the per-token pattern, one at a time, then bursts of 8 (one MoE layer's misses).
37
+ const n = Math.max(1, Math.floor(size / recordBytes));
38
+ const at = () => Math.floor(Math.random() * n) * recordBytes;
39
+ // Throughput comes from the total time: browsers coarsen performance.now() (0.1 ms in a worker without
40
+ // cross-origin isolation), which rounds a single small read to zero.
41
+ const one = []; let t1 = performance.now();
42
+ for (let i = 0; i < reads; i++) { const s = performance.now(); await pool.read(at(), recordBytes); one.push(performance.now() - s); }
43
+ const oneMs = performance.now() - t1;
44
+ const burst = [], bursts = Math.max(1, Math.round(reads / 8)); t1 = performance.now();
45
+ for (let i = 0; i < bursts; i++) {
46
+ const s = performance.now();
47
+ await Promise.all(Array.from({ length: 8 }, () => pool.read(at(), recordBytes)));
48
+ burst.push(performance.now() - s);
49
+ }
50
+ const burstMs = performance.now() - t1;
51
+ const mbps = (count, ms) => Math.round((count * recordBytes) / 1e6 / (Math.max(ms, 0.1) / 1000));
52
+ return {
53
+ path, bytes: size, workers, sequentialMBps: Math.round(sequentialMBps),
54
+ random: { recordBytes, reads, p50Ms: +pct(one, 0.5).toFixed(2), p90Ms: +pct(one, 0.9).toFixed(2), MBps: mbps(reads, oneMs) },
55
+ burst8: { bursts, p50Ms: +pct(burst, 0.5).toFixed(2), MBps: mbps(8 * bursts, burstMs) },
56
+ };
57
+ } finally {
58
+ await pool.close();
59
+ if (probe) await removeOpfs('diskformer/.measure').catch(() => {});
60
+ }
61
+ }
src/opfs-reader.js ADDED
@@ -0,0 +1,258 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* opfs-reader.js — read model weights from the Origin Private File System (OPFS) fast.
2
+ *
3
+ * Only OPFS has the fast path: FileSystemSyncAccessHandle, which exists only in dedicated
4
+ * workers. Measured 2026-10-02 on an M4 Pro, Chrome 153: 1.51 GB/s per worker for 1.3 MB
5
+ * random reads, 4.3–5.9 GB/s with 4 read-only workers; File.slice().arrayBuffer() reached
6
+ * only 0.39 GB/s. (~/Code/knowledge/notes/browser-tab-can-stream-weights-from-ssd.md)
7
+ *
8
+ * Shared by the SSD-streaming engines (Qwen3-MoE experts, Gemma 4 PLE rows). No engine
9
+ * dependencies. Works from a page or from inside an engine worker: the reader and writer
10
+ * workers are spawned from inlined Blob URLs, so there is no extra file to vendor.
11
+ *
12
+ * const pool = await OpfsReaderPool.open('localmind-ssd/<model>/experts.bin', { workers: 4 });
13
+ * const { buf } = await pool.read(offset, length, recycled); // one contiguous range
14
+ * const { buf } = await pool.readv([{ offset, length }, …], into); // many small ranges, packed
15
+ * pool.close();
16
+ *
17
+ * const w = await OpfsWriter.open('localmind-ssd/<model>/experts.bin', { truncate: true });
18
+ * await w.write(u8, at); await w.close(); // a read-write handle is exclusive: close it first
19
+ *
20
+ * Buffers move by transfer, never by copy; pass a previously returned buffer back in to
21
+ * recycle it.
22
+ *
23
+ * Inside a dedicated worker, `{ inline: true }` (see canInline()) skips the worker hop: the
24
+ * sync handle lives in the calling thread, and the pool gains readSync() for a per-token
25
+ * critical path where a postMessage round trip would cost more than the read.
26
+ */
27
+
28
+ // The file operations, shared by the worker and the inline (same-thread) mode. Kept
29
+ // self-contained: it is stringified into the worker's Blob source.
30
+ function makeOps() {
31
+ let handle = null;
32
+ const openFile = async (path, create) => {
33
+ const parts = path.split('/').filter(Boolean);
34
+ let dir = await navigator.storage.getDirectory();
35
+ for (const p of parts.slice(0, -1)) dir = await dir.getDirectoryHandle(p, { create });
36
+ return dir.getFileHandle(parts[parts.length - 1], { create });
37
+ };
38
+ const readFull = (u8, at) => {
39
+ let got = 0;
40
+ while (got < u8.byteLength) {
41
+ const n = handle.read(u8.subarray(got), { at: at + got });
42
+ if (n <= 0) break;
43
+ got += n;
44
+ }
45
+ return got;
46
+ };
47
+ // → [reply, transferList]
48
+ const run = async (m) => {
49
+ try {
50
+ if (m.op === 'open') {
51
+ const fh = await openFile(m.path, !!m.create);
52
+ handle = await fh.createSyncAccessHandle(m.mode === 'readwrite' ? {} : { mode: 'read-only' });
53
+ if (m.truncate) handle.truncate(0);
54
+ return [{ id: m.id, size: handle.getSize() }, []];
55
+ }
56
+ if (m.op === 'read') {
57
+ const t0 = performance.now();
58
+ const got = readFull(new Uint8Array(m.buf, 0, m.length), m.at);
59
+ return [{ id: m.id, buf: m.buf, got, ms: performance.now() - t0 }, [m.buf]];
60
+ }
61
+ if (m.op === 'readv') {
62
+ const t0 = performance.now();
63
+ let pos = 0, got = 0;
64
+ for (let i = 0; i < m.ranges.length; i += 2) {
65
+ const len = m.ranges[i + 1];
66
+ got += readFull(new Uint8Array(m.buf, pos, len), m.ranges[i]);
67
+ pos += len;
68
+ }
69
+ return [{ id: m.id, buf: m.buf, got, ms: performance.now() - t0 }, [m.buf]];
70
+ }
71
+ if (m.op === 'write') {
72
+ const u8 = new Uint8Array(m.buf, 0, m.length);
73
+ let put = 0;
74
+ while (put < u8.byteLength) put += handle.write(u8.subarray(put), { at: m.at + put });
75
+ return [{ id: m.id, buf: m.buf, put }, [m.buf]];
76
+ }
77
+ if (m.op === 'truncate') { handle.truncate(m.size); return [{ id: m.id, size: handle.getSize() }, []]; }
78
+ if (m.op === 'flush') { handle.flush(); return [{ id: m.id, size: handle.getSize() }, []]; }
79
+ if (m.op === 'close') {
80
+ if (handle) { if (m.flush) handle.flush(); handle.close(); }
81
+ handle = null;
82
+ return [{ id: m.id, closed: true }, []];
83
+ }
84
+ return [{ id: m.id, error: `unknown op ${m.op}` }, []];
85
+ } catch (err) {
86
+ return [{ id: m.id, error: `${(err && err.name) || 'Error'}: ${(err && err.message) || err}`, buf: m.buf }, m.buf ? [m.buf] : []];
87
+ }
88
+ };
89
+ run.sync = {
90
+ read: (buf, length, at) => readFull(new Uint8Array(buf, 0, length), at),
91
+ write: (u8, at) => { let put = 0; while (put < u8.byteLength) put += handle.write(u8.subarray(put), { at: at + put }); return put; },
92
+ };
93
+ return run;
94
+ }
95
+
96
+ let workerUrl = null;
97
+ const spawn = () => {
98
+ if (!workerUrl) {
99
+ const src = `const run = (${makeOps.toString()})();\nself.onmessage = async (e) => { const [r, t] = await run(e.data); self.postMessage(r, t); };`;
100
+ workerUrl = URL.createObjectURL(new Blob([src], { type: 'text/javascript' }));
101
+ }
102
+ return new Worker(workerUrl);
103
+ };
104
+
105
+ // Inline mode needs a sync access handle in this thread: only a dedicated worker has one.
106
+ export const canInline = () => typeof FileSystemSyncAccessHandle !== 'undefined'
107
+ && typeof DedicatedWorkerGlobalScope !== 'undefined' && self instanceof DedicatedWorkerGlobalScope;
108
+
109
+ // Same interface as Channel, but the file operations run in this thread.
110
+ class InlineChannel {
111
+ constructor() { this.run = makeOps(); this.inflight = 0; this.nextId = 1; }
112
+ async call(msg) {
113
+ this.inflight++;
114
+ try {
115
+ const [r] = await this.run({ ...msg, id: this.nextId++ });
116
+ if (r.error) { const err = new Error(r.error); err.buf = r.buf; throw err; }
117
+ return r;
118
+ } finally { this.inflight--; }
119
+ }
120
+ terminate() {}
121
+ }
122
+
123
+ // One worker, request/response matched by id; requests run in arrival order.
124
+ class Channel {
125
+ constructor() {
126
+ this.worker = spawn();
127
+ this.pending = new Map();
128
+ this.nextId = 1;
129
+ this.inflight = 0;
130
+ this.worker.onmessage = (e) => {
131
+ const m = e.data, p = this.pending.get(m.id);
132
+ if (!p) return;
133
+ this.pending.delete(m.id);
134
+ this.inflight--;
135
+ if (m.error) { const err = new Error(m.error); err.buf = m.buf; p.reject(err); } else p.resolve(m);
136
+ };
137
+ this.worker.onerror = (e) => {
138
+ for (const p of this.pending.values()) p.reject(new Error(`OPFS worker failed: ${e.message || e}`));
139
+ this.pending.clear();
140
+ };
141
+ }
142
+ call(msg, transfer = []) {
143
+ const id = this.nextId++;
144
+ this.inflight++;
145
+ return new Promise((resolve, reject) => {
146
+ this.pending.set(id, { resolve, reject });
147
+ this.worker.postMessage({ ...msg, id }, transfer);
148
+ });
149
+ }
150
+ terminate() { this.worker.terminate(); }
151
+ }
152
+
153
+ const bufferFor = (into, length) => (into && into.byteLength >= length ? into : new ArrayBuffer(length));
154
+
155
+ export class OpfsReaderPool {
156
+ // inline: true → no workers; reads run on this thread's own sync handle (dedicated
157
+ // worker only, see canInline()), and readSync() is available for a critical path.
158
+ static async open(path, { workers = 4, inline = false } = {}) {
159
+ const pool = new OpfsReaderPool();
160
+ pool.path = path;
161
+ pool.inline = inline;
162
+ pool.channels = [];
163
+ if (inline) pool.channels.push(new InlineChannel());
164
+ else for (let i = 0; i < workers; i++) pool.channels.push(new Channel());
165
+ const sizes = await Promise.all(pool.channels.map((c) => c.call({ op: 'open', path })));
166
+ pool.size = sizes[0].size;
167
+ return pool;
168
+ }
169
+ pick() {
170
+ let best = this.channels[0];
171
+ for (const c of this.channels) if (c.inflight < best.inflight) best = c;
172
+ return best;
173
+ }
174
+ // → { buf, got, ms }; `buf` is `into` when it was large enough, else a new ArrayBuffer.
175
+ read(offset, length, into) {
176
+ const buf = bufferFor(into, length);
177
+ return this.pick().call({ op: 'read', at: offset, length, buf }, [buf]);
178
+ }
179
+ // Inline mode only: a synchronous read, no promise. → bytes read.
180
+ readSync(offset, length, into) {
181
+ if (!this.inline) throw new Error('readSync needs OpfsReaderPool.open(path, { inline: true })');
182
+ return this.channels[0].run.sync.read(into, length, offset);
183
+ }
184
+ // ranges: [{ offset, length }]. The bytes land back-to-back in `buf`, in order.
185
+ readv(ranges, into) {
186
+ let total = 0;
187
+ const flat = new Array(ranges.length * 2);
188
+ ranges.forEach((r, i) => { flat[2 * i] = r.offset; flat[2 * i + 1] = r.length; total += r.length; });
189
+ const buf = bufferFor(into, total);
190
+ return this.pick().call({ op: 'readv', ranges: flat, buf }, [buf]);
191
+ }
192
+ async close() {
193
+ await Promise.all(this.channels.map((c) => c.call({ op: 'close' }).catch(() => {})));
194
+ for (const c of this.channels) c.terminate();
195
+ this.channels = [];
196
+ }
197
+ }
198
+
199
+ export class OpfsWriter {
200
+ // inline: true → the read-write sync handle lives in this thread (dedicated worker only).
201
+ static async open(path, { truncate = false, inline = false } = {}) {
202
+ const w = new OpfsWriter();
203
+ w.path = path;
204
+ w.ch = inline ? new InlineChannel() : new Channel();
205
+ w.size = (await w.ch.call({ op: 'open', path, create: true, mode: 'readwrite', truncate })).size;
206
+ return w;
207
+ }
208
+ // Copies `u8` into a transferable buffer only when it is a view on part of a larger one.
209
+ // Resolves with the (recycled) ArrayBuffer once the bytes are written.
210
+ async write(u8, at) {
211
+ let buf = u8.buffer, length = u8.byteLength;
212
+ if (this.ch instanceof InlineChannel) { this.ch.run.sync.write(u8, at); return buf; }
213
+ if (u8.byteOffset !== 0 || buf.byteLength !== length || (typeof SharedArrayBuffer !== 'undefined' && buf instanceof SharedArrayBuffer)) {
214
+ buf = new ArrayBuffer(length);
215
+ new Uint8Array(buf).set(u8);
216
+ }
217
+ const r = await this.ch.call({ op: 'write', at, length, buf }, [buf]);
218
+ return r.buf;
219
+ }
220
+ // Like write(), but takes ownership of `buf` (no copy) — the caller gets it back when done.
221
+ async writeOwned(buf, length, at) {
222
+ const r = await this.ch.call({ op: 'write', at, length, buf }, [buf]);
223
+ return r.buf;
224
+ }
225
+ async truncate(size) { return (await this.ch.call({ op: 'truncate', size })).size; }
226
+ async flush() { return (await this.ch.call({ op: 'flush' })).size; }
227
+ async close() { await this.ch.call({ op: 'close', flush: true }); this.ch.terminate(); }
228
+ }
229
+
230
+ export async function opfsDir(path, { create = false } = {}) {
231
+ let dir = await navigator.storage.getDirectory();
232
+ for (const p of path.split('/').filter(Boolean)) dir = await dir.getDirectoryHandle(p, { create });
233
+ return dir;
234
+ }
235
+
236
+ export async function readOpfsText(path) {
237
+ const parts = path.split('/').filter(Boolean);
238
+ try {
239
+ const dir = await opfsDir(parts.slice(0, -1).join('/'));
240
+ const f = await (await dir.getFileHandle(parts[parts.length - 1])).getFile();
241
+ return await f.text();
242
+ } catch (_) { return null; }
243
+ }
244
+
245
+ export async function writeOpfsText(path, text) {
246
+ const parts = path.split('/').filter(Boolean);
247
+ const dir = await opfsDir(parts.slice(0, -1).join('/'), { create: true });
248
+ const fh = await dir.getFileHandle(parts[parts.length - 1], { create: true });
249
+ const w = await fh.createWritable();
250
+ await w.write(text);
251
+ await w.close();
252
+ }
253
+
254
+ export async function removeOpfs(path) {
255
+ const parts = path.split('/').filter(Boolean);
256
+ const dir = await opfsDir(parts.slice(0, -1).join('/'));
257
+ await dir.removeEntry(parts[parts.length - 1], { recursive: true });
258
+ }
src/rows.js ADDED
@@ -0,0 +1,230 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* rows.js — rows of a large lookup table (an embedding-style table) served from OPFS into a small GPU-resident
2
+ * cache. Extracted from LocalMind's ple-opfs.js (2026-10-05), where it serves Gemma 4 E2B/E4B's per-layer embedding
3
+ * table from disk; the Gemma-specific helpers stay in LocalMind. `root` is a new option (LocalMind passes
4
+ * 'localmind-ssd').
5
+ *
6
+ * RowFile fixed-size rows in one OPFS file plus manifest.json; synchronous reads through opfs-reader.js in
7
+ * inline mode, so it must live in a dedicated worker.
8
+ * RowCache `slots` rows on the GPU, split into planes (byte ranges of a row that live in separate GPU buffers),
9
+ * an O(1) LRU, and a CPU id→slot mirror (optionally a GPU u32 map). lookup(ids) makes every id
10
+ * resident and returns its slot, writing missing rows with queue.writeBuffer.
11
+ */
12
+
13
+ import { OpfsReaderPool, OpfsWriter, canInline, readOpfsText, writeOpfsText, removeOpfs } from './opfs-reader.js';
14
+
15
+ export const ROOT = 'diskformer';
16
+ const FORMAT = 'localmind-rows/1'; // unchanged from LocalMind, so a store it wrote stays valid
17
+
18
+ // FNV-1a over a few samples of each source buffer and their lengths: cheap enough to run on every
19
+ // load, and it changes when the model's weights change.
20
+ export function fingerprint(buffers, sample = 1 << 16) {
21
+ let h = 0x811c9dc5;
22
+ const mix = (b) => { h ^= b; h = Math.imul(h, 0x01000193) >>> 0; };
23
+ for (const u8 of buffers) {
24
+ for (const v of [u8.length & 0xff, (u8.length >>> 8) & 0xff, (u8.length >>> 16) & 0xff, (u8.length >>> 24) & 0xff]) mix(v);
25
+ const starts = [0, Math.max(0, (u8.length >> 1) - (sample >> 1)), Math.max(0, u8.length - sample)];
26
+ for (const s of starts) for (let i = s, e = Math.min(u8.length, s + sample); i < e; i++) mix(u8[i]);
27
+ }
28
+ return h.toString(16).padStart(8, '0');
29
+ }
30
+
31
+ // Fixed-size rows in one OPFS file, <root>/<key>/<name>, with manifest.json beside it.
32
+ // The file I/O is opfs-reader.js in inline mode: the sync handle lives in this dedicated worker,
33
+ // so a row read is a plain synchronous call with no worker hop.
34
+ export class RowFile {
35
+ #reader = null;
36
+
37
+ constructor({ key, name, rowBytes, rows, root = ROOT }) {
38
+ this.path = `${root}/${key}/${name}`;
39
+ this.manifestPath = `${root}/${key}/manifest.json`;
40
+ this.name = name;
41
+ this.rowBytes = rowBytes;
42
+ this.rows = rows;
43
+ this.manifest = null;
44
+ }
45
+
46
+ static async open({ key, name = 'rows.bin', rowBytes, rows, root = ROOT }) {
47
+ if (!canInline()) throw new Error('RowFile needs a dedicated worker (FileSystemSyncAccessHandle)');
48
+ const f = new RowFile({ key, name, rowBytes, rows, root });
49
+ try { f.manifest = JSON.parse(await readOpfsText(f.manifestPath)); } catch (_) { f.manifest = null; }
50
+ return f;
51
+ }
52
+
53
+ matches(fp) {
54
+ const m = this.manifest;
55
+ return !!m && m.format === FORMAT && m.complete === true && m.name === this.name &&
56
+ m.rowBytes === this.rowBytes && m.rows === this.rows && m.fingerprint === fp;
57
+ }
58
+
59
+ // fill(row0, count, dst) writes `count` rows starting at `row0` into dst (count * rowBytes bytes).
60
+ // The manifest is removed first and written last, so an interrupted write is never taken as valid.
61
+ async write(fp, fill, { batchRows = 8192, extra = {} } = {}) {
62
+ this.close();
63
+ await removeOpfs(this.manifestPath).catch(() => {});
64
+ this.manifest = null;
65
+ const total = this.rows * this.rowBytes;
66
+ const t0 = performance.now();
67
+ const w = await OpfsWriter.open(this.path, { truncate: true, inline: true });
68
+ try {
69
+ const buf = new Uint8Array(batchRows * this.rowBytes);
70
+ for (let row0 = 0; row0 < this.rows; row0 += batchRows) {
71
+ const n = Math.min(batchRows, this.rows - row0);
72
+ const view = buf.subarray(0, n * this.rowBytes);
73
+ fill(row0, n, view);
74
+ await w.write(view, row0 * this.rowBytes);
75
+ }
76
+ } finally { await w.close(); }
77
+ const writeMs = performance.now() - t0;
78
+ this.manifest = { format: FORMAT, name: this.name, rowBytes: this.rowBytes, rows: this.rows, bytes: total,
79
+ fingerprint: fp, complete: true, writtenAt: new Date().toISOString(), writeMs: Math.round(writeMs), ...extra };
80
+ await writeOpfsText(this.manifestPath, JSON.stringify(this.manifest, null, 1));
81
+ return writeMs;
82
+ }
83
+
84
+ async openRead() {
85
+ if (this.#reader) return;
86
+ const r = await OpfsReaderPool.open(this.path, { inline: true });
87
+ if (r.size !== this.rows * this.rowBytes) {
88
+ await r.close();
89
+ throw new Error(`RowFile ${this.name}: size ${r.size} != ${this.rows * this.rowBytes}`);
90
+ }
91
+ this.#reader = r;
92
+ }
93
+
94
+ // Reads rows [row, row + count) to the start of dst (a Uint8Array that starts its buffer).
95
+ readRows(row, count, dst) {
96
+ if (dst.byteOffset !== 0) throw new Error('RowFile.readRows: dst must start at offset 0 of its buffer');
97
+ const len = count * this.rowBytes;
98
+ const got = this.#reader.readSync(row * this.rowBytes, len, dst.buffer);
99
+ if (got !== len) throw new Error(`RowFile ${this.name}: read ${got} of ${len} bytes at row ${row}`);
100
+ }
101
+
102
+ // Closes the read handle. In inline mode the handle closes synchronously inside this call.
103
+ close() {
104
+ if (this.#reader) { this.#reader.close(); this.#reader = null; }
105
+ }
106
+ }
107
+
108
+ export class RowCache {
109
+ // file: an open RowFile. slots: rows held on the GPU. planes: [{ offset, bytes, buffer }], each
110
+ // row byte range [offset, offset + bytes) is stored at buffer[slot * bytes]. queue: GPUQueue.
111
+ // mapBuffer (optional): a GPU u32[file.rows] copy of the id->slot map, 0xFFFFFFFF when absent.
112
+ constructor({ file, slots, planes, queue, mapBuffer = null }) {
113
+ this.file = file;
114
+ this.slots = slots;
115
+ this.planes = planes;
116
+ this.queue = queue;
117
+ this.mapBuffer = mapBuffer;
118
+ this.slotOf = new Int32Array(file.rows).fill(-1);
119
+ this.idOf = new Int32Array(slots).fill(-1);
120
+ // LRU as a doubly linked list over slots; head = least recently used. Starts as 0..slots-1.
121
+ this.prev = new Int32Array(slots);
122
+ this.next = new Int32Array(slots);
123
+ for (let s = 0; s < slots; s++) { this.prev[s] = s - 1; this.next[s] = s + 1 < slots ? s + 1 : -1; }
124
+ this.head = 0;
125
+ this.tail = slots - 1;
126
+ this.stamp = new Uint32Array(slots);
127
+ this.clock = 0;
128
+ this.row = new Uint8Array(file.rowBytes);
129
+ this.u32 = new Uint32Array(1);
130
+ this.stats = { lookups: 0, ids: 0, misses: 0, readMs: 0, replays: 0 };
131
+ this.missLog = []; // the first 4,096 ids installed after load, for analysis
132
+ }
133
+
134
+ #touch(s) {
135
+ if (s === this.tail) return;
136
+ const p = this.prev[s], n = this.next[s];
137
+ if (p >= 0) this.next[p] = n; else this.head = n;
138
+ this.prev[n] = p;
139
+ this.prev[s] = this.tail;
140
+ this.next[s] = -1;
141
+ this.next[this.tail] = s;
142
+ this.tail = s;
143
+ }
144
+
145
+ #mapWrite(id, slot) {
146
+ if (!this.mapBuffer) return;
147
+ this.u32[0] = slot >>> 0;
148
+ this.queue.writeBuffer(this.mapBuffer, id * 4, this.u32);
149
+ }
150
+
151
+ #install(id, s) {
152
+ const old = this.idOf[s];
153
+ if (old >= 0) { this.slotOf[old] = -1; this.#mapWrite(old, 0xffffffff); }
154
+ const t0 = performance.now();
155
+ this.file.readRows(id, 1, this.row);
156
+ this.stats.readMs += performance.now() - t0;
157
+ for (const p of this.planes) this.queue.writeBuffer(p.buffer, s * p.bytes, this.row, p.offset, p.bytes);
158
+ this.idOf[s] = id;
159
+ this.slotOf[id] = s;
160
+ this.#mapWrite(id, s);
161
+ this.stats.misses++;
162
+ if (this.missLog.length < 4096) this.missLog.push(id);
163
+ }
164
+
165
+ // Makes every id resident and returns its slot. Rows one call uses are never evicted by the same
166
+ // call, so a call may name at most `slots` distinct ids.
167
+ lookup(ids, out = new Uint32Array(ids.length)) {
168
+ const call = ++this.clock;
169
+ this.stats.lookups++;
170
+ this.stats.ids += ids.length;
171
+ for (let i = 0; i < ids.length; i++) {
172
+ const id = ids[i];
173
+ let s = this.slotOf[id];
174
+ if (s < 0) {
175
+ s = this.head;
176
+ if (this.stamp[s] === call) throw new Error(`RowCache: one lookup needs more than ${this.slots} rows`);
177
+ this.#install(id, s);
178
+ }
179
+ this.stamp[s] = call;
180
+ this.#touch(s);
181
+ out[i] = s;
182
+ }
183
+ return out;
184
+ }
185
+
186
+ has(id) { return this.slotOf[id] >= 0; }
187
+
188
+ // Loads rows [row0, row0 + count) into the least recently used slots in large reads (for a warm
189
+ // set at load), then uploads the whole id->slot map once. Returns the milliseconds spent.
190
+ warm(row0, count, batch = 4096) {
191
+ const t0 = performance.now();
192
+ count = Math.min(count, this.slots);
193
+ const rb = this.file.rowBytes;
194
+ const buf = new Uint8Array(batch * rb);
195
+ const planeBufs = this.planes.map((p) => new Uint8Array(batch * p.bytes));
196
+ for (let r = row0; r < row0 + count; r += batch) {
197
+ const n = Math.min(batch, row0 + count - r);
198
+ this.file.readRows(r, n, buf);
199
+ // Fresh slots come from the LRU head in order; collect them and write contiguous runs.
200
+ const slots = new Int32Array(n);
201
+ for (let i = 0; i < n; i++) {
202
+ const id = r + i;
203
+ let s = this.slotOf[id];
204
+ if (s < 0) {
205
+ s = this.head;
206
+ const old = this.idOf[s];
207
+ if (old >= 0) this.slotOf[old] = -1;
208
+ this.idOf[s] = id;
209
+ this.slotOf[id] = s;
210
+ }
211
+ this.#touch(s);
212
+ slots[i] = s;
213
+ }
214
+ for (let pi = 0; pi < this.planes.length; pi++) {
215
+ const p = this.planes[pi], dst = planeBufs[pi];
216
+ for (let i = 0; i < n; i++) dst.set(buf.subarray(i * rb + p.offset, i * rb + p.offset + p.bytes), i * p.bytes);
217
+ let i = 0;
218
+ while (i < n) {
219
+ let j = i + 1;
220
+ while (j < n && slots[j] === slots[j - 1] + 1) j++;
221
+ this.queue.writeBuffer(p.buffer, slots[i] * p.bytes, dst, i * p.bytes, (j - i) * p.bytes);
222
+ i = j;
223
+ }
224
+ }
225
+ }
226
+ // slotOf's bytes are the GPU map: -1 as an Int32 is 0xFFFFFFFF as a u32.
227
+ if (this.mapBuffer) this.queue.writeBuffer(this.mapBuffer, 0, this.slotOf);
228
+ return performance.now() - t0;
229
+ }
230
+ }
style.css DELETED
@@ -1,28 +0,0 @@
1
- body {
2
- padding: 2rem;
3
- font-family: -apple-system, BlinkMacSystemFont, "Arial", sans-serif;
4
- }
5
-
6
- h1 {
7
- font-size: 16px;
8
- margin-top: 0;
9
- }
10
-
11
- p {
12
- color: rgb(107, 114, 128);
13
- font-size: 15px;
14
- margin-bottom: 10px;
15
- margin-top: 5px;
16
- }
17
-
18
- .card {
19
- max-width: 620px;
20
- margin: 0 auto;
21
- padding: 16px;
22
- border: 1px solid lightgray;
23
- border-radius: 16px;
24
- }
25
-
26
- .card p:last-child {
27
- margin-bottom: 0;
28
- }